From 9a33cee2ac7145f97ece37c2f77345a04d7467bf Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:58:11 +0800 Subject: [PATCH] [Perf][SM70][Spec Decode] Use the DFlash2 draft FP8 projections for 2-4 concurrent requests (#1289) The SM70 DFlash2 draft quantizes its four large decoder projections to channel FP8 for M=8, one drafting request. With 2-4 requests the decoder sees M=16..32 query rows and fell back to an FP16 torch.mm over the original matrices, which reads twice the bytes. fp8_qpn8_gemm_sm70_out already covers M<=32, so route M=8..32 in steps of eight there, using split 16 where M>16 needs split 12 or 16 (gate_up). Each request's rows now see the same FP8 weights it uses when drafting alone; with an unchanged split they are bit-identical to its M=8 result. Prefill, context KV, partial blocks and M>32 keep the FP16 matrices. Draft projection, 4x V100 TP4 shard, M32: gate_up 156 -> 77 us, down 77 -> 40 us, qkv 41 -> 41 us, o 20 -> 23 us. Qwen3.8-27B NVFP4 + DFlash2, 4x V100-SXM2-16GB, packaged profile, official-style tasks (GSM8K / MT-Bench / HumanEval, T=1.0, thinking, 16 prompts each) with 4 concurrent clients, aggregate tok/s: 282 -> 303, 337 -> 414, 501 -> 587. Acceptance length at c4 4.78/3.14/4.73 -> 4.95/3.15/4.41, against 4.90/3.23/4.62 for the same tasks at c1 (single-request path, unchanged and bit-identical). c1 is unchanged. Co-authored-by: Claude Opus 5.5 --- .../quantization/test_sm70_dflash2_fp8_m8.py | 15 +++++++++---- .../layers/quantization/sm70_dflash2_fp8.py | 21 +++++++++++++------ 2 files changed, 26 insertions(+), 10 deletions(-) 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,