diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index eaa25d6f3..776229b8c 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ba474ac5f..d4ae58608 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py index fc16e4466..53a115966 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/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 )