@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 77e8b02c8d

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

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>

Fangzhou Ai committed Sep 30, 2026 at 03:07 UTC 77e8b02c8d69f25867227fb9883797996fbfcd7b

8 files changed +204 -11

tests/kernels/moe/test_rocm_aiter_moe.py

+109

@@ -638,6 +638,115 @@ def test_aiter_moe_padding_env_var(

638 assert envs.VLLM_ROCM_MOE_PADDING is moe_padding

639

640

641 +# a4w4 (FP4 activation) opt-in gating test ---------------------------------

642 +

643 +

644 +@pytest.mark.parametrize(

645 + ("env_value", "model_type", "expected"),

646 + [

647 + # Unset/"0" is opt-in-only: a4w4 never turns on by itself, even on

648 + # a validated model type.

649 + (None, "deepseek_v41", False),

650 + ("0", "deepseek_v41", False),

651 + ("1", "deepseek_v41", True),

652 + ("1", "deepseek_v41_text", True),

653 + ],

654 +)

655 +def test_aiter_moe_a4w4_dsv4_is_opt_in(

656 + env_value: str | None,

657 + model_type: str,

658 + expected: bool,

659 + monkeypatch: pytest.MonkeyPatch,

660 +):

661 + """vllm-project/vllm#58819's a4w4 MoE activation path must stay a plain

662 + explicit opt-in (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1), with no

663 + default-on behavior, but is still only permitted on the model types

664 + this flag has been validated for."""

665 + import types

666 +

667 + from tests.kernels.moe.utils import make_dummy_moe_config

668 + from vllm._aiter_ops import rocm_aiter_ops

669 +

670 + _assert_aiter_supported()

671 +

672 + fake_vllm_config = types.SimpleNamespace(

673 + model_config=types.SimpleNamespace(

674 + hf_config=types.SimpleNamespace(model_type=model_type)

675 + )

676 + )

677 +

678 + with monkeypatch.context() as mp:

679 + mp.setenv("VLLM_ROCM_USE_AITER", "1")

680 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

681 + if env_value is None:

682 + mp.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False)

683 + else:

684 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", env_value)

685 + _reload_envs()

686 + rocm_aiter_ops.refresh_env_variables()

687 +

688 + with mock.patch(

689 + "vllm.config.get_current_vllm_config_or_none",

690 + return_value=fake_vllm_config,

691 + ):

692 + moe_config = make_dummy_moe_config(

693 + num_experts=4,

694 + experts_per_token=2,

695 + hidden_dim=256,

696 + intermediate_size=512,

697 + in_dtype=torch.bfloat16,

698 + )

699 +

700 + assert moe_config.use_mxfp4_w4a4_dsv4 is expected

701 +

702 +

703 +@pytest.mark.parametrize("model_type", ["deepseek_v4", "gpt_oss", None])

704 +def test_aiter_moe_a4w4_dsv4_rejects_unvalidated_model_type(

705 + model_type: str | None,

706 + monkeypatch: pytest.MonkeyPatch,

707 +):

708 + """Forcing a4w4 on a model type this flag hasn't been GSM8K/AgentX

709 + validated for must raise: it also flips the MXFP4 weight-shuffle layout

710 + (ATOM SEPARATED vs INTERLEAVE), which silently produces garbled output

711 + on the wrong model rather than just being a slower path — the same bug

712 + class vllm-project/vllm#58819's own second commit fixed for DeepSeek

713 + V4.1 itself."""

714 + import types

715 +

716 + from tests.kernels.moe.utils import make_dummy_moe_config

717 + from vllm._aiter_ops import rocm_aiter_ops

718 +

719 + _assert_aiter_supported()

720 +

721 + fake_vllm_config = types.SimpleNamespace(

722 + model_config=types.SimpleNamespace(

723 + hf_config=types.SimpleNamespace(model_type=model_type)

724 + )

725 + )

726 +

727 + with monkeypatch.context() as mp:

728 + mp.setenv("VLLM_ROCM_USE_AITER", "1")

729 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

730 + mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1")

731 + _reload_envs()

732 + rocm_aiter_ops.refresh_env_variables()

