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 (