[MoE] [MoE Refactor] Add moe kernel oracle abc 37753 (#43461)
Signed-off-by: Qiuyang Yue <yueqiuyang1389@gmail.com> Signed-off-by: qyYue1389 <yueqiuyang1389@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Qiuyang Yue committed
Jun 25, 2026 at 06:34 UTC
9bfd878a48fb92e5e28267f07ad1b7d844ebf160
5 files changed
+280
-6
tests/kernels/moe/test_moe_kernel_oracle.py
new
+51
@@ -0,0 +1,51 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the MoEKernelOracle ABC introduced in PR series for #37753.
4
+
5
+This file contains a single canonical demonstration that
6
+`UnquantizedMoEKernelOracle` methods delegate one-to-one to the
7
+existing module-level functions in `oracle/unquantized.py`. Each method
8
+on `UnquantizedMoEKernelOracle` follows the same `return module_fn(args)`
9
+pattern, so verifying delegation for one method (`make_kernel`) gives
10
+high confidence in the rest.
11
+"""
12
+
13
+from unittest.mock import patch
14
+
15
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
16
+from vllm.model_executor.layers.fused_moe.oracle import UnquantizedMoEKernelOracle
17
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
18
+ UnquantizedMoeBackend,
19
+)
20
+
21
+
22
+class TestUnquantizedDelegation:
23
+ """UnquantizedMoEKernelOracle methods must delegate to the existing
24
+ module-level functions; behaviour is bit-identical."""
25
+
26
+ def test_make_kernel_delegates(self) -> None:
27
+ quant_config = object()
28
+ moe_config = object()
29
+ experts_cls = TritonExperts
30
+ sentinel_kernel = object()
31
+
32
+ with patch(
33
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized."
34
+ "make_unquantized_moe_kernel",
35
+ return_value=sentinel_kernel,
36
+ ) as mocked:
37
+ out = UnquantizedMoEKernelOracle().make_kernel(
38
+ quant_config,
39
+ moe_config,
40
+ UnquantizedMoeBackend.TRITON,
41
+ experts_cls,
42
+ )
43
+
44
+ mocked.assert_called_once_with(
45
+ quant_config,
46
+ moe_config,
47
+ UnquantizedMoeBackend.TRITON,
48
+ experts_cls,
49
+ None, # routing_tables default
50
+ )
51
+ assert out is sentinel_kernel
vllm/model_executor/layers/fused_moe/oracle/__init__.py
+10
@@ -1,2 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle
5
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
6
+ UnquantizedMoEKernelOracle,
7
+)
8
+
9
+__all__ = [
10
+ "MoEKernelOracle",
11
+ "UnquantizedMoEKernelOracle",
12
+]
vllm/model_executor/layers/fused_moe/oracle/base.py
new
+141
@@ -0,0 +1,141 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Abstract base class for MoE kernel oracles.
5
+
6
+Each MoE oracle (unquantized / fp8 / nvfp4 / mxfp4 / mxfp8 / int8 /
7
+int_wna16) is responsible for selecting the right MoE kernel backend for a
8
+given (model, hardware, deployment-config) tuple. The current
9
+implementation expresses this responsibility as module-level functions
10
+that follow an informal convention.
11
+
12
+This module declares the abstract contract; concrete oracles inherit from
13
+`MoEKernelOracle` and provide the platform-specific behaviour.
14
+
15
+This is the first PR in the series suggested by @robertgshaw2-redhat in
16
+PR #37776 (see issue #37753). It intentionally only introduces the ABC;
17
+follow-up PRs migrate each oracle to inherit from it. The single concrete
18
+subclass shipped here (`UnquantizedMoEKernelOracle`) delegates to the
19
+existing module-level functions to keep behaviour bit-identical with
20
+pre-class code.
21
+"""
22
+
23
+from abc import ABC, abstractmethod
24
+from enum import Enum
25
+from typing import TYPE_CHECKING, Generic, TypeVar
26
+
27
+import torch
28
+
29
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
30
+from vllm.config.kernel import MoEBackend
31
+from vllm.model_executor.layers.fused_moe.config import (
32
+ FusedMoEConfig,
33
+ FusedMoEQuantConfig,
34
+)
35
+
36
+if TYPE_CHECKING:
37
+ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
38
+
39
+BackendT = TypeVar("BackendT", bound=Enum)
40
+
41
+
42
+class MoEKernelOracle(ABC, Generic[BackendT]):
43
+ """Abstract base for MoE kernel-selection oracles.
44
+
45
+ Concrete oracles MUST implement: `backend_enum_cls`,
46
+ `get_priority_backends`, `backend_to_kernel_cls`, `map_backend`,
47
+ `select_backend`, `make_kernel`.
48
+
49
+ Concrete oracles MAY override: `convert_to_kernel_format`,
50
+ `make_quant_config`. The base class provides default implementations
51
+ that are appropriate for oracles which do not need them
52
+ (e.g. `make_quant_config` raises on the unquantized oracle).
53
+ """
54
+
55
+ @abstractmethod
56
+ def backend_enum_cls(self) -> type[BackendT]:
57
+ """Return the concrete `Enum` class enumerating this oracle's
58
+ backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`)."""
59
+
60
+ @abstractmethod
61
+ def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]:
62
+ """Return platform-appropriate backends in priority order for
63
+ this `moe_config`."""
64
+
65
+ @abstractmethod
66
+ def backend_to_kernel_cls(self, backend: BackendT) -> type[mk.FusedMoEExperts]:
67
+ """Map a backend enum value to its concrete `FusedMoEExperts`
68
+ subclass."""
69
+
70
+ @abstractmethod
71
+ def map_backend(self, runner_backend: MoEBackend) -> BackendT:
72
+ """Map a user-facing `MoEBackend` (from the runner config) to
73
+ this oracle's enum."""
74
+
75
+ @abstractmethod
76
+ def select_backend(
77
+ self,
78
+ moe_config: FusedMoEConfig,
79
+ weight_key: "QuantKey | None" = None,
80
+ activation_key: "QuantKey | None" = None,
81
+ ) -> tuple[BackendT, type[mk.FusedMoEExperts] | None]:
82
+ """Primary entry point: choose the best supported backend for
83
+ the given `moe_config`.
84
+
85
+ `weight_key` / `activation_key` carry the quantization scheme of
86
+ the weights and activations and are consumed by quantized oracles
87
+ (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized
88
+ oracle ignores them. Subclasses with additional selection inputs
89
+ (e.g. int_wna16 needs `weight_bits`, fp8 needs
90
+ `allow_vllm_cutlass`) widen the signature in their override; a
91
+ per-oracle config object is the longer-term target tracked in
92
+ the #37753 follow-up PRs.
93
+ """
94
+
95
+ @abstractmethod
96
+ def make_kernel(
97
+ self,
98
+ quant_config: FusedMoEQuantConfig,
99
+ moe_config: FusedMoEConfig,
100
+ backend: BackendT,
101
+ experts_cls: type[mk.FusedMoEExperts],
102
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
103
+ ) -> mk.FusedMoEKernel:
104
+ """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts
105
+ combinator) for the chosen backend."""
106
+
107
+ def convert_to_kernel_format(
108
+ self,
109
+ backend: BackendT,
110
+ moe_config: FusedMoEConfig,
111
+ w13_weight: torch.Tensor,
112
+ w2_weight: torch.Tensor,
113
+ ) -> tuple[torch.Tensor, torch.Tensor]:
114
+ """Shuffle weights into the layout expected by `backend`.
115
+
116
+ Default implementation returns the inputs unchanged. Oracles
117
+ whose backends need weight permutation should override this
118
+ (e.g. `UnquantizedMoEKernelOracle` handles AITER and FlashInfer
119
+ layouts).
120
+
121
+ `moe_config` carries MoE-layer state (e.g. `is_act_and_mul`)
122
+ that the conversion needs without coupling the oracle to a
123
+ `Module` reference. Quantized oracles whose conversion
124
+ additionally needs scales / zero-points / block shapes will
125
+ override with a wider signature (and ultimately a per-oracle
126
+ config object — tracked in the #37753 follow-up PRs).
127
+ """
128
+ return w13_weight, w2_weight
129
+
130
+ def make_quant_config(self, *args, **kwargs) -> FusedMoEQuantConfig:
131
+ """Build a `FusedMoEQuantConfig` for this oracle.
132
+
133
+ Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with
134
+ the appropriate signature for their quantization scheme.
135
+ Unquantized oracles inherit the default, which raises because
136
+ there is no quantization-specific config to build.
137
+ """
138
+ raise NotImplementedError(
139
+ f"{type(self).__name__} does not implement make_quant_config; "
140
+ "this oracle has no quantization-specific config to build."
141
+ )
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+77
-5
@@ -2,9 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from enum import Enum
5
+from typing import TYPE_CHECKING
6
7
import torch
7
-from torch.nn import Module
8
9
import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
20
)
21
+from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle
22
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
23
align_moe_weights_for_fi,
24
convert_moe_weights_to_flashinfer_trtllm_block_layout,
@@ -25,6 +26,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
26
)
27
from vllm.platforms import current_platform
28
29
+if TYPE_CHECKING:
30
+ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
31
+
32
logger = init_logger(__name__)
33
34
@@ -256,7 +260,7 @@ def select_unquantized_moe_backend(
260
261
def convert_to_unquantized_kernel_format(
262
unquantized_backend: UnquantizedMoeBackend,
259
- layer: Module,
263
+ moe_config: FusedMoEConfig,
264
w13_weight: torch.Tensor,
265
w2_weight: torch.Tensor,
266
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -264,13 +268,13 @@ def convert_to_unquantized_kernel_format(
268
w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
269
270
elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
267
- if layer.moe_config.is_act_and_mul:
271
+ if moe_config.is_act_and_mul:
272
# Swap halves to arrange as [w3; w1] (kernel expectation)
273
# Non-gated MoE: w13 is a single projection, no need to swap.
274
w13_weight = swap_w13_to_w31(w13_weight)
275
276
elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
273
- is_act_and_mul = layer.moe_config.is_act_and_mul
277
+ is_act_and_mul = moe_config.is_act_and_mul
278
if not is_act_and_mul:
279
# Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK
280
# weight layout uses block_k=128). Pad along the intermediate dim when
@@ -278,7 +282,7 @@ def convert_to_unquantized_kernel_format(
282
w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi(
283
w13_weight, w2_weight, is_act_and_mul, min_alignment=128
284
)
281
- layer.moe_config.intermediate_size_per_partition = padded_intermediate
285
+ moe_config.intermediate_size_per_partition = padded_intermediate
286
287
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
288
w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
@@ -333,3 +337,71 @@ def make_unquantized_moe_kernel(
337
)
338
339
return kernel
340
+
341
+
342
+# ---------------------------------------------------------------------------
343
+# Class-based view (first PR of the #37753 series; see oracle/base.py).
344
+# Methods delegate to the module-level functions above so behaviour is
345
+# bit-identical with pre-class code.
346
+# ---------------------------------------------------------------------------
347
+
348
+
349
+class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]):
350
+ """Class-based view of the unquantized MoE kernel oracle.
351
+
352
+ Each method delegates to its module-level counterpart so that
353
+ instantiating and calling this class is bit-identical to calling
354
+ the standalone functions. Follow-up PRs may move logic from the
355
+ module-level functions into these methods.
356
+ """
357
+
358
+ def backend_enum_cls(self) -> type[UnquantizedMoeBackend]:
359
+ return UnquantizedMoeBackend
360
+
361
+ def get_priority_backends(
362
+ self, moe_config: FusedMoEConfig
363
+ ) -> list[UnquantizedMoeBackend]:
364
+ return _get_priority_backends(moe_config)
365
+
366
+ def backend_to_kernel_cls(
367
+ self, backend: UnquantizedMoeBackend
368
+ ) -> type[mk.FusedMoEExperts]:
369
+ return backend_to_kernel_cls(backend)
370
+
371
+ def map_backend(self, runner_backend: MoEBackend) -> UnquantizedMoeBackend:
372
+ return map_unquantized_backend(runner_backend)
373
+
374
+ def select_backend(
375
+ self,
376
+ moe_config: FusedMoEConfig,
377
+ weight_key: "QuantKey | None" = None,
378
+ activation_key: "QuantKey | None" = None,
379
+ ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
380
+ assert weight_key is None and activation_key is None, (
381
+ "Weights and activations will never be quantized for "
382
+ "UnquantizedMoEKernelOracle"
383
+ )
384
+ return select_unquantized_moe_backend(moe_config)
385
+
386
+ def convert_to_kernel_format(
387
+ self,
388
+ backend: UnquantizedMoeBackend,
389
+ moe_config: FusedMoEConfig,
390
+ w13_weight: torch.Tensor,
391
+ w2_weight: torch.Tensor,
392
+ ) -> tuple[torch.Tensor, torch.Tensor]:
393
+ return convert_to_unquantized_kernel_format(
394
+ backend, moe_config, w13_weight, w2_weight
395
+ )
396
+
397
+ def make_kernel(
398
+ self,
399
+ quant_config: FusedMoEQuantConfig,
400
+ moe_config: FusedMoEConfig,
401
+ backend: UnquantizedMoeBackend,
402
+ experts_cls: type[mk.FusedMoEExperts],
403
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
404
+ ) -> mk.FusedMoEKernel:
405
+ return make_unquantized_moe_kernel(
406
+ quant_config, moe_config, backend, experts_cls, routing_tables
407
+ )
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1
-1
@@ -161,7 +161,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
161
# Shuffle weights to runtime format.
162
w13_new, w2_new = convert_to_unquantized_kernel_format(
163
self.unquantized_backend,
164
- layer=layer,
164
+ moe_config=layer.moe_config,
165
w13_weight=w13,
166
w2_weight=w2,
167
)