[ROCm] [BugFix] Fix AITER MXFP4 ASM-GEMM crash on unfused shared experts (#55213)
Signed-off-by: Colin Zeng <Colin.Zeng@amd.com>
Colin Z committed
Sep 9, 2026 at 22:23 UTC
588a813a60448f6f71dfa0e0ecd52cc0bacc7329
3 files changed
+114
-22
tests/kernels/quantization/test_rocm_mxfp4.py
+61
@@ -637,3 +637,64 @@ def test_aiter_fp4_gemm_skinny_shapes(M, N, K):
637
pass_rate=pass_rate,
638
max_violation_factor=GEMM_MAX_VIOLATION_FACTOR,
639
)
640
+
641
+
642
+@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
643
+@pytest.mark.parametrize(
644
+ ("shape", "supported"),
645
+ [
646
+ ((2560, 80), True), # shared_expert gate/up_proj (in=2560 -> sn=80)
647
+ ((1280, 80), True), # gate/up_proj under TP=2 (out sharded)
648
+ ((2560, 8), True), # minimal aligned columns
649
+ ((32, 8), True), # minimal aligned rows and columns
650
+ ((2560, 20), False), # shared_expert down_proj (in=640 -> sn=20)
651
+ ((2560, 10), False), # down_proj under TP=2 (in 640->320 -> sn=10)
652
+ ((30, 8), False), # rows not a multiple of 32
653
+ ],
654
+)
655
+def test_asm_fp4_scale_swizzle_supported_shape_rules(shape, supported):
656
+ """The ASM swizzle needs rows % 32 == 0 and columns % 8 == 0."""
657
+ from vllm.model_executor.kernels.linear.mxfp4.aiter import (
658
+ _asm_fp4_scale_swizzle_supported,
659
+ )
660
+
661
+ weight_scale = torch.empty(shape, dtype=torch.uint8, device="cpu")
662
+ assert _asm_fp4_scale_swizzle_supported(weight_scale) is supported
663
+
664
+
665
+@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
666
+def test_asm_fp4_scale_swizzle_rejects_non_2d():
667
+ from vllm.model_executor.kernels.linear.mxfp4.aiter import (
668
+ _asm_fp4_scale_swizzle_supported,
669
+ )
670
+
671
+ assert not _asm_fp4_scale_swizzle_supported(
672
+ torch.empty(2560, dtype=torch.uint8, device="cpu")
673
+ )
674
+ assert not _asm_fp4_scale_swizzle_supported(
675
+ torch.empty(80, 2, 16, dtype=torch.uint8, device="cpu")
676
+ )
677
+
678
+
679
+@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
680
+def test_aiter_mxfp4_process_weights_falls_back_to_triton_for_misaligned_scale():
681
+ """A misaligned weight_scale must disable ASM and take the Triton path."""
682
+ from torch.nn.parameter import Parameter
683
+
684
+ from vllm.model_executor.kernels.linear.mxfp4.aiter import AiterMxfp4LinearKernel
685
+
686
+ kernel = object.__new__(AiterMxfp4LinearKernel)
687
+ kernel.use_asm_gemm = True
688
+ kernel.out_dtype = torch.bfloat16
689
+
690
+ layer = torch.nn.Module()
691
+ layer.weight_scale = Parameter(
692
+ torch.zeros(2560, 10, dtype=torch.uint8, device="cpu"), requires_grad=False
693
+ )
694
+
695
+ kernel.process_weights_after_loading(layer)
696
+
697
+ assert kernel.use_asm_gemm is False
698
+ # Triton path stores the transposed, contiguous scale.
699
+ assert tuple(layer.weight_scale.shape) == (10, 2560)
700
+ assert layer.weight_scale.is_contiguous()
vllm/model_executor/kernels/linear/mxfp4/aiter.py
+29
@@ -16,6 +16,22 @@ from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
16
17
logger = init_logger(__name__)
18
19
+
20
+_ASM_FP4_SCALE_ROW_MULTIPLE = 32
21
+_ASM_FP4_SCALE_COL_MULTIPLE = 8
22
+
23
+
24
+def _asm_fp4_scale_swizzle_supported(weight_scale: torch.Tensor) -> bool:
25
+ # The ASM swizzle reshapes weight_scale into a fixed tile layout, requiring
26
+ # scale rows divisible by 32 and scale columns divisible by 8
27
+ if weight_scale.ndim != 2:
28
+ return False
29
+ sm, sn = weight_scale.shape
30
+ return (
31
+ sm % _ASM_FP4_SCALE_ROW_MULTIPLE == 0 and sn % _ASM_FP4_SCALE_COL_MULTIPLE == 0
32
+ )
33
+
34
+
35
# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP
36
# which can force the engine core to spawn instead of fork.
37
# is_aiter_found_and_supported() checks platform + arch + library availability via
@@ -169,6 +185,19 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel):
185
return True, None
186
187
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
188
+ if self.use_asm_gemm and not _asm_fp4_scale_swizzle_supported(
189
+ layer.weight_scale.data
190
+ ):
191
+ logger.warning_once(
192
+ "AITER ASM FP4 GEMM requires weight_scale dims divisible by "
193
+ "(%d, %d), but this layer has weight_scale shape %s. Falling "
194
+ "back to the AITER Triton FP4 GEMM for this layer.",
195
+ _ASM_FP4_SCALE_ROW_MULTIPLE,
196
+ _ASM_FP4_SCALE_COL_MULTIPLE,
197
+ tuple(layer.weight_scale.data.shape),
198
+ )
199
+ self.use_asm_gemm = False
200
+
201
if self.use_asm_gemm:
202
from aiter.ops.shuffle import shuffle_weight
203
vllm/model_executor/models/qwen3_next.py
+24
-22
@@ -141,17 +141,30 @@ class Qwen3NextSparseMoeBlock(nn.Module):
141
142
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
143
144
- shared_expert_group_size = get_quark_ocp_mx_group_size(
145
- quant_config,
146
- f"{prefix}.shared_expert.down_proj",
147
- )
148
- self.replicate_shared_expert = _should_replicate_misaligned_shared_expert(
149
- config.shared_expert_intermediate_size,
150
- self.tp_size,
151
- shared_expert_group_size,
152
- parallel_config.enable_expert_parallel,
153
- self.is_sequence_parallel,
154
- )
144
+ # Resolve shared-expert fusion first (when enabled, TP alignment constraint no
145
+ # longer applies)
146
+ self.is_fused_shared_expert_enabled = False
147
+ if config.shared_expert_intermediate_size > 0:
148
+ self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
149
+ quant_config,
150
+ prefix,
151
+ shared_expert_name="shared_expert",
152
+ )
153
+
154
+ if self.is_fused_shared_expert_enabled:
155
+ self.replicate_shared_expert = False
156
+ else:
157
+ shared_expert_group_size = get_quark_ocp_mx_group_size(
158
+ quant_config,
159
+ f"{prefix}.shared_expert.down_proj",
160
+ )
161
+ self.replicate_shared_expert = _should_replicate_misaligned_shared_expert(
162
+ config.shared_expert_intermediate_size,
163
+ self.tp_size,
164
+ shared_expert_group_size,
165
+ parallel_config.enable_expert_parallel,
166
+ self.is_sequence_parallel,
167
+ )
168
if self.tp_size > config.num_experts:
169
raise ValueError(
170
f"Tensor parallel size {self.tp_size} is greater than "
@@ -183,17 +196,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):
196
prefix=f"{prefix}.shared_expert_gate",
197
)
198
186
- self.is_fused_shared_expert_enabled = False
187
- if (
188
- config.shared_expert_intermediate_size > 0
189
- and not self.replicate_shared_expert
190
- ):
191
- self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
192
- quant_config,
193
- prefix,
194
- shared_expert_name="shared_expert",
195
- )
196
-
199
if (
200
self.is_fused_shared_expert_enabled
201
or config.shared_expert_intermediate_size <= 0