diff --git a/tests/kernels/quantization/test_sm70_dflash2_fp8_m8.py b/tests/kernels/quantization/test_sm70_dflash2_fp8_m8.py index 6afe9618b..22a329bd0 100644 --- a/tests/kernels/quantization/test_sm70_dflash2_fp8_m8.py +++ b/tests/kernels/quantization/test_sm70_dflash2_fp8_m8.py @@ -42,19 +42,25 @@ def test_dynamic_prefill_and_batch_keep_fp16_graph_route(): fullgraph=True, dynamic=True, ) - for rows in [12, 8, 7, 32, 8, 128]: + for rows in [12, 8, 7, 32, 16, 8, 24, 40, 128]: x = torch.randn(rows, 5120, device="cuda", dtype=torch.float16) * 0.1 actual = compiled(x) assert torch.isfinite(actual).all() - if rows != 8: + if rows % 8 or rows > 32: # The compiled dynamic graph must preserve the original matrix, - # rather than silently applying FP8 to prefill or concurrency. + # rather than silently applying FP8 to prefill or partial blocks. expected = F.linear(x, original) assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) else: expected = apply_dflash2_fp8_m8(layer, x, None) assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) assert not torch.equal(actual, F.linear(x, original)) + # Each request's eight rows match what it gets drafting alone. + for row in range(0, rows, 8): + alone = apply_dflash2_fp8_m8(layer, x[row : row + 8], None) + assert torch.equal( + actual[row : row + 8].view(torch.int16), alone.view(torch.int16) + ) assert len(compilations) == 1 x = torch.randn(8, 5120, device="cuda", dtype=torch.float16) * 0.1 graph = torch.cuda.CUDAGraph() @@ -67,7 +73,8 @@ def test_dynamic_prefill_and_batch_keep_fp16_graph_route(): output.view(torch.int16), apply_dflash2_fp8_m8(layer, x, None).view(torch.int16), ) - assert apply_dflash2_fp8_m8(layer, x.repeat(4, 1), None) is None + assert apply_dflash2_fp8_m8(layer, x.repeat(5, 1), None) is None + assert apply_dflash2_fp8_m8(layer, x[:7], None) is None assert apply_dflash2_fp8_m8(layer, x.T.contiguous().T, None) is None assert apply_dflash2_fp8_m8(layer, x.unsqueeze(0), None) is None diff --git a/vllm/model_executor/layers/quantization/sm70_dflash2_fp8.py b/vllm/model_executor/layers/quantization/sm70_dflash2_fp8.py index eb2fded7d..7d11e6c2b 100644 --- a/vllm/model_executor/layers/quantization/sm70_dflash2_fp8.py +++ b/vllm/model_executor/layers/quantization/sm70_dflash2_fp8.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Channel-FP8 M8 projection route for the SM70 DFlash2 decoder. -Only the four large decoder projections are quantized. The existing FP16 -matrices remain available for context preparation and other batch sizes. The -FP8 layout replaces the FP16 M8 packing. Norms and the LM head are outside this -route. +Only the four large decoder projections are quantized. The decoder sees eight +query rows per drafting request, so batches of up to four requests (M=8..32 in +steps of eight) use the same FP8 weights a single request does. The existing +FP16 matrices remain available for context preparation and other row counts. +The FP8 layout replaces the FP16 M8 packing. Norms and the LM head are outside +this route. """ import torch @@ -26,6 +28,10 @@ _SHAPES = { } +def _fp8_rows(rows: int) -> bool: + return 0 < rows <= 32 and rows % 8 == 0 + + def _dispatch( x: torch.Tensor, weight: torch.Tensor, @@ -34,7 +40,10 @@ def _dispatch( split: int, ) -> torch.Tensor: output = x.new_empty((x.shape[0], weight.shape[0])) - if x.shape[0] == 8: + if _fp8_rows(x.shape[0]): + # The M=17..32 kernels take split 12 or 16. + if x.shape[0] > 16 and split < 12: + split = 16 ops.fp8_qpn8_gemm_sm70_out(output, x, codes, scales, split, 2, True, False) else: # This is the same mm_out used by the old FP16 M8 dispatch fallback. @@ -113,7 +122,7 @@ def apply_dflash2_fp8_m8( or not x.is_contiguous() ): return None - if not torch.compiler.is_compiling() and x.shape[0] != 8: + if not torch.compiler.is_compiling() and not _fp8_rows(x.shape[0]): return None return torch.ops.vllm.sm70_dflash2_fp8_m8_dispatch( x,