733 +

734 + with (

735 + mock.patch(

736 + "vllm.config.get_current_vllm_config_or_none",

737 + return_value=fake_vllm_config,

738 + ),

739 + pytest.raises(ValueError, match="VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4"),

740 + ):

741 + make_dummy_moe_config(

742 + num_experts=4,

743 + experts_per_token=2,

744 + hidden_dim=256,

745 + intermediate_size=512,

746 + in_dtype=torch.bfloat16,

747 + )

748 +

749 +

750 # Dispatch-policy forwarding test ------------------------------------------

751

752

vllm/_aiter_ops.py

+25

@@ -345,6 +345,7 @@ def _rocm_aiter_fused_moe_impl(

345 shared_w1_scale: torch.Tensor | None = None,

346 shared_w2_scale: torch.Tensor | None = None,

347 shared_expert_id: int = -1,

348 + q_dtype_a: torch.dtype | None = None,

349 ) -> torch.Tensor:

350 has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args(

351 shared_w1,

@@ -378,6 +379,12 @@ def _rocm_aiter_fused_moe_impl(

379 shared_w2_scale=shared_w2_scale,

380 shared_expert_id=shared_expert_id,

381 )

382 + if q_dtype_a is not None:

383 + # DeepSeek V4.1 a4w4 override (use_mxfp4_w4a4_dsv4 in

384 + # rocm_aiter_moe.py). rocm_aiter_ops.fused_moe_supports_quant_dtype_a()

385 + # is checked at config time, so this is only reached on an AITER

386 + # build new enough to accept it.

387 + extra_kwargs["quant_dtype_a"] = q_dtype_a

388

389 return fused_moe(

390 hidden_states,

@@ -435,6 +442,7 @@ def _rocm_aiter_fused_moe_fake(

442 shared_w1_scale: torch.Tensor | None = None,

443 shared_w2_scale: torch.Tensor | None = None,

444 shared_expert_id: int = -1,

445 + q_dtype_a: torch.dtype | None = None,

446 ) -> torch.Tensor:

447 if output_dtype is not None:

448 return torch.empty_like(hidden_states, dtype=output_dtype)

@@ -2485,6 +2493,21 @@ class rocm_aiter_ops:

2493

2494 return "gate_mode" in inspect.signature(fused_moe).parameters

2495

2496 + @classmethod

2497 + @if_aiter_supported

2498 + @functools.cache

2499 + def fused_moe_supports_quant_dtype_a(cls) -> bool:

2500 + """Probe whether the installed aiter.fused_moe accepts `quant_dtype_a`.

2501 +

2502 + Added in https://github.com/ROCm/aiter/pull/5439 (unreleased at

2503 + merge time). Older AITER can't override the activation quant dtype.

2504 + """

2505 + import inspect

2506 +

2507 + from aiter.fused_moe import fused_moe

2508 +

2509 + return "quant_dtype_a" in inspect.signature(fused_moe).parameters

2510 +

2511 @staticmethod

2512 def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool:

2513 """Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract."""

@@ -3005,6 +3028,7 @@ class rocm_aiter_ops:

3028 shared_w1_scale: torch.Tensor | None = None,

3029 shared_w2_scale: torch.Tensor | None = None,

3030 shared_expert_id: int = -1,

3031 + q_dtype_a: torch.dtype | None = None,

3032 ) -> torch.Tensor:

3033 return torch.ops.vllm.rocm_aiter_fused_moe(

3034 hidden_states,

@@ -3036,6 +3060,7 @@ class rocm_aiter_ops:

3060 shared_w1_scale,

3061 shared_w2_scale,

3062 shared_expert_id,

3063 + q_dtype_a,

3064 )

3065

3066 @staticmethod

vllm/envs.py

@@ -141,6 +141,7 @@ if TYPE_CHECKING:

141 VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False

142 VLLM_ROCM_USE_AITER_MOE: bool = True

143 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0

144 + VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: bool | None = None

145 VLLM_ROCM_USE_AITER_MOE_SITUV2: Literal["auto", "a4w4", "a8w4", "a16w4"] = "auto"

146 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

147 VLLM_ROCM_USE_AITER_MLA: bool = True

@@ -1291,6 +1292,13 @@ environment_variables: dict[str, Callable[[], Any]] = {

1292 ["auto", "a4w4", "a8w4", "a16w4", "0", "1"],

1293 case_sensitive=False,

1294 ),

1295 + # Opt-in switch for a4w4 (FP4 activation) MoE on DeepSeek V4.1, AITER

1296 + # MXFP4 backend. Default is a8w4 (FP8); set to "1" to enable a4w4

1297 + # ("true" is not accepted -- only "0"/"1"). Raises if set for other

1298 + # models.

1299 + "VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4": lambda: maybe_convert_bool(

1300 + os.getenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4")

1301 + ),

1302 # MoE sorting dispatch policy for AITER fused MoE kernels.

1303 # 0 = auto (default): single-pass for small batches, multi-pass

1304 # for large batches

vllm/model_executor/layers/fused_moe/config.py

+35

@@ -6,6 +6,7 @@ from typing import Union

7 import torch

9 +import vllm.envs as envs

10 from vllm.config import ParallelConfig, SchedulerConfig

11 from vllm.config.kernel import MoEBackend

12 from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank

@@ -1247,6 +1248,12 @@ class FusedMoEParallelConfig:

1248 )

