[ROCm][gpt-oss] Hybrid CDNA4 swizzle gate for A8W4 MoE (#44804)

Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>

vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py

@@ -113,6 +113,12 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
     from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
     from aiter.ops.triton.quant_moe import downcast_to_static_fp8
 
+    from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+        should_use_cdna4_mx_scale_swizzle,
+    )
+
+    _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
+
     assert quant_config.w1_precision is not None, (
         "w1_precision in quant config can't be None"
     )
@@ -135,7 +141,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
         routing_data,
         gather_indx=gather_indx,
         gammas=gammas if apply_router_weight_on_input else None,
-        swizzle_mx_scale="CDNA4_SCALE",
+        swizzle_mx_scale=_swizzle_mx_scale,
         out_dtype=torch.float8_e4m3fn,
         apply_swiglu=True,
         alpha=swiglu_alpha,
@@ -155,7 +161,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
         routing_data,
         scatter_indx=scatter_indx,
         gammas=None if apply_router_weight_on_input else gammas,
-        swizzle_mx_scale="CDNA4_SCALE",
+        swizzle_mx_scale=_swizzle_mx_scale,
         unpadded_N=unpadded_N_w2,
         unpadded_K=unpadded_K_w2,
     )

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

@@ -19,6 +19,20 @@ logger = init_logger(__name__)
 CK_MXFP4_MOE_DIM_ALIGNMENT = 256
 
 
+def should_use_cdna4_mx_scale_swizzle() -> bool:
+    """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950.
+
+    CDNA4 swizzle requires BLOCK_K%256==0; at TP>=4 the A8W4 dispatch
+    picks BK<256 tiles for the smaller per-rank shapes, so swizzle must
+    be off. Used by both the weight-load swizzle in `_swizzle_mxfp4` and
+    the kernel-argument gate in `aiter_mxfp4_w4a8_moe`; they must agree.
+    """
+    from vllm.distributed import get_tensor_model_parallel_world_size
+    from vllm.platforms.rocm import on_gfx950
+
+    return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
+
+
 def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
     """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
     assert has_triton_kernels()
@@ -44,10 +58,8 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
         value_layout = StridedLayout
         scale_layout = StridedLayout
     elif current_platform.is_rocm():
-        from vllm.platforms.rocm import on_gfx950
-
         value_layout = StridedLayout
-        if on_gfx950():
+        if should_use_cdna4_mx_scale_swizzle():
             try:
                 # triton < 3.6
                 from triton_kernels.tensor_details.layout import GFX950MXScaleLayout