diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index cc2adc31fcd..7c3fe5831f3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/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, ) diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 51b7b29551d..db88ba273cd 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/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