1249

1250

1251 +# Model types validated for VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1. This flag

1252 +# also changes the MXFP4 weight shuffle layout, so using it on other models

1253 +# can produce garbled output.

1254 +_AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES = ("deepseek_v41", "deepseek_v41_text")

1255 +

1256 +

1257 # Adapted from pplx-kernels tests/all_to_all_utils.py

1258 @dataclass

1259 class FusedMoEConfig:

@@ -1314,6 +1321,11 @@ class FusedMoEConfig:

1321 tp_shard_with_padding: bool = False

1322 rocm_aiter_fmoe_enabled: bool = False

1323 aiter_fmoe_shared_expert_enabled: bool = False

1324 + # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on

1325 + # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected

1326 + # for any other model type. Resolved here, not in the forward path,

1327 + # because get_current_vllm_config() isn't set there.

1328 + use_mxfp4_w4a4_dsv4: bool = False

1329

1330 def __post_init__(self):

1331 from vllm._aiter_ops import rocm_aiter_ops

@@ -1345,6 +1357,29 @@ class FusedMoEConfig:

1357 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

1358 )

1359

1360 + if self.rocm_aiter_fmoe_enabled and envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4:

1361 + from vllm.config import get_current_vllm_config_or_none

1362 +

1363 + vllm_config = get_current_vllm_config_or_none()

1364 + model_type = (

1365 + getattr(vllm_config.model_config.hf_config, "model_type", None)

1366 + if vllm_config is not None

1367 + else None

1368 + )

1369 + if model_type not in _AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES:

1370 + raise ValueError(

1371 + f"VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1 only supports "

1372 + f"model_type in {_AITER_MOE_A4W4_DSV4_VALIDATED_MODEL_TYPES}, "

1373 + f"got {model_type!r}. Unset this env var for this model."

1374 + )

1375 + if not rocm_aiter_ops.fused_moe_supports_quant_dtype_a():

1376 + raise ValueError(

1377 + "VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1 needs an AITER build "

1378 + "with fused_moe(quant_dtype_a=...) support "

1379 + "(ROCm/aiter#5439+). Upgrade AITER or unset this env var."

1380 + )

1381 + self.use_mxfp4_w4a4_dsv4 = True

1382 +

1383 if self.use_mori_kernels:

1384 assert self.rocm_aiter_fmoe_enabled, (

1385 "Mori needs to be used with aiter fused_moe for now."

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

+14 -7

@@ -375,15 +375,13 @@ def rocm_aiter_fused_experts(

375 intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)

376 )

377

378 - # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs

379 - # for interleaved vs separated gate and up weights.

380 - # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by

381 - # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`,

382 - # which always sets `is_guinterleave=True`.

383 - # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling.

378 + # AITER's stage1 GEMM needs gate_mode to match how weights were

