[Leonccaa]

Change

On SM70 the paged QSA scorer loads every 64-key tile of the compressed key cache once per query row. In a multi-request decode batch the rows of one request share a page table, so this PR adds a request-segmented scorer that loads each tile once per segment of up to four rows of one request and multiplies it by those rows together: four rows × four heads form one 128 × 16 query block, the paged kernel's own 64 × 128 by 128 × 16 dot, with the same 2 warps, the same pipelined tile loop and the same per-row 16-lane reduction (the other rows' lanes are zeroed). Scores, visible lengths and the selected top-512 blocks are bitwise equal to the paged scorer.

qsa_select_paged_tokens uses it for batches with two or more requests, at most 256 rows per logits chunk, at least two rows per request on average, FP16 H4/D128 and int32 metadata. Single-request batches keep the shared-key and cuBLAS paths. One-row-per-request batches (MTP drafting steps, also when padded for CUDA graphs) keep the paged scorer, where one-row segments measured up to 15 % slower. Segment descriptors are built on the device, so the call is CUDA-graph safe.

VLLM_SM70_QSA_REQUEST_SEGMENTS switches it. The ReVolta SM70 profile turns it on; an explicit value wins. The variable is not a torch.compile factor (the route is chosen inside the QSA custom op), so registering or switching it does not recompile a deployment. The acceleration report gains a qsa_request_segments entry.

Validation

CPU: Triton-interpreter checks of the segment builder and scorer against the paged kernel and a NumPy oracle (11), dispatch tests (18), the compile-factor test and the profile tests pass; pre-commit and mypy 3.12 pass on the changed files.

GPU, one V100 PCIe at 250 W:

• 36 kernel cases (C1/M8, C4/M16, C4/M20, C8/M40 at 32K / 128K / 256K tokens, segment limits 1 / 2 / 4), eagerly and through CUDA graphs replayed after queries, keys, page tables, positions, lengths and request runs change: FP32 score words, visible lengths and top-512 IDs identical to the paged scorer; guard buffers intact.
• qsa_select_paged_tokens with the variable off and on over 30 shapes (decode batches, one row per request, ragged rows, a prefill chunk next to decode rows with and without the host request split, two logits chunks, CUDA-graph padding rows): selected token indices identical, and the shapes that must keep the paged scorer do.
• Whole selection per layer (scoring, top-512, expansion), new / paged: 4 requests × 4 rows 0.39–0.40, 4 × 5 rows 0.47–0.50, 8 × 4 rows 0.32–0.36, 8 × 5 rows 0.49–0.55, 32K–256K tokens.

Engine, our 4 × V100 PCIe-switch server at 250 W, the 0.1.5 release image with this PR's three runtime files, variable off and on in one window with the same compiled artifacts (no recompilation, same KV pool):

• C1 greedy and sampled output bitwise identical (single requests do not use the new scorer).
• At 4 and 8 concurrent requests the engine is not run-to-run reproducible, and two flag-off processes differ from each other as much as off and on do. DFlash2 short prompts: 118 near-tie flips off vs on, 118 off vs off. Long contexts: 63 of 84 requests diverge in both pairs, median first divergence at token 153 and 154.
• Decode throughput while all requests decode, on / off, greedy / sampled. DFlash2 (against the mean of two off runs): 4 × 32K 1.03 / 1.07, 4 × 128K 1.15 / 1.14, 4 × 256K 1.21 / 1.21, 8 × 32K 1.01 / 1.07, 8 × 64K 1.08 / 1.10. MTP4 (one off run): 4 × 32K 1.03 / 1.15, 4 × 128K 1.14 / 1.13, 4 × 256K 1.24 / 1.21, 8 × 64K 1.08 / 1.08. Two off runs agree within 1 % at 4 × 128K / 256K and within 2–8 % at 8 requests.
• Short-prompt speed matrix (512-token answers): C1 unchanged; DFlash2 C8 sampled 1.036.

Variants tried and not taken: 8-row segments (two query blocks) spill at 2 warps; 4 or 8 warps fit but change the reduction order (about 3.7 % of score words differ) and were slower than this scorer in every measured shape; a two-pass 8-row variant at 2 warps is bitwise but not faster than four-row segments.

Attribution

The request-segment idea, the first prototype and its CPU interpreter checks come from Codex in the research PRs #56 and #57. Claude Code reworked the kernel until it was bitwise on SM70 (loop-invariant grouped queries inside the paged kernel's pipelined tile loop), integrated and tested it, and ran the GPU and engine validation. Commits are signed off under the DCO.

🤖 Generated with Claude Code