Skip to content

vllm.model_executor.layers.fused_moe.oracle.base

Abstract base class for MoE kernel oracles.

Classes:

Functions:

  • make_moe_kernel –

    Assemble a FusedMoEKernel from a prepare/finalize and an experts part.

MoEKernelOracle

Bases: ABC, Generic[BackendT]

Abstract base for MoE kernel-selection oracles.

Concrete oracles provide name and backend_to_kernel_cls today; the remaining methods below are the #37753 migration target and are still implemented per module, so they are declared but not enforced.

Methods:

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
class MoEKernelOracle(ABC, Generic[BackendT]):
    """Abstract base for MoE kernel-selection oracles.

    Concrete oracles provide `name` and `backend_to_kernel_cls` today; the
    remaining methods below are the #37753 migration target and are still
    implemented per module, so they are declared but not enforced.
    """

    # Display name for the shared log helpers below; subclasses must set it.
    name: ClassVar[str]

    def __init__(
        self,
        moe_config: FusedMoEConfig,
        weight_key: "QuantKey | None" = None,
        activation_key: "QuantKey | None" = None,
    ) -> None:
        # Deployment inputs shared by every backend probe of one selection.
        self.moe_config = moe_config
        self.weight_key = weight_key
        self.activation_key = activation_key

    def backend_enum_cls(self) -> type[BackendT]:
        """Return the concrete `Enum` class enumerating this oracle's
        backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`)."""
        raise NotImplementedError

    def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]:
        """Return platform-appropriate backends in priority order for
        this `moe_config`."""
        raise NotImplementedError

    def backend_to_kernel_cls(
        self, backend: BackendT
    ) -> list[type[mk.FusedMoEExperts]]:
        """Map a backend enum value to its concrete `FusedMoEExperts`
        subclasses, in selection priority order.

        Implementations must forward to the module-level function at call time,
        not alias it with `= staticmethod(fn)`, which snapshots it at import.
        """
        raise NotImplementedError

    def map_backend(self, runner_backend: MoEBackend) -> BackendT:
        """Map a user-facing `MoEBackend` (from the runner config) to
        this oracle's enum."""
        raise NotImplementedError

    def select_backend(
        self,
        moe_config: FusedMoEConfig,
        weight_key: "QuantKey | None" = None,
        activation_key: "QuantKey | None" = None,
    ) -> tuple[BackendT, type[mk.FusedMoEExperts] | None]:
        """Primary entry point: choose the best supported backend for
        the given `moe_config`.

        `weight_key` / `activation_key` carry the quantization scheme of
        the weights and activations and are consumed by quantized oracles
        (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized
        oracle ignores them. Subclasses with additional selection inputs
        (e.g. int_wna16 needs `weight_bits`, fp8 needs
        `allow_vllm_cutlass`) widen the signature in their override; a
        per-oracle config object is the longer-term target tracked in
        the #37753 follow-up PRs.
        """
        raise NotImplementedError

    def make_kernel(
        self,
        quant_config: FusedMoEQuantConfig,
        moe_config: FusedMoEConfig,
        backend: BackendT,
        experts_cls: type[mk.FusedMoEExperts],
        routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
    ) -> mk.FusedMoEKernel:
        """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts
        combinator) for the chosen backend."""
        raise NotImplementedError

    def convert_to_kernel_format(
        self,
        backend: BackendT,
        moe_config: FusedMoEConfig,
        w13_weight: torch.Tensor,
        w2_weight: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Shuffle weights into the layout expected by `backend`.

        Default implementation returns the inputs unchanged. Oracles
        whose backends need weight permutation should override this
        (e.g. `UnquantizedMoEKernelOracle` handles AITER and FlashInfer
        layouts).

        `moe_config` carries MoE-layer state (e.g. `is_act_and_mul`)
        that the conversion needs without coupling the oracle to a
        `Module` reference. Quantized oracles whose conversion
        additionally needs scales / zero-points / block shapes will
        override with a wider signature (and ultimately a per-oracle
        config object — tracked in the #37753 follow-up PRs).
        """
        return w13_weight, w2_weight

    def make_quant_config(self, *args, **kwargs) -> FusedMoEQuantConfig:
        """Build a `FusedMoEQuantConfig` for this oracle.

        Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with
        the appropriate signature for their quantization scheme.
        Unquantized oracles inherit the default, which raises because
        there is no quantization-specific config to build.
        """
        raise NotImplementedError(
            f"{type(self).__name__} does not implement make_quant_config; "
            "this oracle has no quantization-specific config to build."
        )

    def make_log_backend(
        self,
        backend: BackendT,
        available_backends: list[BackendT] | None = None,
        is_lora: bool = False,
    ) -> str:
        """Message for the backend a select loop settled on."""
        available = (
            f" out of potential backends: {[b.value for b in available_backends]}"
            if available_backends is not None
            else ""
        )
        lora = "LoRA " if is_lora else ""
        return f"Using {backend.value} {self.name} MoE {lora}backend{available}."

    def make_log_unsupported(
        self, backend: BackendT, reason: str | None, is_lora: bool = False
    ) -> str:
        """Message for a backend that does not support the configuration."""
        since = f" since {reason}" if reason else ""
        lora = "LoRA " if is_lora else ""
        return (
            f"{self.name} MoE {lora}backend '{backend.value}' does not support the "
            f"deployment configuration{since}."
        )

    def try_backend(
        self,
        backend: BackendT,
        kernel_classes: list[type[mk.FusedMoEExperts]] | None = None,
        activation_key: "QuantKey | None" = _UNSET,
    ) -> tuple[type[mk.FusedMoEExperts] | None, list[Rejection]]:
        """Probe `backend`'s candidate kernel classes, first supported wins.

        A miss logs a DEBUG line and yields to the next candidate; whether that
        is fatal stays with the caller. `activation_key` overrides this oracle's
        own key for oracles whose activation key varies per backend.

        Returns:
            `(kernel_cls, rejected)`: the winning class, or `None`, plus every
            `(kernel_cls, reason)` probed and rejected before it, in order, so
            that callers collecting reasons for a report keep one entry per
            candidate rather than only the last.

        """
        rejected: list[Rejection] = []
        act_quant_key = (
            self.activation_key if activation_key is _UNSET else activation_key
        )
        if kernel_classes is None:
            kernel_classes = self.backend_to_kernel_cls(backend)
        for kernel_cls in kernel_classes:
            supported, reason = kernel_cls.is_supported_config(
                kernel_cls,
                self.moe_config,
                self.weight_key,
                act_quant_key,
                self.moe_config.activation_format,
            )
            if supported:
                return kernel_cls, rejected
            logger.debug_once(self.make_log_unsupported(backend, reason))
            rejected.append((kernel_cls, reason))
        return None, rejected

    def return_or_raise(
        self,
        backend: BackendT,
        available_backends: list[BackendT] | None = None,
        kernel_classes: list[type[mk.FusedMoEExperts]] | None = None,
        activation_key: "QuantKey | None" = _UNSET,
        is_lora: bool = False,
    ) -> tuple[BackendT, type[mk.FusedMoEExperts]]:
        """Return `backend` with its first supported kernel class, or raise.

        Raises:
            ValueError: no candidate kernel class supports the deployment
                configuration; the message carries the last rejection reason.

        """
        kernel_cls, rejected = self.try_backend(backend, kernel_classes, activation_key)
        if kernel_cls is None:
            reason = rejected[-1][1] if rejected else None
            raise ValueError(self.make_log_unsupported(backend, reason, is_lora))
        logger.info_once(self.make_log_backend(backend, available_backends, is_lora))
        return backend, kernel_cls

backend_enum_cls()

Return the concrete Enum class enumerating this oracle's backends (e.g. UnquantizedMoeBackend, Fp8MoeBackend).

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
59
60
61
62
def backend_enum_cls(self) -> type[BackendT]:
    """Return the concrete `Enum` class enumerating this oracle's
    backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`)."""
    raise NotImplementedError

backend_to_kernel_cls(backend)

Map a backend enum value to its concrete FusedMoEExperts subclasses, in selection priority order.

Implementations must forward to the module-level function at call time, not alias it with = staticmethod(fn), which snapshots it at import.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
69
70
71
72
73
74
75
76
77
78
def backend_to_kernel_cls(
    self, backend: BackendT
) -> list[type[mk.FusedMoEExperts]]:
    """Map a backend enum value to its concrete `FusedMoEExperts`
    subclasses, in selection priority order.

    Implementations must forward to the module-level function at call time,
    not alias it with `= staticmethod(fn)`, which snapshots it at import.
    """
    raise NotImplementedError

convert_to_kernel_format(backend, moe_config, w13_weight, w2_weight)

Shuffle weights into the layout expected by backend.

Default implementation returns the inputs unchanged. Oracles whose backends need weight permutation should override this (e.g. UnquantizedMoEKernelOracle handles AITER and FlashInfer layouts).

moe_config carries MoE-layer state (e.g. is_act_and_mul) that the conversion needs without coupling the oracle to a Module reference. Quantized oracles whose conversion additionally needs scales / zero-points / block shapes will override with a wider signature (and ultimately a per-oracle config object — tracked in the #37753 follow-up PRs).

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def convert_to_kernel_format(
    self,
    backend: BackendT,
    moe_config: FusedMoEConfig,
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Shuffle weights into the layout expected by `backend`.

    Default implementation returns the inputs unchanged. Oracles
    whose backends need weight permutation should override this
    (e.g. `UnquantizedMoEKernelOracle` handles AITER and FlashInfer
    layouts).

    `moe_config` carries MoE-layer state (e.g. `is_act_and_mul`)
    that the conversion needs without coupling the oracle to a
    `Module` reference. Quantized oracles whose conversion
    additionally needs scales / zero-points / block shapes will
    override with a wider signature (and ultimately a per-oracle
    config object — tracked in the #37753 follow-up PRs).
    """
    return w13_weight, w2_weight

get_priority_backends(moe_config)

Return platform-appropriate backends in priority order for this moe_config.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
64
65
66
67
def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]:
    """Return platform-appropriate backends in priority order for
    this `moe_config`."""
    raise NotImplementedError

make_kernel(quant_config, moe_config, backend, experts_cls, routing_tables=None)

Construct the FusedMoEKernel (Prepare/Finalize + Experts combinator) for the chosen backend.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
105
106
107
108
109
110
111
112
113
114
115
def make_kernel(
    self,
    quant_config: FusedMoEQuantConfig,
    moe_config: FusedMoEConfig,
    backend: BackendT,
    experts_cls: type[mk.FusedMoEExperts],
    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEKernel:
    """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts
    combinator) for the chosen backend."""
    raise NotImplementedError

make_log_backend(backend, available_backends=None, is_lora=False)

Message for the backend a select loop settled on.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def make_log_backend(
    self,
    backend: BackendT,
    available_backends: list[BackendT] | None = None,
    is_lora: bool = False,
) -> str:
    """Message for the backend a select loop settled on."""
    available = (
        f" out of potential backends: {[b.value for b in available_backends]}"
        if available_backends is not None
        else ""
    )
    lora = "LoRA " if is_lora else ""
    return f"Using {backend.value} {self.name} MoE {lora}backend{available}."

make_log_unsupported(backend, reason, is_lora=False)

Message for a backend that does not support the configuration.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
168
169
170
171
172
173
174
175
176
177
def make_log_unsupported(
    self, backend: BackendT, reason: str | None, is_lora: bool = False
) -> str:
    """Message for a backend that does not support the configuration."""
    since = f" since {reason}" if reason else ""
    lora = "LoRA " if is_lora else ""
    return (
        f"{self.name} MoE {lora}backend '{backend.value}' does not support the "
        f"deployment configuration{since}."
    )

make_quant_config(*args, **kwargs)

Build a FusedMoEQuantConfig for this oracle.

Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with the appropriate signature for their quantization scheme. Unquantized oracles inherit the default, which raises because there is no quantization-specific config to build.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
140
141
142
143
144
145
146
147
148
149
150
151
def make_quant_config(self, *args, **kwargs) -> FusedMoEQuantConfig:
    """Build a `FusedMoEQuantConfig` for this oracle.

    Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with
    the appropriate signature for their quantization scheme.
    Unquantized oracles inherit the default, which raises because
    there is no quantization-specific config to build.
    """
    raise NotImplementedError(
        f"{type(self).__name__} does not implement make_quant_config; "
        "this oracle has no quantization-specific config to build."
    )

map_backend(runner_backend)

Map a user-facing MoEBackend (from the runner config) to this oracle's enum.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
80
81
82
83
def map_backend(self, runner_backend: MoEBackend) -> BackendT:
    """Map a user-facing `MoEBackend` (from the runner config) to
    this oracle's enum."""
    raise NotImplementedError

return_or_raise(backend, available_backends=None, kernel_classes=None, activation_key=_UNSET, is_lora=False)

Return backend with its first supported kernel class, or raise.

Raises:

  • ValueError –

    no candidate kernel class supports the deployment configuration; the message carries the last rejection reason.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def return_or_raise(
    self,
    backend: BackendT,
    available_backends: list[BackendT] | None = None,
    kernel_classes: list[type[mk.FusedMoEExperts]] | None = None,
    activation_key: "QuantKey | None" = _UNSET,
    is_lora: bool = False,
) -> tuple[BackendT, type[mk.FusedMoEExperts]]:
    """Return `backend` with its first supported kernel class, or raise.

    Raises:
        ValueError: no candidate kernel class supports the deployment
            configuration; the message carries the last rejection reason.

    """
    kernel_cls, rejected = self.try_backend(backend, kernel_classes, activation_key)
    if kernel_cls is None:
        reason = rejected[-1][1] if rejected else None
        raise ValueError(self.make_log_unsupported(backend, reason, is_lora))
    logger.info_once(self.make_log_backend(backend, available_backends, is_lora))
    return backend, kernel_cls

select_backend(moe_config, weight_key=None, activation_key=None)

Primary entry point: choose the best supported backend for the given moe_config.

weight_key / activation_key carry the quantization scheme of the weights and activations and are consumed by quantized oracles (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized oracle ignores them. Subclasses with additional selection inputs (e.g. int_wna16 needs weight_bits, fp8 needs allow_vllm_cutlass) widen the signature in their override; a per-oracle config object is the longer-term target tracked in the #37753 follow-up PRs.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def select_backend(
    self,
    moe_config: FusedMoEConfig,
    weight_key: "QuantKey | None" = None,
    activation_key: "QuantKey | None" = None,
) -> tuple[BackendT, type[mk.FusedMoEExperts] | None]:
    """Primary entry point: choose the best supported backend for
    the given `moe_config`.

    `weight_key` / `activation_key` carry the quantization scheme of
    the weights and activations and are consumed by quantized oracles
    (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized
    oracle ignores them. Subclasses with additional selection inputs
    (e.g. int_wna16 needs `weight_bits`, fp8 needs
    `allow_vllm_cutlass`) widen the signature in their override; a
    per-oracle config object is the longer-term target tracked in
    the #37753 follow-up PRs.
    """
    raise NotImplementedError

try_backend(backend, kernel_classes=None, activation_key=_UNSET)

Probe backend's candidate kernel classes, first supported wins.

A miss logs a DEBUG line and yields to the next candidate; whether that is fatal stays with the caller. activation_key overrides this oracle's own key for oracles whose activation key varies per backend.

Returns:

  • type[FusedMoEExperts] | None –

    (kernel_cls, rejected): the winning class, or None, plus every

  • list[Rejection] –

    (kernel_cls, reason) probed and rejected before it, in order, so

  • tuple[type[FusedMoEExperts] | None, list[Rejection]] –

    that callers collecting reasons for a report keep one entry per

  • tuple[type[FusedMoEExperts] | None, list[Rejection]] –

    candidate rather than only the last.

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def try_backend(
    self,
    backend: BackendT,
    kernel_classes: list[type[mk.FusedMoEExperts]] | None = None,
    activation_key: "QuantKey | None" = _UNSET,
) -> tuple[type[mk.FusedMoEExperts] | None, list[Rejection]]:
    """Probe `backend`'s candidate kernel classes, first supported wins.

    A miss logs a DEBUG line and yields to the next candidate; whether that
    is fatal stays with the caller. `activation_key` overrides this oracle's
    own key for oracles whose activation key varies per backend.

    Returns:
        `(kernel_cls, rejected)`: the winning class, or `None`, plus every
        `(kernel_cls, reason)` probed and rejected before it, in order, so
        that callers collecting reasons for a report keep one entry per
        candidate rather than only the last.

    """
    rejected: list[Rejection] = []
    act_quant_key = (
        self.activation_key if activation_key is _UNSET else activation_key
    )
    if kernel_classes is None:
        kernel_classes = self.backend_to_kernel_cls(backend)
    for kernel_cls in kernel_classes:
        supported, reason = kernel_cls.is_supported_config(
            kernel_cls,
            self.moe_config,
            self.weight_key,
            act_quant_key,
            self.moe_config.activation_format,
        )
        if supported:
            return kernel_cls, rejected
        logger.debug_once(self.make_log_unsupported(backend, reason))
        rejected.append((kernel_cls, reason))
    return None, rejected

make_moe_kernel(moe_quant_config, moe_config, experts_cls, routing_tables=None, prepare_finalize_kwargs=None, experts_kwargs=None, experts_log=None)

Assemble a FusedMoEKernel from a prepare/finalize and an experts part.

This is the step every MoE oracle runs once its backend is chosen, so it lives here rather than being copy-pasted into each oracle module.

Parameters:

  • moe_quant_config

    (FusedMoEQuantConfig) –

    Quantization configuration of the layer.

  • moe_config

    (FusedMoEConfig) –

    MoE layer configuration.

  • experts_cls

    (type[FusedMoEExperts]) –

    Selected experts implementation.

  • routing_tables

    (tuple[Tensor, Tensor, Tensor] | None, default: None ) –

    Optional pre-computed routing tables (EP).

  • prepare_finalize_kwargs

    (dict[str, Any] | None, default: None ) –

    Extra kwargs for maybe_make_prepare_finalize (e.g. input_dtype).

  • experts_kwargs

    (dict[str, Any] | None, default: None ) –

    Extra kwargs for the experts constructor (e.g. the W4A8 b_strides1 / b_strides2 / group_size).

  • experts_log

    (str | None, default: None ) –

    Optional info_once template for the selected experts class, e.g. "Using %s MoE backend".

Source code in vllm/model_executor/layers/fused_moe/oracle/base.py 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
def make_moe_kernel(
    moe_quant_config: FusedMoEQuantConfig,
    moe_config: FusedMoEConfig,
    experts_cls: type[mk.FusedMoEExperts],
    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
    prepare_finalize_kwargs: dict[str, Any] | None = None,
    experts_kwargs: dict[str, Any] | None = None,
    experts_log: str | None = None,
) -> mk.FusedMoEKernel:
    """Assemble a `FusedMoEKernel` from a prepare/finalize and an experts part.

    This is the step every MoE oracle runs once its backend is chosen, so it
    lives here rather than being copy-pasted into each oracle module.

    Args:
        moe_quant_config: Quantization configuration of the layer.
        moe_config: MoE layer configuration.
        experts_cls: Selected experts implementation.
        routing_tables: Optional pre-computed routing tables (EP).
        prepare_finalize_kwargs: Extra kwargs for `maybe_make_prepare_finalize`
            (e.g. `input_dtype`).
        experts_kwargs: Extra kwargs for the experts constructor (e.g. the W4A8
            `b_strides1` / `b_strides2` / `group_size`).
        experts_log: Optional `info_once` template for the selected experts
            class, e.g. `"Using %s MoE backend"`.

    """
    prepare_finalize = maybe_make_prepare_finalize(
        moe=moe_config,
        quant_config=moe_quant_config,
        routing_tables=routing_tables,
        allow_new_interface=True,
        use_monolithic=experts_cls.is_monolithic(),
        **(prepare_finalize_kwargs or {}),
    )
    assert prepare_finalize is not None
    if not experts_cls.activation_format().is_superset(
        prepare_finalize.activation_format
    ):
        raise ValueError(
            f"{experts_cls.__name__} expects "
            f"{experts_cls.activation_format().value[0]} activations but the selected "
            f"prepare/finalize uses {prepare_finalize.activation_format.value[0]}"
        )

    logger.info_once("Using %s", prepare_finalize.__class__.__name__)
    if experts_log is not None:
        logger.info_once(experts_log, experts_cls.__name__)

    extra_kwargs = dict(experts_kwargs or {})
    if prepare_finalize.activation_format.is_batched:
        max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
        assert max_num_tokens is not None
        extra_kwargs["max_num_tokens"] = max_num_tokens
        extra_kwargs["num_dispatchers"] = prepare_finalize.num_dispatchers()

    experts = experts_cls(
        moe_config=moe_config,
        quant_config=moe_quant_config,
        **extra_kwargs,
    )
    return mk.FusedMoEKernel(prepare_finalize, experts)