diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index e4ae7f897e35..92893748e20f 100644
--- a/.buildkite/test-amd.yaml
+++ b/.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
diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md
index 6c71aa1f46d4..2528bc9ea459 100644
--- a/docs/design/moe_kernel_features.md
+++ b/docs/design/moe_kernel_features.md
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
| marlin | standard,batched | 3 / N/A | 3 / N/A | silu,swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] |
| trtllm | standard | mxfp4,nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],[`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,fp8 | G(32),G(128),A,T | silu, gelu,swigluoai | Y | N | `rocm_aiter_fused_experts`,`AiterExperts` |
+| rocm aiter moe | standard | none,mxfp4,fp8 | G(32),G(128),A,T | silu, gelu,swigluoai,situ | Y | Y | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.rocm_aiter_fused_experts],[`AiterUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterUnquantizedExperts],[`AiterFp8BlockExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8BlockExperts],[`AiterFp8TensorExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8TensorExperts],[`AiterFp8ChannelExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8ChannelExperts],[`AiterMxfp4Experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterMxfp4Experts] |
| cpu_moe | standard | N/A | N/A | silu, gelu,gelu_tanh,swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] |
| naive batched4 | batched | int8,fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py
index 5f6349a32880..1f79b62bf8c8 100644
--- a/tests/kernels/moe/modular_kernel_tools/common.py
+++ b/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)
diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
index 6ca644c5fce1..56119ffddeb7 100644
--- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py
+++ b/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(
diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py
index f29c22085240..da643befecd9 100644
--- a/tests/kernels/moe/test_modular_kernel_combinations.py
+++ b/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]
diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py
index 9869b6bc5467..f06afc8bd63e 100644
--- a/tests/kernels/moe/test_mori_moe.py
+++ b/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(
diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py
index abd46753515e..af2f7beb1bd5 100644
--- a/tests/kernels/moe/test_rocm_aiter_moe.py
+++ b/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")
diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py
index bc438c66931b..a6f7f3080583 100644
--- a/tests/kernels/moe/test_topk_softplus_sqrt.py
+++ b/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"),
],
)
diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py
index 120fbd152668..e68310f98d0d 100644
--- a/vllm/model_executor/layers/fused_moe/__init__.py
+++ b/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",
diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
index 07eddbf4a092..5c506eae123a 100644
--- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+++ b/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()
diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py
index 446825c59047..074acd2a0942 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py
+++ b/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 (
diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
index cbbc39b94310..88e80c755d1f 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
@@ -284,10 +284,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 (
@@ -305,10 +305,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
diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py
index 25b55670d1f1..1af0450dfd13 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+++ b/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 (