[ROCm][Bugfix] Fix GPT-OSS Quark MXFP4 MoE loading - emulation buffer not block-aligned (#43979)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Jul 18, 2026 at 11:49 UTC
f12b80c6efa116c7c4020ee611d622ab7030ca40
3 files changed
+59
-7
tests/kernels/moe/test_ocp_mx_moe.py
+46
@@ -1512,3 +1512,49 @@ def test_rocm_mxfp4_moe_oracle(
1512
1513
# Check accuracy using per-backend thresholds
1514
check_accuracy(ref, out, atol=0.1, rtol=config["rtol"], percent=config["percent"])
1515
+
1516
+
1517
+# -----------------------------------------------------------------------------
1518
+# MXFP4 emulation size-rounding tests
1519
+# -----------------------------------------------------------------------------
1520
+# Emulation needs each per-partition dim rounded up to OCP_MX_BLOCK_SIZE (32);
1521
+# a non-block-aligned shard (e.g. GPT-OSS 2880 // 4 = 720) otherwise truncates
1522
+# the scale buffer and fails weight loading.
1523
+# NOTE: gated to ROCm since it is the emulation backend's current target;
1524
+# remove this skip if the backend is enabled on non-ROCm platforms.
1525
+@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm")
1526
+@pytest.mark.parametrize(
1527
+ "hidden_size,intermediate_size,expected_hidden,expected_intermediate",
1528
+ [
1529
+ (2880, 720, 2880, 736), # GPT-OSS TP=4 shard: 720 -> round_up(720, 32)
1530
+ (2880, 360, 2880, 384), # GPT-OSS TP=8 shard: 360 -> 384
1531
+ (2880, 2880, 2880, 2880), # already block-aligned: unchanged
1532
+ (90, 90, 96, 96), # both dims unaligned
1533
+ ],
1534
+)
1535
+def test_mxfp4_emulation_rounds_up_to_block_size(
1536
+ hidden_size: int,
1537
+ intermediate_size: int,
1538
+ expected_hidden: int,
1539
+ expected_intermediate: int,
1540
+):
1541
+ """Emulation must block-align per-partition dims to OCP_MX_BLOCK_SIZE."""
1542
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1543
+ Mxfp4MoeBackend,
1544
+ mxfp4_round_up_hidden_size_and_intermediate_size,
1545
+ )
1546
+ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
1547
+ OCP_MX_BLOCK_SIZE,
1548
+ )
1549
+
1550
+ rounded_hidden, rounded_intermediate = (
1551
+ mxfp4_round_up_hidden_size_and_intermediate_size(
1552
+ Mxfp4MoeBackend.EMULATION, hidden_size, intermediate_size
1553
+ )
1554
+ )
1555
+
1556
+ assert rounded_hidden == expected_hidden
1557
+ assert rounded_intermediate == expected_intermediate
1558
+ # The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate.
1559
+ assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0
1560
+ assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+10
-1
@@ -27,6 +27,9 @@ from vllm.model_executor.layers.fused_moe.config import (
27
ocp_mx_moe_quant_config,
28
)
29
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
30
+from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
31
+ OCP_MX_BLOCK_SIZE,
32
+)
33
from vllm.model_executor.layers.quantization.utils.quant_utils import (
34
QuantKey,
35
kFp8Dynamic128Sym,
@@ -633,7 +636,13 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
636
backend: Mxfp4MoeBackend, hidden_size: int, intermediate_size: int
637
) -> tuple[int, int]:
638
"""Round up hidden_size and intermediate_size based on backend requirements."""
636
- if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
639
+ if backend == Mxfp4MoeBackend.EMULATION:
640
+ # Emulation has no kernel tile; it only needs OCP MX block alignment so the
641
+ # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated
642
+ # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720).
643
+ intermediate_size = round_up(intermediate_size, OCP_MX_BLOCK_SIZE)
644
+ hidden_size = round_up(hidden_size, OCP_MX_BLOCK_SIZE)
645
+ elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
646
# DeepGEMM requires M/N/K alignment
647
intermediate_size = round_up(intermediate_size, 128)
648
hidden_size = round_up(hidden_size, 128)
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3
-6
@@ -1063,12 +1063,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1063
act_dtype=act_dtype,
1064
moe_parallel_config=moe_parallel_config,
1065
)
1066
- # In case quantization emulation backend is used, there is no need to apply
1067
- # MXFP4-specific padding logic as the compute happens in higher precision.
1068
- if (
1069
- self.mxfp4_backend is not None
1070
- and self.mxfp4_backend != Mxfp4MoeBackend.EMULATION
1071
- ):
1066
+ # Round per-partition sizes up to each backend's requirement. Emulation is
1067
+ # handled inside the helper too (OCP MX block alignment), so no special-case.
1068
+ if self.mxfp4_backend is not None:
1069
hidden_size, intermediate_size_per_partition = (
1070
mxfp4_round_up_hidden_size_and_intermediate_size(
1071
self.mxfp4_backend, hidden_size, intermediate_size_per_partition