A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][gpt-oss] Hybrid CDNA4 swizzle gate for A8W4 MoE (#44804)
Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>
xiaohuguo2023 committed
Jun 10, 2026 at 07:59 UTC
bb78168b210c4ff84cd02cd3412e21e595b0f603
2 files changed
+23
-5
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+8
-2
@@ -113,6 +113,12 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
113
from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
114
from aiter.ops.triton.quant_moe import downcast_to_static_fp8
115
116
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
117
+ should_use_cdna4_mx_scale_swizzle,
118
+ )
119
+
120
+ _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
121
+
122
assert quant_config.w1_precision is not None, (
123
"w1_precision in quant config can't be None"
124
)
@@ -135,7 +141,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
141
routing_data,
142
gather_indx=gather_indx,
143
gammas=gammas if apply_router_weight_on_input else None,
138
- swizzle_mx_scale="CDNA4_SCALE",
144
+ swizzle_mx_scale=_swizzle_mx_scale,
145
out_dtype=torch.float8_e4m3fn,
146
apply_swiglu=True,
147
alpha=swiglu_alpha,
@@ -155,7 +161,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
161
routing_data,
162
scatter_indx=scatter_indx,
163
gammas=None if apply_router_weight_on_input else gammas,
158
- swizzle_mx_scale="CDNA4_SCALE",
164
+ swizzle_mx_scale=_swizzle_mx_scale,
165
unpadded_N=unpadded_N_w2,
166
unpadded_K=unpadded_K_w2,
167
)
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+15
-3
@@ -19,6 +19,20 @@ logger = init_logger(__name__)
19
CK_MXFP4_MOE_DIM_ALIGNMENT = 256
20
21
22
+def should_use_cdna4_mx_scale_swizzle() -> bool:
23
+ """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950.
24
+
25
+ CDNA4 swizzle requires BLOCK_K%256==0; at TP>=4 the A8W4 dispatch
26
+ picks BK<256 tiles for the smaller per-rank shapes, so swizzle must
27
+ be off. Used by both the weight-load swizzle in `_swizzle_mxfp4` and
28
+ the kernel-argument gate in `aiter_mxfp4_w4a8_moe`; they must agree.
29
+ """
30
+ from vllm.distributed import get_tensor_model_parallel_world_size
31
+ from vllm.platforms.rocm import on_gfx950
32
+
33
+ return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
34
+
35
+
36
def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
37
"""weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
38
assert has_triton_kernels()
@@ -44,10 +58,8 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
58
value_layout = StridedLayout
59
scale_layout = StridedLayout
60
elif current_platform.is_rocm():
47
- from vllm.platforms.rocm import on_gfx950
48
-
61
value_layout = StridedLayout
50
- if on_gfx950():
62
+ if should_use_cdna4_mx_scale_swizzle():
63
try:
64
# triton < 3.6
65
from triton_kernels.tensor_details.layout import GFX950MXScaleLayout