[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 <noreply@anthropic.com>

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
 

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,