[AMD][Fix] Fix aiter bpreshuffle GEMM for output sizes it cannot dispatch for qwen3.5 mxfp-attn-fp8-v2 TP4 (#37564)

Co-authored-by: HAI <hixiao@gmail.com>

python/sglang/srt/layers/quantization/fp8.py

@@ -67,6 +67,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     requant_block_scale_ue8m0_for_deepgemm,
     resolve_mxfp8_dense_gemm_backend,
     unshuffle_aiter_fp8_weight,
+    use_aiter_bpreshuffle_gemm,
 )
 from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
 from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin
@@ -934,7 +935,8 @@ class Fp8LinearMethod(LinearMethodBase):
                     weight_scale = weight_scale.t().contiguous()
                     if _use_aiter and self.use_aiter_fp8_per_token:
                         self.use_per_token_if_dynamic = True
-                        qweight = shuffle_weight(qweight.contiguous(), (16, 16))
+                        if use_aiter_bpreshuffle_gemm(qweight.shape[0]):
+                            qweight = shuffle_weight(qweight.contiguous(), (16, 16))
                 else:
                     # per-tensor quantization
                     qweight, weight_scale = input_to_float8(layer.weight)
@@ -990,7 +992,8 @@ class Fp8LinearMethod(LinearMethodBase):
                                 weight=weight,
                                 weight_scale=weight_scale,
                             )
-                        weight = shuffle_weight(weight.contiguous(), (16, 16))
+                        if use_aiter_bpreshuffle_gemm(weight.shape[0]):
+                            weight = shuffle_weight(weight.contiguous(), (16, 16))
                 else:
                     # Dequant -> Quant with max scale so we can run per tensor.
                     weight = layer.weight

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -1859,6 +1859,15 @@ def _apply_fallback_scaled_mm(
     return output.to(dtype=input_dtype)
 
 
+def use_aiter_bpreshuffle_gemm(output_size: int) -> bool:
+    # aiter's CK gemm_a8w8_bpreshuffle instances are GemmSpecialization::Default
+    # (pre-shuffled weights are never N-padded) with NPerBlock=64, so any N that
+    # is not a multiple of 64 raises "This GEMM is not supported!". Measured on
+    # gfx950 for M=16384/N=32/K=4096, torch._scaled_mm rowwise runs that shape in
+    # 14us against 90us for the cktile instance that does accept it.
+    return _use_aiter and output_size % 64 == 0
+
+
 def apply_fp8_linear_bmm_flashinfer(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -2069,7 +2078,10 @@ def apply_fp8_linear(
         # into this sector means use dynamic per-token-per-channel quant
         # per-token scale quant for input matrix, every row(one token) have one scale factor
         # per-channel scale quant for weight matrix, every col(one channel) have one scale factor
-        if _use_aiter:
+        # Must agree with the load-time predicate that decides whether the
+        # weight was pre-shuffled; an unshuffled weight through the aiter path
+        # (or a shuffled one through torch._scaled_mm) silently returns garbage.
+        if use_aiter_bpreshuffle_gemm(weight.shape[1]):
             # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype)
             # XQ -> input tensor, shape = (m, k)
             # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf

python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py

@@ -15,6 +15,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     apply_fp8_linear,
     cutlass_fp8_supported,
     normalize_e4m3fn_to_e4m3fnuz,
+    use_aiter_bpreshuffle_gemm,
 )
 from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme
 from sglang.srt.layers.quantization.utils import requantize_with_max_scale
@@ -95,7 +96,7 @@ class QuarkW8A8Fp8(QuarkLinearScheme):
                 weight_scale = layer.weight_scale.data
             if self.per_token:
                 weight_scale = weight_scale.view(-1, 1)
-            if _use_aiter:
+            if use_aiter_bpreshuffle_gemm(weight.shape[0]):
                 layer.weight = Parameter(
                     shuffle_weight(weight, (16, 16)).t(), requires_grad=False
                 )