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

.buildkite/test-amd.yaml

@@ -1528,7 +1528,32 @@ steps:
   commands:
   - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+
+- label: ":amd: (MI300) MoE Kernels Multi-GPU" # TBD
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - csrc/rocm/
+  - tests/kernels/moe/test_modular_kernel_combinations.py
+  - tests/kernels/moe/modular_kernel_tools/
+  - tests/kernels/utils.py
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/distributed/device_communicators/
+  - vllm/distributed/parallel_state.py
+  - vllm/envs.py
+  - vllm/config
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "multigpu"
 
 - label: ":amd: (MI300) vLLM IR" # TBD
   timeout_in_minutes: 50

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/common.py

@@ -167,6 +167,9 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
         vllm_config.model_config = SimpleNamespace(
             enforce_eager=True,
             is_moe=True,
+            enable_sleep_mode=False,
+            sleep_mode_backend="cumem",
+            sleep_mode_offload_cudagraph=False,
         )
         vllm_config.parallel_config.data_parallel_size = self.world_size
         vllm_config.parallel_config.enable_expert_parallel = True
@@ -722,6 +725,9 @@ def _maybe_convert_weights_for_experts(
     rank_weights: WeightTensors,
 ) -> WeightTensors:
     """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK)."""
+    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+        AiterExperts,
+    )
     from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
         Fp8MoeBackend,
         convert_to_fp8_moe_kernel_format,
@@ -734,7 +740,7 @@ def _maybe_convert_weights_for_experts(
     # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor
     # and per-token schemes. Those keep JIT-compiling until
     # https://github.com/ROCm/aiter/issues/5766 is fixed.
-    if fe_name == "AiterExperts" and (
+    if issubclass(fe_type, AiterExperts) and (
         config.quant_dtype is None or config.quant_block_shape is not None
     ):
         return _shuffle_weights_for_aiter(rank_weights)

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,13 @@
 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,
+    AiterFp8BlockExperts,
+    AiterFp8ChannelExperts,
+    AiterFp8TensorExperts,
+    AiterUnquantizedExperts,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8DynamicTensorSym,
     kFp8StaticTensorSym,
@@ -80,10 +86,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()
@@ -101,7 +107,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), (
@@ -156,7 +162,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(
@@ -172,7 +178,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")
             ):
@@ -210,7 +216,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:
@@ -270,7 +276,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
@@ -591,7 +597,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]
@@ -671,7 +677,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]
@@ -738,6 +744,20 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
 ]
 
 
+def _padding_matrix_fe_cls(
+    quant_config: TestMoEQuantConfig | None,
+) -> type[AiterExperts]:
+    """Map a _PADDING_QUANT_CONFIGS entry to the AiterExperts leaf class that
+    supports it (mirrors Config.fp8_quant_key_pair's scheme derivation)."""
+    if quant_config is None:
+        return AiterUnquantizedExperts
+    if quant_config.block_shape is not None:
+        return AiterFp8BlockExperts
+    if quant_config.per_out_ch_quant:
+        return AiterFp8ChannelExperts
+    return AiterFp8TensorExperts
+
+
 def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor:
     """Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real
     gate/up halves of `real_per_half` size each."""
@@ -950,7 +970,7 @@ def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig |
         dtype=torch.bfloat16,
         quant_config=quant_config,
         prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
-        fused_experts_type=AiterExperts,
+        fused_experts_type=_padding_matrix_fe_cls(quant_config),
         world_size=1,
         hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None,
         intermediate_size_per_partition_unpadded=(
@@ -1113,7 +1133,7 @@ def test_aiter_moe_token_padding_garbage_rows(
         dtype=torch.bfloat16,
         quant_config=quant_config,
         prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
-        fused_experts_type=AiterExperts,
+        fused_experts_type=_padding_matrix_fe_cls(quant_config),
         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",