diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 97762fd98..037ae9412 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -751,7 +751,27 @@ def flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback( # fp8_blockscale_gemm_sm90 requires: N % 64 == 0, K % 128 == 0 shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 - if not (shape_supported and dtype_supported): + # Keep this backend to 1 <= M < 32, mirroring vLLM's + # FlashInferFp8DeepGEMMDynamicBlockScaledKernel. fp8_blockscale_gemm_sm90 is + # one entry point over two kernels and only the M < 32 swapAB half is worth + # taking: + # M >= 32 picks the non-swapAB kernel, which is slower than DeepGEMM (worst + # just above the threshold) and, on some checkpoints, less accurate. + # M == 0 hard-fails inside the kernel ("Check failed: (input_ptr != + # nullptr)"). Empty batches are a normal steady-state input, not an edge + # case: DP attention hands an idle rank a zero-token forward so the + # collectives stay in sync (ScheduleBatch.prepare_for_idle). + # Same shape of guard as the gfx95 CK M bound below. + m_supported = 1 <= input.view(-1, input.shape[-1]).shape[0] < 32 + + if not m_supported and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + # DeepGEMM covers both ends and falls back to triton on its own for + # shapes it cannot serve. + return deepgemm_w8a8_block_fp8_linear_with_fallback( + input, weight, block_size, weight_scale, input_scale, bias + ) + + if not (shape_supported and dtype_supported and m_supported): if weight_scale.dtype == torch.int32: weight_scale = _unpack_ue8m0_scale_for_triton( weight_scale, weight.shape, block_size