[ROCm][MoE] Fix gfx950 block scale swizzle for AITER Triton MXFP4 W4A16 (#54773)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Sep 1, 2026 at 21:48 UTC
73723b707fe4b5ba99271ef3cbd3d3ad20d36000
1 file changed
+10
-2
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
+10
-2
@@ -152,6 +152,9 @@ def aiter_triton_kernel_w4a16_moe_forward(
152
score_mode: str | None = None,
153
):
154
assert quant_config is not None and rocm_aiter_ops.is_enabled()
155
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
156
+ should_use_cdna4_mx_scale_swizzle,
157
+ )
158
from vllm.platforms.rocm import on_gfx1250
159
160
try:
@@ -245,6 +248,11 @@ def aiter_triton_kernel_w4a16_moe_forward(
248
# SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu.
249
swiglu_add_residual = activation != MoEActivation.SILU
250
251
+ # Weight loading already swizzled the block scales on gfx950, so the kernel
252
+ # has to be told; the load-time swizzle and this gate must agree or the
253
+ # kernel indexes a swizzled buffer as if it were linear.
254
+ swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
255
+
256
intermediate = moe_gemm_a16w4(
257
hidden_states,
258
w1_data,
@@ -256,7 +264,7 @@ def aiter_triton_kernel_w4a16_moe_forward(
264
routing_data,
265
gather_indx=gather_idx,
266
gammas=gammas if apply_router_weight_on_input else None,
259
- swizzle_mx_scale=None,
267
+ swizzle_mx_scale=swz,
268
apply_swiglu=True,
269
alpha=swiglu_alpha,
270
limit=swiglu_limit,
@@ -276,7 +284,7 @@ def aiter_triton_kernel_w4a16_moe_forward(
284
routing_data,
285
scatter_indx=scatter_idx,
286
gammas=None if apply_router_weight_on_input else gammas,
279
- swizzle_mx_scale=None,
287
+ swizzle_mx_scale=swz,
288
unpadded_N=unpadded_N_w2,
289
unpadded_K=unpadded_K_w2,
290
)