iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Add opt-in a4w4 (FP4 activation) MoE for DeepSeek V4.1 on AITER (#58819)

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
Signed-off-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com>
Co-authored-by: Cursor Agent (Claude Sonnet 5) <cursoragent@cursor.com>
This commit is contained in:
Fangzhou Ai
2026-09-30 15:07:44 +08:00
committed by GitHub
co-authored by Cursor Agent
parent e5f3d08d72
commit 77e8b02c8d
8 changed files with 204 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_rocm_aiter_moe.py
+109
View File
@@ -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 ------------------------------------------
vllm/_aiter_ops.py
+25
View File
@@ -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
vllm/envs.py
+8
View File
@@ -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
vllm/model_executor/layers/fused_moe/config.py
+35
View File
@@ -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."
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+14 -7
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+11 -4
View File
@@ -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,
)
vllm/model_executor/layers/quantization/mxfp4.py
+1
View File
@@ -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,
)
)
vllm/model_executor/layers/quantization/online/mxfp4.py
+1
View File
@@ -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,
)
)