[Quant] Keep the flashinfer_deepgemm FP8 GEMM to 1 <= M < 32 (#32843)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
729081e14f
commit
00a219f6c9
@@ -751,7 +751,27 @@ def flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback(
|
|||||||
# fp8_blockscale_gemm_sm90 requires: N % 64 == 0, K % 128 == 0
|
# fp8_blockscale_gemm_sm90 requires: N % 64 == 0, K % 128 == 0
|
||||||
shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 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:
|
if weight_scale.dtype == torch.int32:
|
||||||
weight_scale = _unpack_ue8m0_scale_for_triton(
|
weight_scale = _unpack_ue8m0_scale_for_triton(
|
||||||
weight_scale, weight.shape, block_size
|
weight_scale, weight.shape, block_size
|
||||||
|
|||||||
Reference in New Issue
Block a user