[Fix][ROCm] MXFP4 MoE round-up inflates TP-sharded expert weights on CDNA3, starving KV cache (#56359)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
vllmellm committed
Sep 17, 2026 at 22:52 UTC
3affd35f3164ba17503ba63cd8733becd1b03bf1
1 file changed
+25
-12
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+25
-12
@@ -769,18 +769,31 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
769
intermediate_size = round_up(intermediate_size, 128)
770
hidden_size = round_up(hidden_size, 128)
771
elif current_platform.is_rocm():
772
- if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and (
773
- activation == MoEActivation.SITU or activation == MoEActivation.SILU
774
- ):
775
- # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size
776
- # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a
777
- # no-op for K3's shapes) rather than the generic ROCm 256 round-up,
778
- # which would inflate weights and OOM.
779
- intermediate_size = round_up(intermediate_size, 128)
780
- hidden_size = round_up(hidden_size, 128)
781
- else:
782
- intermediate_size = round_up(intermediate_size, 256)
783
- hidden_size = round_up(hidden_size, 256)
772
+ from vllm.platforms.rocm import get_cdna_version
773
+
774
+ is_situ_or_silu = activation in (
775
+ MoEActivation.SITU,
776
+ MoEActivation.SILU,
777
+ )
778
+
779
+ # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size
780
+ # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128
781
+ # rather than the generic ROCm 256 round-up, which would inflate
782
+ # weights and OOM.
783
+ aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
784
+
785
+ # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs.
786
+ # CDNA4's F16xMXFP4 configuration uses block_k=256.
787
+ triton_uses_128 = (
788
+ backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4
789
+ )
790
+
791
+ alignment = (
792
+ 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256
793
+ )
794
+
795
+ intermediate_size = round_up(intermediate_size, alignment)
796
+ hidden_size = round_up(hidden_size, alignment)
797
elif backend == Mxfp4MoeBackend.CPU:
798
# CPU AMX kernel uses BLOCK_N=32, align to 32
799
intermediate_size = round_up(intermediate_size, 32)