[ROCm][MoE] Split AiterExperts by weight-quantization key

docs/design/moe_kernel_features.md

@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
 | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] |
-| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
+| rocm aiter moe | standard | none,</br>mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai,</br>situ | Y | Y | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.rocm_aiter_fused_experts],</br>[`AiterUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterUnquantizedExperts],</br>[`AiterFp8BlockExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8BlockExperts],</br>[`AiterFp8TensorExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8TensorExperts],</br>[`AiterFp8ChannelExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8ChannelExperts],</br>[`AiterMxfp4Experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterMxfp4Experts] |
 | cpu_moe | standard | N/A | N/A | silu, gelu,</br>gelu_tanh,</br>swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],</br>[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],</br>[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] |
 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
 

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -93,6 +93,7 @@ class ExpertInfo:
 ]
 common_float_and_int_types = common_float_types + [torch.int8]
 nvfp4_types = ["nvfp4"]
+mxfp4_types = ["mxfp4"]
 fp8_types = [current_platform.fp8_dtype()]
 
 
@@ -318,20 +319,37 @@ def expert_info(kind) -> ExpertInfo:
 
 if has_aiter():
     from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-        AiterExperts,
+        AiterFp8BlockExperts,
+        AiterFp8ChannelExperts,
+        AiterFp8TensorExperts,
+        AiterMxfp4Experts,
+        AiterUnquantizedExperts,
     )
 
-    register_experts(
-        AiterExperts,
-        standard_format,
-        # AiterExperts also supports the fully-unquantized (None, None)
-        # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8.
-        common_float_types,
-        blocked_quantization_support=True,
-        needs_aiter=True,
-    )
+    # Registration dtypes gate which quant configs the modular-kernel test
+    # harness will ever pair with a given class (see Config.is_valid()):
+    # they must match what each subclass's _supports_quant_scheme() actually
+    # declares, or the harness will generate a case the class rejects.
+    for aiter_experts_cls, aiter_supported_dtypes in (
+        (AiterUnquantizedExperts, common_float_types),
+        (AiterFp8BlockExperts, fp8_types),
+        (AiterFp8TensorExperts, fp8_types),
+        (AiterFp8ChannelExperts, fp8_types),
+        (AiterMxfp4Experts, mxfp4_types),
+    ):
+        register_experts(
+            aiter_experts_cls,
+            standard_format,
+            aiter_supported_dtypes,
+            blocked_quantization_support=True,
+            needs_aiter=True,
+        )
 else:
-    AiterExperts = None
+    AiterUnquantizedExperts = None
+    AiterFp8BlockExperts = None
+    AiterFp8TensorExperts = None
+    AiterFp8ChannelExperts = None
+    AiterMxfp4Experts = None
 
 if has_deep_gemm() and is_deep_gemm_supported():
     register_experts(

tests/kernels/moe/test_modular_kernel_combinations.py

@@ -16,7 +16,10 @@
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+    AiterExperts,
+    AiterUnquantizedExperts,
+)
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
 from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm
@@ -76,10 +79,10 @@ def assert_aiter_quant_scheme_case(config: Config) -> None:
     indirectly through the general quant-config sweep.
     See https://github.com/vllm-project/vllm/issues/54966."""
     fe_cls = config.fused_experts_type
-    if fe_cls is not AiterExperts:
+    if not issubclass(fe_cls, AiterExperts):
         return
 
-    if config.quant_config is None:
+    if config.quant_config is None or config.quant_dtype is None:
         w_key, a_key = None, None
     else:
         w_key, a_key = config.fp8_quant_key_pair()
@@ -97,7 +100,7 @@ def assert_aiter_activation_case(config: Config) -> None:
     AiterExperts being reached only indirectly through the general
     activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
     fe_cls = config.fused_experts_type
-    if fe_cls is not AiterExperts:
+    if not issubclass(fe_cls, AiterExperts):
         return
 
     assert fe_cls._supports_activation(config.activation), (
@@ -152,7 +155,7 @@ def rank_worker(
             if (
                 topk == 1
                 and config.supports_apply_weight_on_input()
-                and config.fused_experts_type is AiterExperts
+                and issubclass(config.fused_experts_type, AiterExperts)
                 and config.quant_block_shape is not None
             ):
                 print(
@@ -168,7 +171,7 @@ def rank_worker(
             # https://github.com/vllm-project/vllm/issues/57029
             if (
                 config.world_size > 1
-                and config.fused_experts_type is AiterExperts
+                and issubclass(config.fused_experts_type, AiterExperts)
                 and getattr(config.prepare_finalize_type, "__name__", "")
                 in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize")
             ):
@@ -206,7 +209,7 @@ def rank_worker(
 
             is_aiter_fp8 = (
                 _cp.is_rocm()
-                and config.fused_experts_type is AiterExperts
+                and issubclass(config.fused_experts_type, AiterExperts)
                 and config.quant_config is not None
             )
             if is_aiter_fp8:
@@ -266,7 +269,7 @@ def is_nyi_config(config: Config) -> bool:
             return True
 
     if config.activation != MoEActivation.SILU:
-        if config.fused_experts_type is not AiterExperts:
+        if not issubclass(config.fused_experts_type, AiterExperts):
             return True  # AITER-only for this axis, for now
         if config.quant_dtype is not None:
             return True  # unquantized-only for this axis, for now
@@ -587,7 +590,7 @@ def test_aiter_moe_sorting_backend_dispatch_env_matrix(
         dtype=torch.bfloat16,
         quant_config=None,
         prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
-        fused_experts_type=AiterExperts,
+        fused_experts_type=AiterUnquantizedExperts,
         world_size=1,
     )
     assert config.is_valid()[0]
@@ -667,7 +670,7 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
         dtype=torch.bfloat16,
         quant_config=None,
         prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
-        fused_experts_type=AiterExperts,
+        fused_experts_type=AiterUnquantizedExperts,
         world_size=1,
     )
     assert config.is_valid()[0]

tests/kernels/moe/test_mori_moe.py

@@ -42,7 +42,10 @@
 
 def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
     from vllm._aiter_ops import rocm_aiter_ops
-    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
+    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+        AiterFp8ChannelExperts,
+        AiterUnquantizedExperts,
+    )
     from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
         MoriPrepareAndFinalize,
     )
@@ -94,7 +97,8 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
     assert isinstance(prepare, MoriPrepareAndFinalize)
     assert prepare.use_fp8_dispatch == fp8
     assert prepare.num_dispatchers() == 2
-    experts = AiterExperts(config, quant)
+    experts_cls = AiterFp8ChannelExperts if fp8 else AiterUnquantizedExperts
+    experts = experts_cls(config, quant)
     assert not experts.expects_unquantized_inputs
     kernel = FusedMoEKernel(prepare, experts)
     local_w1, local_w2 = rocm_aiter_ops.shuffle_weights(

tests/kernels/moe/test_rocm_aiter_moe.py

@@ -848,13 +848,13 @@ def test_activation_method_enum_values():
 def test_aiter_mxfp4_quant_scheme_support_matches_gfx950():
     """AITER MXFP4 MoE support should stay gfx950-only."""
     from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-        AiterExperts,
+        AiterMxfp4Experts,
     )
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         kMxfp4Static,
     )
 
-    assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950()
+    assert AiterMxfp4Experts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950()
 
 
 @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")

tests/kernels/moe/test_topk_softplus_sqrt.py

@@ -87,7 +87,7 @@ def _torch_topk_softplus_sqrt(
     "backend,expected_experts",
     [
         ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"),
-        ("aiter", "AiterExperts"),
+        ("aiter", "AiterMxfp4Experts"),
         ("triton_unfused", "UnfusedOAITritonExperts"),
     ],
 )

vllm/model_executor/layers/fused_moe/__init__.py

@@ -106,7 +106,11 @@ def get_config() -> dict[str, Any] | None:
         BatchedTritonExperts,
     )
     from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-        AiterExperts,
+        AiterFp8BlockExperts,
+        AiterFp8ChannelExperts,
+        AiterFp8TensorExperts,
+        AiterMxfp4Experts,
+        AiterUnquantizedExperts,
     )
     from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
         TritonOrDeepGemmExperts,
@@ -132,7 +136,11 @@ def get_config() -> dict[str, Any] | None:
     )
 
     __all__ += [
-        "AiterExperts",
+        "AiterUnquantizedExperts",
+        "AiterFp8BlockExperts",
+        "AiterFp8TensorExperts",
+        "AiterFp8ChannelExperts",
+        "AiterMxfp4Experts",
         "fused_topk",
         "fused_experts",
         "get_config_file_name",

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -482,29 +482,6 @@ def _supports_current_device() -> bool:
     def _supports_no_act_and_mul() -> bool:
         return False
 
-    @staticmethod
-    def _supports_quant_scheme(
-        weight_key: QuantKey | None,
-        activation_key: QuantKey | None,
-    ) -> bool:
-        SUPPORTED_W_A = [
-            (None, None),
-            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
-            (kFp8StaticTensorSym, kFp8StaticTensorSym),
-            (kFp8StaticTensorSym, kFp8DynamicTensorSym),
-            (kFp8StaticChannelSym, kFp8DynamicTokenSym),
-            (kMxfp4Static, None),
-            (kMxfp4Static, kMxfp4Dynamic),
-        ]
-        if (weight_key, activation_key) not in SUPPORTED_W_A:
-            return False
-        if weight_key == kMxfp4Static:
-            from vllm.platforms.rocm import on_gfx950, on_gfx1250
-
-            if not on_gfx950() or on_gfx1250():
-                return False
-        return True
-
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
         return activation in [
@@ -597,3 +574,77 @@ def apply(
             output.set_(result)
         else:
             output.copy_(result)
+
+
+class AiterUnquantizedExperts(AiterExperts):
+    """Unquantized (bf16/fp16) weights and activations."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [(None, None)]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+
+class AiterFp8BlockExperts(AiterExperts):
+    """128x128 block-quantized weight scale, 128-grouped dynamic
+    activation scale."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+
+class AiterFp8TensorExperts(AiterExperts):
+    """Per-tensor static weight scale, either per-tensor static or
+    per-tensor dynamic activation scale."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [
+            (kFp8StaticTensorSym, kFp8StaticTensorSym),
+            (kFp8StaticTensorSym, kFp8DynamicTensorSym),
+        ]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+
+class AiterFp8ChannelExperts(AiterExperts):
+    """Per-output-channel static weight scale, per-token dynamic
+    activation scale."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [(kFp8StaticChannelSym, kFp8DynamicTokenSym)]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+
+class AiterMxfp4Experts(AiterExperts):
+    """MXFP4-quantized weights, with either unquantized (bf16/fp16) or
+    MXFP4-quantized dynamic activations. gfx950-only (not gfx1250)."""
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [
+            (kMxfp4Static, None),
+            (kMxfp4Static, kMxfp4Dynamic),
+        ]
+        if (weight_key, activation_key) not in SUPPORTED_W_A:
+            return False
+        from vllm.platforms.rocm import on_gfx950, on_gfx1250
+
+        return on_gfx950() and not on_gfx1250()

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -211,10 +211,12 @@ def backend_to_kernel_cls(
 
     elif backend == Fp8MoeBackend.AITER:
         from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-            AiterExperts,
+            AiterFp8BlockExperts,
+            AiterFp8ChannelExperts,
+            AiterFp8TensorExperts,
         )
 
-        return [AiterExperts]
+        return [AiterFp8BlockExperts, AiterFp8TensorExperts, AiterFp8ChannelExperts]
 
     elif backend == Fp8MoeBackend.VLLM_CUTLASS:
         from vllm.model_executor.layers.fused_moe.experts.triton_cutlass_moe import (

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -263,10 +263,10 @@ def backend_to_kernel_cls(
 
     elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
         from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-            AiterExperts,
+            AiterMxfp4Experts,
         )
 
-        return [AiterExperts]
+        return [AiterMxfp4Experts]
 
     elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
         from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
@@ -284,10 +284,10 @@ def backend_to_kernel_cls(
 
     elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
         from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-            AiterExperts,
+            AiterMxfp4Experts,
         )
 
-        return [AiterExperts]
+        return [AiterMxfp4Experts]
 
     elif backend == Mxfp4MoeBackend.XPU:
         from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

@@ -126,10 +126,10 @@ def backend_to_kernel_cls(
 
     elif backend == UnquantizedMoeBackend.AITER:
         from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
-            AiterExperts,
+            AiterUnquantizedExperts,
         )
 
-        return [AiterExperts]
+        return [AiterUnquantizedExperts]
 
     elif backend == UnquantizedMoeBackend.TRITON:
         from vllm.model_executor.layers.fused_moe.experts.triton_moe import (