diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index fbc8586fd3..05692852e1 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -638,6 +638,115 @@ def test_aiter_moe_padding_env_var( assert envs.VLLM_ROCM_MOE_PADDING is moe_padding +# a4w4 (FP4 activation) opt-in gating test --------------------------------- + + +@pytest.mark.parametrize( + ("env_value", "model_type", "expected"), + [ + # Unset/"0" is opt-in-only: a4w4 never turns on by itself, even on + # a validated model type. + (None, "deepseek_v41", False), + ("0", "deepseek_v41", False), + ("1", "deepseek_v41", True), + ("1", "deepseek_v41_text", True), + ], +) +def test_aiter_moe_a4w4_dsv4_is_opt_in( + env_value: str | None, + model_type: str, + expected: bool, + monkeypatch: pytest.MonkeyPatch, +): + """vllm-project/vllm#58819's a4w4 MoE activation path must stay a plain + explicit opt-in (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1), with no + default-on behavior, but is still only permitted on the model types + this flag has been validated for.""" + import types + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + fake_vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace( + hf_config=types.SimpleNamespace(model_type=model_type) + ) + ) + + with monkeypatch.context() as mp: + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + if env_value is None: + mp.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) + else: + mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", env_value) + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + with mock.patch( + "vllm.config.get_current_vllm_config_or_none", + return_value=fake_vllm_config, + ): + moe_config = make_dummy_moe_config( + num_experts=4, + experts_per_token=2, + hidden_dim=256, + intermediate_size=512, + in_dtype=torch.bfloat16, + ) + + assert moe_config.use_mxfp4_w4a4_dsv4 is expected + + +@pytest.mark.parametrize("model_type", ["deepseek_v4", "gpt_oss", None]) +def test_aiter_moe_a4w4_dsv4_rejects_unvalidated_model_type( + model_type: str | None, + monkeypatch: pytest.MonkeyPatch, +): + """Forcing a4w4 on a model type this flag hasn't been GSM8K/AgentX + validated for must raise: it also flips the MXFP4 weight-shuffle layout + (ATOM SEPARATED vs INTERLEAVE), which silently produces garbled output + on the wrong model rather than just being a slower path — the same bug + class vllm-project/vllm#58819's own second commit fixed for DeepSeek + V4.1 itself.""" + import types + + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + fake_vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace( + hf_config=types.SimpleNamespace(model_type=model_type) + ) + ) + + with monkeypatch.context() as mp: + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + with ( + mock.patch( + "vllm.config.get_current_vllm_config_or_none", + return_value=fake_vllm_config, + ), + pytest.raises(ValueError, match="VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4"), + ): + make_dummy_moe_config( + num_experts=4, + experts_per_token=2, + hidden_dim=256, + intermediate_size=512, + in_dtype=torch.bfloat16, + ) + + # Dispatch-policy forwarding test ------------------------------------------ diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 906349af9a..f2d313bdd0 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -345,6 +345,7 @@ def _rocm_aiter_fused_moe_impl( shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + q_dtype_a: torch.dtype | None = None, ) -> torch.Tensor: has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args( shared_w1, @@ -378,6 +379,12 @@ def _rocm_aiter_fused_moe_impl( shared_w2_scale=shared_w2_scale, shared_expert_id=shared_expert_id, ) + if q_dtype_a is not None: + # DeepSeek V4.1 a4w4 override (use_mxfp4_w4a4_dsv4 in + # rocm_aiter_moe.py). rocm_aiter_ops.fused_moe_supports_quant_dtype_a() + # is checked at config time, so this is only reached on an AITER + # build new enough to accept it. + extra_kwargs["quant_dtype_a"] = q_dtype_a return fused_moe( hidden_states, @@ -435,6 +442,7 @@ def _rocm_aiter_fused_moe_fake( shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + q_dtype_a: torch.dtype | None = None, ) -> torch.Tensor: if output_dtype is not None: return torch.empty_like(hidden_states, dtype=output_dtype) @@ -2485,6 +2493,21 @@ class rocm_aiter_ops: return "gate_mode" in inspect.signature(fused_moe).parameters + @classmethod + @if_aiter_supported + @functools.cache + def fused_moe_supports_quant_dtype_a(cls) -> bool: + """Probe whether the installed aiter.fused_moe accepts `quant_dtype_a`. + + Added in https://github.com/ROCm/aiter/pull/5439 (unreleased at + merge time). Older AITER can't override the activation quant dtype. + """ + import inspect + + from aiter.fused_moe import fused_moe + + return "quant_dtype_a" in inspect.signature(fused_moe).parameters + @staticmethod def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool: """Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract.""" @@ -3005,6 +3028,7 @@ class rocm_aiter_ops: shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + q_dtype_a: torch.dtype | None = None, ) -> torch.Tensor: return torch.ops.vllm.rocm_aiter_fused_moe( hidden_states, @@ -3036,6 +3060,7 @@ class rocm_aiter_ops: shared_w1_scale, shared_w2_scale, shared_expert_id, + q_dtype_a, ) @staticmethod diff --git a/vllm/envs.py b/vllm/envs.py index 428d6d3c94..67e0207492 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -141,6 +141,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 + VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: bool | None = None VLLM_ROCM_USE_AITER_MOE_SITUV2: Literal["auto", "a4w4", "a8w4", "a16w4"] = "auto" VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True @@ -1291,6 +1292,13 @@ environment_variables: dict[str, Callable[[], Any]] = { ["auto", "a4w4", "a8w4", "a16w4", "0", "1"], case_sensitive=False, ), + # Opt-in switch for a4w4 (FP4 activation) MoE on DeepSeek V4.1, AITER + # MXFP4 backend. Default is a8w4 (FP8); set to "1" to enable a4w4 + # ("true" is not accepted -- only "0"/"1"). Raises if set for other + # models. + "VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4": lambda: maybe_convert_bool( + os.getenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4") + ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass # for large batches diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 11e2a57520..302ed121fe 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -6,6 +6,7 @@ from typing import Union import torch +import vllm.envs as envs from vllm.config import ParallelConfig, SchedulerConfig from vllm.config.kernel import MoEBackend from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank @@ -1247,6 +1248,12 @@ class FusedMoEParallelConfig: ) +# Model types validated for VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1. This flag +# also changes the MXFP4 weight shuffle layout, so using it on other models +# can produce garbled output. +_AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES = ("deepseek_v41", "deepseek_v41_text") + + # Adapted from pplx-kernels tests/all_to_all_utils.py @dataclass class FusedMoEConfig: @@ -1314,6 +1321,11 @@ class FusedMoEConfig: tp_shard_with_padding: bool = False rocm_aiter_fmoe_enabled: bool = False aiter_fmoe_shared_expert_enabled: bool = False + # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on + # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected + # for any other model type. Resolved here, not in the forward path, + # because get_current_vllm_config() isn't set there. + use_mxfp4_w4a4_dsv4: bool = False def __post_init__(self): from vllm._aiter_ops import rocm_aiter_ops @@ -1345,6 +1357,29 @@ class FusedMoEConfig: rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() ) + if self.rocm_aiter_fmoe_enabled and envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + model_type = ( + getattr(vllm_config.model_config.hf_config, "model_type", None) + if vllm_config is not None + else None + ) + if model_type not in _AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES: + raise ValueError( + f"VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1 only supports " + f"model_type in {_AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES}, " + f"got {model_type!r}. Unset this env var for this model." + ) + if not rocm_aiter_ops.fused_moe_supports_quant_dtype_a(): + raise ValueError( + "VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1 needs an AITER build " + "with fused_moe(quant_dtype_a=...) support " + "(ROCm/aiter#5439+). Upgrade AITER or unset this env var." + ) + self.use_mxfp4_w4a4_dsv4 = True + if self.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( "Mori needs to be used with aiter fused_moe for now." 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 4f2f760ad3..07eddbf4a0 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 @@ -375,15 +375,13 @@ def rocm_aiter_fused_experts( intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) ) - # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs - # for interleaved vs separated gate and up weights. - # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by - # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`, - # which always sets `is_guinterleave=True`. - # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling. + # AITER's stage1 GEMM needs gate_mode to match how weights were + # shuffled at load time (oracle/mxfp4.py). gpt-oss uses INTERLEAVE; + # DeepSeek V4.1 a4w4 uses SEPARATED (see below). from aiter.ops.flydsl.moe_common import GateMode gate_mode = "" + q_dtype_a = None if activation == MoEActivation.SITU: gate_mode = ( GateMode.INTERLEAVE.value @@ -391,7 +389,15 @@ def rocm_aiter_fused_experts( else GateMode.SEPARATED.value ) elif quant_config.use_mxfp4_w4a16: - gate_mode = GateMode.INTERLEAVE.value + if moe_config.use_mxfp4_w4a4_dsv4: + # Opt-in a4w4 for DeepSeek V4.1 (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1). + # Must match the SEPARATED weight shuffle from load time. + from aiter import dtypes + + gate_mode = GateMode.SEPARATED.value + q_dtype_a = dtypes.fp4x2 + else: + gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: gate_mode = ( GateMode.INTERLEAVE.value @@ -433,6 +439,7 @@ def rocm_aiter_fused_experts( shared_w1_scale=shared_w1_scale, shared_w2_scale=shared_w2_scale, shared_expert_id=shared_expert_id, + q_dtype_a=q_dtype_a, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 56ba3981b7..27d59ba2e0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1382,6 +1382,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias: torch.Tensor | None = None, _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, activation: MoEActivation | None = None, + use_separated_a4w4: bool = False, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -1626,10 +1627,16 @@ def convert_weight_to_mxfp4_moe_kernel_format( from aiter.ops.shuffle import shuffle_scale as _shuf_s from aiter.ops.shuffle import shuffle_weight as _shuf_w + # DeepSeek V4.1 a4w4 uses ATOM's SEPARATED gate/up layout instead of + # the default INTERLEAVE shuffle (INTERLEAVE + fp4x2 has no tuned + # kernel and produces garbage output). Must match GateMode.SEPARATED + # in rocm_aiter_moe.py. + is_guinterleave = not use_separated_a4w4 + w13_weight = torch.nn.Parameter( _shuf_w( w13_weight.data.view(torch.float4_e2m1fn_x2), - is_guinterleave=True, + is_guinterleave=is_guinterleave, gate_up=True, ), requires_grad=False, @@ -1637,14 +1644,14 @@ def convert_weight_to_mxfp4_moe_kernel_format( shuffled_w13_scale = _shuf_s( w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]), num_experts, - True, + is_guinterleave, True, ) w2_weight = torch.nn.Parameter( _shuf_w( w2_weight.data.view(torch.float4_e2m1fn_x2), - is_guinterleave=True, + is_guinterleave=is_guinterleave, gate_up=False, ), requires_grad=False, @@ -1653,7 +1660,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( shuffled_w2_scale = _shuf_s( w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]), num_experts, - True, + is_guinterleave, False, ) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index dcba98c7e4..f9c7f37ff3 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -736,6 +736,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w2_bias=w2_bias, _cache_permute_indices=self._cache_permute_indices, activation=self.moe.activation, + use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4, ) ) diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index 0115988e92..72914a9726 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -220,6 +220,7 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): w2_weight_scale=w2_scale, w13_bias=w13_bias, w2_bias=w2_bias, + use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4, ) )