379 + # shuffled at load time (oracle/mxfp4.py). gpt-oss uses INTERLEAVE;

380 + # DeepSeek V4.1 a4w4 uses SEPARATED (see below).

381 from aiter.ops.flydsl.moe_common import GateMode

382

383 gate_mode = ""

384 + q_dtype_a = None

385 if activation == MoEActivation.SITU:

386 gate_mode = (

387 GateMode.INTERLEAVE.value

@@ -391,7 +389,15 @@ def rocm_aiter_fused_experts(

389 else GateMode.SEPARATED.value

390 )

391 elif quant_config.use_mxfp4_w4a16:

394 - gate_mode = GateMode.INTERLEAVE.value

392 + if moe_config.use_mxfp4_w4a4_dsv4:

393 + # Opt-in a4w4 for DeepSeek V4.1 (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1).

394 + # Must match the SEPARATED weight shuffle from load time.

395 + from aiter import dtypes

396 +

397 + gate_mode = GateMode.SEPARATED.value

398 + q_dtype_a = dtypes.fp4x2

399 + else:

400 + gate_mode = GateMode.INTERLEAVE.value

401 elif activation_interleave is not None:

402 gate_mode = (

403 GateMode.INTERLEAVE.value

@@ -433,6 +439,7 @@ def rocm_aiter_fused_experts(

439 shared_w1_scale=shared_w1_scale,

440 shared_w2_scale=shared_w2_scale,

441 shared_expert_id=shared_expert_id,

442 + q_dtype_a=q_dtype_a,

443 )

444

445

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+11 -4

@@ -1382,6 +1382,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1382 w2_bias: torch.Tensor | None = None,

1383 _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,

1384 activation: MoEActivation | None = None,

1385 + use_separated_a4w4: bool = False,

1386 ) -> tuple[

1387 torch.Tensor,

1388 torch.Tensor,

@@ -1626,10 +1627,16 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1627 from aiter.ops.shuffle import shuffle_scale as _shuf_s

1628 from aiter.ops.shuffle import shuffle_weight as _shuf_w

1629

1630 + # DeepSeek V4.1 a4w4 uses ATOM's SEPARATED gate/up layout instead of

1631 + # the default INTERLEAVE shuffle (INTERLEAVE + fp4x2 has no tuned

1632 + # kernel and produces garbage output). Must match GateMode.SEPARATED

1633 + # in rocm_aiter_moe.py.

1634 + is_guinterleave = not use_separated_a4w4

1635 +

1636 w13_weight = torch.nn.Parameter(

1637 _shuf_w(

1638 w13_weight.data.view(torch.float4_e2m1fn_x2),

1632 - is_guinterleave=True,

1639 + is_guinterleave=is_guinterleave,

1640 gate_up=True,

1641 ),

1642 requires_grad=False,

@@ -1637,14 +1644,14 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1644 shuffled_w13_scale = _shuf_s(

1645 w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]),

1646 num_experts,

1640 - True,

1647 + is_guinterleave,

1648 True,

1649 )

1650

1651 w2_weight = torch.nn.Parameter(

1652 _shuf_w(

1653 w2_weight.data.view(torch.float4_e2m1fn_x2),

1647 - is_guinterleave=True,

1654 + is_guinterleave=is_guinterleave,

1655 gate_up=False,

1656 ),

1657 requires_grad=False,

@@ -1653,7 +1660,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(

1660 shuffled_w2_scale = _shuf_s(

1661 w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]),

1662 num_experts,

1656 - True,

1663 + is_guinterleave,

1664 False,

1665 )

1666

vllm/model_executor/layers/quantization/mxfp4.py

@@ -736,6 +736,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

736 w2_bias=w2_bias,

737 _cache_permute_indices=self._cache_permute_indices,

738 activation=self.moe.activation,

739 + use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4,

740 )

741 )

742

vllm/model_executor/layers/quantization/online/mxfp4.py

@@ -220,6 +220,7 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):

220 w2_weight_scale=w2_scale,

221 w13_bias=w13_bias,

222 w2_bias=w2_bias,

223 + use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4,

224 )

225 )

226