[ROCm][DSV4] Preserve native MXFP4 TP8 shard allocation (#51473)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Fangzhou Ai committed
Aug 11, 2026 at 02:57 UTC
1044d19998df96408832b33ccc7120716c615969
2 files changed
+63
-3
tests/kernels/moe/test_ocp_mx_moe.py
+57
-3
@@ -2,13 +2,17 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import types
5
-from dataclasses import dataclass
5
+from dataclasses import dataclass, replace
6
7
import pytest
8
import torch
9
10
from tests.kernels.moe.utils import check_accuracy
11
-from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
11
+from vllm._aiter_ops import (
12
+ is_aiter_found,
13
+ is_aiter_found_and_supported,
14
+ rocm_aiter_ops,
15
+)
16
from vllm.platforms import current_platform
17
from vllm.utils.flashinfer import has_flashinfer
18
@@ -1598,8 +1602,58 @@ def test_rocm_mxfp4_moe_oracle(
1602
1603
1604
# -----------------------------------------------------------------------------
1601
-# MXFP4 emulation size-rounding tests
1605
+# MXFP4 size-rounding tests
1606
# -----------------------------------------------------------------------------
1607
+
1608
+
1609
+@pytest.mark.skipif(not ROCM_AVAILABLE, reason="ROCm-specific test")
1610
+@pytest.mark.skipif(
1611
+ not is_aiter_found_and_supported(), reason="AITER is not installed or supported"
1612
+)
1613
+def test_aiter_mxfp4_bf16_silu_preserves_native_tp_shard():
1614
+ """Keep 128-aligned SiLU shards at their native widths."""
1615
+ from vllm.model_executor.layers.fused_moe import FusedMoEConfig
1616
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
1617
+ from vllm.model_executor.layers.fused_moe.config import (
1618
+ FusedMoEParallelConfig,
1619
+ RoutingMethodType,
1620
+ )
1621
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1622
+ Mxfp4MoeBackend,
1623
+ )
1624
+ from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
1625
+
1626
+ moe_parallel_config = replace(
1627
+ FusedMoEParallelConfig.make_no_parallel(),
1628
+ tp_size=8,
1629
+ )
1630
+ moe_config = FusedMoEConfig(
1631
+ num_experts=384,
1632
+ experts_per_token=6,
1633
+ hidden_dim=7168,
1634
+ intermediate_size=3072,
1635
+ num_local_experts=384,
1636
+ num_logical_experts=384,
1637
+ activation=MoEActivation.SILU,
1638
+ device="cpu",
1639
+ routing_method=RoutingMethodType.Renormalize,
1640
+ moe_parallel_config=moe_parallel_config,
1641
+ in_dtype=torch.bfloat16,
1642
+ )
1643
+ method = object.__new__(Mxfp4MoEMethod)
1644
+ method.moe = moe_config
1645
+ method.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16
1646
+
1647
+ rounded_shape = method.maybe_roundup_sizes(
1648
+ hidden_size=7168,
1649
+ intermediate_size_per_partition=moe_config.intermediate_size_per_partition,
1650
+ act_dtype=torch.bfloat16,
1651
+ moe_parallel_config=moe_config.moe_parallel_config,
1652
+ )
1653
+
1654
+ assert rounded_shape == (7168, 384)
1655
+
1656
+
1657
# Emulation needs each per-partition dim rounded up to OCP_MX_BLOCK_SIZE (32);
1658
# a non-block-aligned shard (e.g. GPT-OSS 2880 // 4 = 720) otherwise truncates
1659
# the scale buffer and fails weight loading.
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+6
@@ -651,6 +651,12 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
651
# which would inflate weights and OOM.
652
intermediate_size = round_up(intermediate_size, 128)
653
hidden_size = round_up(hidden_size, 128)
654
+ elif (
655
+ backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SILU
656
+ ):
657
+ # AITER's A16W4 SiLU kernel handles native dimensions aligned to 128.
658
+ intermediate_size = round_up(intermediate_size, 128)
659
+ hidden_size = round_up(hidden_size, 128)
660
elif backend == Mxfp4MoeBackend.EMULATION:
661
# Emulation has no kernel tile; it only needs OCP MX block alignment so the
662
# per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated