From e15dff8d8ba85b7e3e742e3c67a03021f83aa7d9 Mon Sep 17 00:00:00 2001 From: jiaxwang Date: Tue, 21 Jul 2026 06:37:14 +0000 Subject: [PATCH] [ROCm] Support MiniMax-M3 NVFP4 SwiGLU-OAI Signed-off-by: jiaxwang --- tests/quantization/test_quark.py | 126 ++++++++++++++++++ .../model_executor/layers/fused_moe/config.py | 4 + .../layers/fused_moe/oracle/nvfp4.py | 5 + .../layers/quantization/quark/quark_moe.py | 3 + 4 files changed, 138 insertions(+) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 9622944670cf..28d21e1ce487 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -69,6 +69,132 @@ def enable_pickle(monkeypatch): monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") +def test_quark_nvfp4_moe_forwards_swiglu_params(): + from unittest.mock import MagicMock, patch + + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.quantization.quark import quark_moe + from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkNvfp4MoEMethod, + ) + + moe_config = FusedMoEConfig( + num_experts=2, + experts_per_token=1, + hidden_dim=16, + intermediate_size=32, + num_local_experts=2, + num_logical_experts=2, + activation=MoEActivation.SWIGLUOAI, + device="cpu", + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + max_num_tokens=8, + intermediate_size_per_partition=32, + ) + + layer = torch.nn.Module() + layer.w13_weight_scale = torch.ones(2) + layer.w2_weight_scale = torch.ones(2) + layer.w13_weight_scale_2 = torch.ones(2) + layer.w2_weight_scale_2 = torch.ones(2) + layer.w13_input_scale_2 = torch.ones(2) + layer.w2_input_scale_2 = torch.ones(2) + layer.swiglu_limit = 7.0 + layer.swiglu_alpha = 1.702 + layer.swiglu_beta = 1.0 + + make_quant_config = MagicMock(return_value=object()) + with ( + patch.object( + quark_moe, + "select_nvfp4_moe_backend", + return_value=(object(), object()), + ), + patch.object( + quark_moe, + "make_nvfp4_moe_quant_config", + make_quant_config, + ), + ): + method = QuarkNvfp4MoEMethod({}, {}, moe_config, MagicMock()) + method.get_fused_moe_quant_config(layer) + + _, kwargs = make_quant_config.call_args + assert kwargs["swiglu_limit"] == 7.0 + assert kwargs["swiglu_alpha"] == 1.702 + assert kwargs["swiglu_beta"] == 1.0 + + +def test_quark_nvfp4_swiglu_limit_allows_emulation_backend(monkeypatch): + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.oracle import nvfp4 + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + select_nvfp4_moe_backend, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kNvfp4Dynamic, + kNvfp4Static, + ) + + class UnsupportedExperts: + @staticmethod + def is_supported_config(cls, config, weight_key, activation_key, fmt): + return False, "unsupported" + + class SupportedExperts: + @staticmethod + def is_supported_config(cls, config, weight_key, activation_key, fmt): + return True, None + + def backend_to_kernel_cls(backend): + if backend == NvFp4MoeBackend.EMULATION: + return [SupportedExperts] + return [UnsupportedExperts] + + config = FusedMoEConfig( + num_experts=2, + experts_per_token=1, + hidden_dim=16, + intermediate_size=32, + num_local_experts=2, + num_logical_experts=2, + activation=MoEActivation.SWIGLUOAI, + device="cpu", + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + max_num_tokens=8, + swiglu_limit=7.0, + intermediate_size_per_partition=32, + ) + + monkeypatch.setattr(nvfp4.envs, "is_set", lambda _: False) + monkeypatch.setattr(nvfp4.envs, "VLLM_TEST_FORCE_FP8_MARLIN", False) + monkeypatch.setattr(nvfp4, "backend_to_kernel_cls", backend_to_kernel_cls) + + backend, experts_cls = select_nvfp4_moe_backend( + config=config, + weight_key=kNvfp4Static, + activation_key=kNvfp4Dynamic, + ) + + assert backend == NvFp4MoeBackend.EMULATION + assert experts_cls is SupportedExperts + + @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) @pytest.mark.parametrize("tp", [1]) def test_quark_fp8_w_per_tensor_a_per_tensor(vllm_runner, kv_cache_dtype, tp): diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 02e6c51116d1..47ebc201bf08 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -815,6 +815,8 @@ def nvfp4_moe_quant_config( w2_bias: torch.Tensor | None = None, is_scale_swizzled: bool = True, gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for mxfp4 activations and nvp4 weights. @@ -834,6 +836,8 @@ def nvfp4_moe_quant_config( block_shape=None, is_scale_swizzled=is_scale_swizzled, gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 7852b4db8474..7a150d500353 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -191,6 +191,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.MARLIN, + NvFp4MoeBackend.EMULATION, } if config.swiglu_limit is not None: @@ -468,6 +469,8 @@ def make_nvfp4_moe_quant_config( a13_scale: torch.Tensor, a2_scale: torch.Tensor, swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: if backend == NvFp4MoeBackend.HUMMING: @@ -495,6 +498,8 @@ def make_nvfp4_moe_quant_config( w1_scale=w13_scale, w2_scale=w2_scale, gemm1_clamp_limit=swiglu_limit, + gemm1_alpha=swiglu_alpha, + gemm1_beta=swiglu_beta, ) # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas. diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 15023d7ca39a..98fa6baf427d 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1575,6 +1575,9 @@ def get_fused_moe_quant_config( w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale_2, a2_scale=layer.w2_input_scale_2, + swiglu_limit=getattr(layer, "swiglu_limit", None), + swiglu_alpha=getattr(layer, "swiglu_alpha", None), + swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, )