[Peuqui]

Update 2026-10-28: rebased onto main c59b3c4, which contains #1222 and #1227; their commits are no longer part of this branch, which now holds 20 commits (the 14 of #1279 and 6 new). The six are unchanged apart from the regenerated reference (git range-diff), with one addition: the fused-softmax commit now also updates the two DeepSeek-V4.1 callers of attend_gathered_keys in vllm/models/deepseek_v41/attention_ops.py (from #1277), which still passed the removed logits buffer and relied on the function to zero unused keys; the CI mypy run on the earlier push caught it. Against main's function on random data with NaN in the unused slots, 99.69-100 % of their outputs are bit-identical and the rest differ by one fp16 step, on a Tesla V100 and a Quadro RTX 8000, and the DeepSeek-V4.1 test files give the same results as on main. Test Plan files on the head before that addition (6dc35fd), Quadro RTX 8000 / Tesla V100, all passing: test_deepseek_v4_sm70_indexer.py 31 / 30 (1 skipped), test_deepseek_v4_sm70_routes.py 30 / 30, test_cudagraph_dispatch.py 51 / 51, test_dsv4_sm70_sparse_bmm.py 12 skipped / 12 passed, test_dsv4_sparse_prefill_bmm.py 17 / 17, test_deepseek_v4_sm70_qnorm_rope_kv_insert.py 4 / 4, test_compressor_kv_cache.py 32 / 32 (6 skipped each), test_sparse_indexer_decode_seq_lens.py 11 / 11. pre-commit run --from-ref origin/main --to-ref HEAD on this head: every hook passes except check-layering, which now reports vllm/config/speculative.py (model coupling 142 -> 144) next to vllm/model_executor/layers/fused_moe/runner/moe_runner.py (platform coupling 37 -> 39); main c59b3c4 reports the same two counts, so this PR adds none. pre-commit run mypy-3.10 --hook-stage manual: passed.

Purpose

Long DeepSeek-V4-Flash prefills on SM70/SM75 pipeline stages spend much of their non-MoE time in memory-bound glue around the lightning indexer and the sparse MLA attention. In a per-stage nsys profile of a 150k-token prefill (PP=5 over three Tesla V100 and two Quadro RTX 8000), the indexer's relu/head-sum epilogue took 14.0 s and 17.6 s of the two RTX stages, the pipeline's slowest, and six elementwise torch ops took 820 of 1434 us of each sparse attention pass on the RTX 8000.

Stacked PR. The branch sits on #1279 (TP prefill, itself on #1278 and #1216). #1222 (cuBLAS decode indexer above 131,072 tokens, same indexer file), on which it was built, is merged. The copy of the #1279 stack is identical to the open PR. Only the last six commits are new here, so please review those (git log 3f48315..32e4372, where 3f48315 is the head of #1279). The Turing indexer kernel uses the m16n8k8 fragment helpers of #1278 and shares the sm_75 source list in CMakeLists.txt and neighbouring lines of csrc/ops.h and csrc/torch_bindings.cpp with #1279.

The six commits make five changes (commit 2 has its test update in a commit of its own):

1. Vectorized indexer prefill epilogue (77fd5d1). BLOCK_K 512 over 2 warps reads eight fp16 scores (16 bytes) per thread instead of two. On the [508 tokens x 64 heads, 1024 keys] score tile it reads at the cards' plain read rate of that tile. The head sum stays fp32; its reduction order changes in the last bits (up to 1.5e-5 on unit-scale scores), so the top-k can differ only at near ties. Prompts up to index_topk * 4 tokens select every key and are unaffected.
2. Fused sparse MLA softmax (459edf6; test update 1ba5d92). The gather writes zeros for unused slots and never reads their indices. One Triton kernel reads the fp32 scores once and writes fp16 probabilities: scale, mask, sink column and the fp32 softmax with libdevice exp, the same expf torch's softmax uses. The fp32 logits and probability buffers go away (140 -> 110 MiB per 128-token pass at 640 keys). 99.2-99.8 % of the outputs are bit-identical, the rest differ by at most 3e-5 (reduction order); the error against float64 is unchanged.
3. Fused lightning-indexer prefill scores on Turing tensor cores (37febaf). indexer_prefill_scores_sm75 keeps one token's queries for 32 heads as m16n8k8 A fragments in registers, stages e4m3 key tiles widened to fp16 in shared memory, rounds the scores to fp16 as the cuBLAS route stores them, and keeps relu, head weights and the head sum in registers; only the [tokens, keys] result reaches memory.
   • Taken when the indexer runs its relu route on an SM75 device (the capability of the query's own device) with 32 or 64 heads. Volta and other head counts keep the cuBLAS route plus the epilogue of commit 1, or the per-head Triton kernel.
   • kernel_config.sm70_sparse.indexer_prefill_cublas keeps its meaning (cuBLAS or the per-head Triton kernel) for the chunks that do not take the fused kernel; its docstring now says so.
   • The source builds for sm_75 only. Each prefill scoring route logs its choice once per device ("SM70 indexer prefill: ... on cuda:N").
4. Indexer logits freed per prefill sub-chunk (8477763). The loop kept the previous sub-chunk's scores alive until the next scorer had allocated its own, so the peak under VLLM_SPARSE_INDEXER_MAX_LOGITS_MB was two sub-chunks while the profile run reserves one. SparseAttnIndexerKpool has the same loop and gets the same release; it was not measured.
5. Compressed KV gather with 2048 two-warp workers per request instead of 128 (437ad09). Same kernel and arithmetic, bit-identical output. This Triton gather also serves the ROCm backend and NVIDIA GPUs without CuteDSL (dequantize_and_gather_k_cache); only the V100 and the RTX 8000 were measured.

Not a duplicate. I compared the changed files with those of all 82 open PRs (complete file lists through the API), searched the added and removed lines of every open PR for this PR's symbols (indexer_prefill_scores, _relu_weight_headsum, _EPILOGUE_BLOCK_K, sparse_prefill_bmm, dequantize_and_gather_k_cache_triton, NUM_WORKERS, VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, del logits, indexer_prefill_cublas), ran a three-way merge with every PR that shares a file, and searched the open PRs for indexer, "sparse MLA", softmax, epilogue, DeepSeek-V4, DSv4, "lightning indexer" and gather. Findings:

• Serve GLM 5.3 Flash and Qwen3.8 Flash Next GGUF across seven heterogeneous GPUs #1260 (Serve GLM 5.3 Flash and Qwen3.8 Flash Next GGUF across seven heterogeneous GPUs; conflicting with main) changes vllm/models/deepseek_v4/sm70/indexer.py and its test file: FP32 cuBLAS prefill scores, head masking in the relu/head-sum epilogue, a per-tile release of the scores and a rewritten comment at _EPILOGUE_BLOCK_K. That overlaps commits 1 and 3 in text and in meaning: commit 1's 16-byte loads read fp16 scores, and commit 3 rounds its scores to fp16 to match the cuBLAS route as it is on main. Whichever lands second has to adapt.
• [Perf][SM70][DeepSeek-V4] Keep the cuBLAS decode indexer above 131,072 tokens of context #1222 was stacked here unchanged and is merged now.
• [Perf][SM70] Use the grouped page4 route for QSA DCP prefill attention #828 shows the indexer and sparse attention files in its GitHub file list only because of its old base; its three commits change Qwen4Exp QSA files only.
• SM70: integrate retained MTP optimizations with upstream runtime owners #903 (vllm/config/sm70_sparse.py), [Model][SM70] DeepSeek-V4.1-Flash-1017 on 8x V100: EXL3 experts, graph attention, DSpark speculative decoding, 1M context (stacks on #1031) #1243 and [Model][SM70] DeepSeek-V4.1-Flash on V100: PP3xTP4 (12 GPUs, 256K) and PP2xTP4 (8 GPUs) #1031 (both list tests/kernels/attention/test_dsv4_sparse_prefill_bmm.py in a build smoke script) already conflict with main; a three-way merge with this branch gives the same conflicting files as with its base 64eb4db.
• Other open PRs touch CMakeLists.txt, csrc/ops.h, csrc/torch_bindings.cpp or the generated reference in other places.

AI assistance was used for this change. I reviewed every line and ran the tests and measurements below.

Test Plan

Native build of this branch in place (TORCH_CUDA_ARCH_LIST=7.0; indexer_prefill_scores_sm75.cu builds for sm_75 only, like the sm_75 sources of #1278 and #1279).

One pytest process per file on a Quadro RTX 8000 (SM75) and on a Tesla V100 (SM70):
python -m pytest tests/kernels/test_deepseek_v4_sm70_indexer.py
python -m pytest tests/models/test_deepseek_v4_sm70_routes.py
python -m pytest tests/kernels/attention/test_dsv4_sparse_prefill_bmm.py
python -m pytest tests/kernels/attention/test_dsv4_sm70_sparse_bmm.py
python -m pytest tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py
python -m pytest tests/kernels/test_compressor_kv_cache.py
python -m pytest tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
python -m pytest tests/v1/cudagraph/test_cudagraph_dispatch.py

pre-commit run --from-ref origin/main --to-ref HEAD; pre-commit run mypy-3.10 --hook-stage manual --from-ref origin/main --to-ref HEAD

Commit 3 in a model, A/B/A/B on an 8-layer DeepSeek-V4-Flash mini checkpoint (layers 0-7 of the checkpoint byte for byte, MTP off), PP=2 with a Tesla V100 as stage 0 and a Quadro RTX 8000 as stage 1, this branch in place with stock Triton, --max-num-batched-tokens 1024, --max-model-len 307200: one uncounted boot, then without and with the fused Turing scorer twice each. "Without" is this branch with the fused branch of sm70_indexer_prefill_logits removed, so the RTX stage takes the cuBLAS route plus the epilogue of commit 1. Per boot: prefill of 8k, 32k, 150k and 280k tokens, decode time per token, and the top-5 logprobs of 16 greedy tokens after a short and after a 4k-token prompt.

One full-model boot of this branch in place (stock Triton, no kernel_config Triton switch): DeepSeek-V4-Flash PP=5 as in the contract below, but without DSpark, since Model Runner V2 on main rejects speculative method 'dspark', and with CUDA graphs for one decode token. Per boot: prefill and decode client, eight quality questions (1500 tokens, read by hand), a long-context probe up to the entry's maximum length, and the route lines of the indexer prefill scorer in the log.

Path coverage

Path / provider	Format	Key conditions (shape, TP/EP, graph, capability)	Expected selection / fallback	Evidence (source / static / executed / measured)	Unverified scope
SM70-family indexer prefill epilogue	fp16 scores	cuBLAS prefill scoring (Volta; Turing chunks outside commit 3)	BLOCK_K 512, 2 warps	executed: test_deepseek_v4_sm70_indexer.py on RTX 8000 and V100; measured: CUDA graphs, model A/B	none
sparse MLA bmm attention, prefill and decode	FP8 KV (DeepSeek-V4)	SM70 family	gather with zeroed slots + fused Triton softmax	executed: test_dsv4_sparse_prefill_bmm.py, test_dsv4_sm70_sparse_bmm.py (decode reference Volta-only), test_deepseek_v4_sm70_routes.py; measured: CUDA graphs, model A/B	none
indexer prefill scores on SM75	e4m3 keys, 32 or 64 heads	relu route, SM75 (query's own device)	indexer_prefill_scores_sm75; otherwise cuBLAS + epilogue or the per-head Triton kernel	executed: test_deepseek_v4_sm70_indexer.py on RTX 8000 (fp64 reference with fp16 score rounding, 32/64 heads, partial tiles); measured: mini checkpoint A/B on the RTX stage (Test Result); executed: full-model boot of this branch, route lines on both RTX stages (Test Result)	SM8x+ (never selected there); head counts other than 32 and 64 (not selected)
indexer logits per sub-chunk	any	VLLM_SPARSE_INDEXER_MAX_LOGITS_MB below one chunk's scores	previous sub-chunk freed before the next is scored	executed: mini checkpoint runs with 1, 2 and 5 sub-chunks per step (SparseAttnIndexer)	SparseAttnIndexerKpool (same change, not run)
compressed KV gather	FP8 KV (DeepSeek-V4)	prefill; Triton gather (SM70 family, ROCm, NVIDIA without CuteDSL)	2048 workers, bit-identical	executed: test_deepseek_v4_sm70_qnorm_rope_kv_insert.py, test_compressor_kv_cache.py on RTX 8000 and V100; measured: CUDA graphs, every grid compared bit for bit with the 128-worker launch	ROCm and NVIDIA GPUs other than V100 and RTX 8000 (not run)

Acceleration and benchmark contract (required for performance changes)

• Default enabled or opt-in: default enabled on the matching capability (per device); no new switch. The fused Turing scorer cannot be switched off separately.
• Required CLI options and environment switches: none for the change. Measured entry: DeepSeek-V4-Flash (MXFP4 experts, block-FP8 linears), PP=5 over the five cards (CUDA_VISIBLE_DEVICES=0,1,4,3,2: RTX 8000 on the first and last stage), VLLM_PP_LAYER_PARTITION=10,8,8,8,9, DSpark with 5 speculative tokens, --dtype half, --kv-cache-dtype fp8, --max-num-seqs 1, --max-model-len 307200, VLLM_USE_V2_MODEL_RUNNER=1, NCCL_P2P_DISABLE=1. The two model A/B series below used --max-num-batched-tokens 512 and --num-gpu-blocks-override 3000; the full boot of this branch used --max-num-batched-tokens 1024 and --num-gpu-blocks-override 3500, without DSpark (Test Plan).
• KV cache dtype used for the benchmark: fp8.
• Wheel SHA or source commit:
  • Model A/B series: fork builds of github.com/Peuqui/1Cat-vLLM. Reference 4da0bcf (the fork's production at the time); changed 5669da1 (commit 1) and e1cb50f (commits 1-2). Both changed builds also carry 6c0ef8e and 0314585 (QPN8 decode split-K by occupancy, M <= 8). Net of the two, that rule only changes Turing grids whose split-16 launch spills fewer than half the SMs into a second wave (output widths of about 4,600 to 5,760); DeepSeek-V4-Flash's block-FP8 projections have widths 512, 1024, 4096, 8192 and 32768.
  • Full boot: this branch at 437ad09 in place. Mini checkpoint runs: an earlier, unpublished state of this branch that differs only in the scope of the route log lines and the wrapping of two strings.
  • Commits 3, 4 and 5 were not in either model A/B series.
• PYTHONPATH, source overlays, or external native libraries used:
  • Model A/B series: the changed arm ran its fork source tree on PYTHONPATH with its own compile cache; both arms used the stock Triton 3.6.0 of the environment.
  • The logits measurements of commit 4 ran on the mini checkpoint with a fork tree, the patched Triton build of a later PR of this series and its tensor-core switch.
  • Mini checkpoint A/B and full boot of this branch: PYTHONPATH on this tree only, stock Triton 3.6.0, no kernel_config Triton switch.
• User-entry route-hit, speed, and output-quality evidence: see Test Result.
• Existing kernel/backend registry and capability-rejection reasons: selection stays in the DeepSeek-V4 SM70 indexer and sparse attention routes, keyed on the capability of the layer's own device. indexer_prefill_scores_sm75 rejects other devices, dtypes, strides and misaligned inputs with an error.
• Per-engine configuration and deprecated-variable compatibility: no setting added or changed; the docstring of sm70_sparse.indexer_prefill_cublas now names the Turing chunks it does not decide.
• New/changed environment metadata and generated reference check: no environment metadata. docs/design/architecture/runtime_reference.md is regenerated (one Sm70SparseConfig anchor moved by the docstring line) and checked by the check-architecture-reference hook.
• Route snapshot command and intentional changed rows: none; the routes are chosen inside the indexer and sparse attention ops.
• SM70 registered variables / unregistered reads / model-parameter restrictions, before -> after: no change.

Test Result

Model A/B/A/B on DeepSeek-V4-Flash PP=5 (Wheel SHA above): order one uncounted cold boot of the changed build, then reference a, changed a, reference b, changed b. Prefill time in s, lower is better (reference a / reference b -> changed a / changed b):

• Epilogue alone (commit 1): 32.7k tokens 14.66 / 14.66 -> 14.16 / 14.15 (-3.4 %); 178k 112.79 / 112.76 -> 93.94 / 93.96 (-16.7 %); 287k 222.35 / 222.50 -> 173.07 / 173.14 (-22.2 %). Decode step unchanged (87-88 ms at 178k, 93-94 ms at 287k). Greedy output and the eight quality answers byte-identical in all five boots.
• Epilogue and fused softmax (commits 1-2): 32.7k 14.67 / 14.66 -> 12.94 / 12.94 (-11.8 %); 178k 112.71 / 112.70 -> 82.20 / 82.14 (-27.1 %); 287k 222.26 / 222.19 -> 149.73 / 149.72 (-32.6 %). Decode step 1.9-3.4 ms shorter (178k 86.9 / 86.8 -> 84.9 / 84.5 ms, 287k 93.6 / 93.8 -> 90.4 / 90.8 ms). Texts byte-identical within each arm and different between the arms (reduction order); all read by hand and equivalent.

Commit 3, mini checkpoint A/B/A/B (Test Plan). Prefill time in s, lower is better (without a / without b -> with a / with b):

• 8k tokens 1.116 / 1.127 -> 1.122 / 1.113 (-0.4 %); 32k 4.130 / 4.163 -> 4.152 / 4.246 (+1.3 %); 150k 22.906 / 22.914 -> 20.871 / 20.881 (-8.9 %); 280k 51.255 / 51.413 -> 43.537 / 43.516 (-15.2 %).
• Decode time per token 16.86 / 17.01 -> 16.98 / 17.05 ms (58.8-59.3 -> 58.7-58.9 tok/s); the indexer decode path is unchanged.
• Greedy tokens identical in all four counted boots for both prompts. The largest top-5 logprob gap against the first "without" boot is 0.062 for the "with" boots and 0.125 for the second "without" boot, so the fused scorer stays within the boot-to-boot spread of the cuBLAS route.
• Peak memory.used of the RTX stage 30,025 MiB with the fused scorer and 30,165 MiB without; the V100 stage 30,374 MiB in both.
• The mini checkpoint's text is not meaningful (8 of 43 layers); only equality and distances are evaluated.

Full boot of this branch (Test Plan), one boot, no A/B:

• Route lines: "SM70 indexer prefill: fused SM75 scorer" on cuda:0 and cuda:4 (the two RTX 8000 stages), "SM70 indexer prefill: cuBLAS route" on cuda:1 to cuda:3 (the V100 stages).
• Eight quality answers (1500-token limit), read by hand: all correct and all ended with stop (four-stroke engine, a Python function for the n-th prime, a speed calculation with the right result of 84 km/h, a question with a misspelled name declined instead of answered with invented facts, an ordering puzzle, a code bug, a summary, an English explanation).
• Long-context probe, 400 decoded tokens per prompt; time to first token in s (lower is better), decode ms per step and tok/s (higher is better; one token per step without DSpark): 32,668 tokens 10.00 s, 114.7 ms, 8.7 tok/s; 178,480 tokens 60.68 s, 197.6 ms, 5.1 tok/s; 287,039 tokens 107.66 s, 260.0 ms, 3.8 tok/s; 302,833 tokens 115.14 s, 268.9 ms, 3.7 tok/s. All four answers are coherent summaries of their prompt.
• These decode numbers are those of main plus this stack without speculative decoding; apart from commit 2 (model A/B above) this PR does not change the decode path.
• KV cache 317,261 tokens. Peak memory.used per card (GPU 0 RTX 8000, 1 V100, 2 RTX 8000, 3 V100, 4 V100), MiB: boot and client 39,745 / 32,202 / 35,575 / 32,202 / 32,202; quality and long-context probe 39,231 / 32,284 / 35,635 / 32,284 / 32,284.

Kernel microbenchmarks, from the commit messages (CUDA graphs, lower is better):

• Indexer epilogue, [508 tokens x 64 heads, 1024 keys] fp16 tile (66.6 MB): RTX 8000 632 -> 118.6 us (105 -> 562 GB/s), V100 362 -> 85.1 us (184 -> 782 GB/s).
• Sparse attention pass, us per call: prefill T128 H64 W640 RTX 8000 1438.3 -> 755.7, V100 1178.7 -> 632.0; W2176 RTX 8000 4666.8 -> 2367.7, V100 3760.3 -> 1983.5; decode T6 H64 (128 + 512 keys) V100 146.5 -> 96.5, SWA only (128 keys) 55.1 -> 36.6.
• Turing indexer scores, RTX 8000, 512 tokens x 64 heads, ms per call, cuBLAS + epilogue -> fused: 1024 keys 0.326 -> 0.151, 4096 keys 1.270 -> 0.586, 18750 keys 5.551 -> 2.305, 37500 keys 11.018 -> 4.518 (57-70 TFLOPS); max error against fp64 equal to the cuBLAS route in all shapes.
• Compressed KV gather, us per call (rows 2048 / 18750 / 37500): RTX 8000 47.2 -> 7.2 / 532.3 -> 107.3 / 1061.9 -> 215.6; V100 45.8 -> 7.7 / 616.4 -> 87.4 / 1229.3 -> 167.7 (about 280 and 345 GB/s at 18750 rows).
• Logits per sub-chunk (DeepSeek-V4 8-layer mini checkpoint, PP=2, 1024-token chunks, 280k-token prompt, V100 stage, torch allocator peak in MiB): budget 160 before the change 16759 (two sub-chunks alive at once, 160 + 113 MiB), after 16646; budget 512 (one sub-chunk of 273 MiB) 16759. Prefill time with budget 512 against 64, A/B/A/B, s: 8k 0.854 / 0.843 vs 0.845 / 0.837, 150k 15.593 / 15.593 vs 15.596 / 15.599, 280k 32.849 / 32.758 vs 32.888 / 32.871.

Unit tests on this branch at 437ad09, native build in place (_C carries 73 sm_70 and 4 sm_75 cubins), one process per file:

• Quadro RTX 8000: test_deepseek_v4_sm70_indexer.py 31 passed; test_deepseek_v4_sm70_routes.py 30 passed; test_dsv4_sparse_prefill_bmm.py 17 passed; test_dsv4_sm70_sparse_bmm.py 12 skipped (its decode reference kernel is Volta-only); test_deepseek_v4_sm70_qnorm_rope_kv_insert.py 2 passed; test_compressor_kv_cache.py 32 passed, 6 skipped; test_sparse_indexer_decode_seq_lens.py 11 passed; test_cudagraph_dispatch.py 51 passed.
• Tesla V100: test_deepseek_v4_sm70_indexer.py 30 passed, 1 skipped (the fused-route case needs SM75); test_deepseek_v4_sm70_routes.py 30 passed; test_dsv4_sparse_prefill_bmm.py 17 passed; test_dsv4_sm70_sparse_bmm.py 12 passed; test_deepseek_v4_sm70_qnorm_rope_kv_insert.py 2 passed; test_compressor_kv_cache.py 32 passed, 6 skipped; test_sparse_indexer_decode_seq_lens.py 11 passed; test_cudagraph_dispatch.py 51 passed.
• The two model-head-count cases of test_deepseek_v4_sm70_indexer.py (32 and 64 heads) compare with an fp64 reference that rounds the per-head scores to fp16 as every prefill route does.

pre-commit run --from-ref origin/main --to-ref HEAD: every hook passes except check-layering, which reports vllm/model_executor/layers/fused_moe/runner/moe_runner.py (platform coupling 37 -> 39); this PR does not touch that file and main reports the same. pre-commit run mypy-3.10 --hook-stage manual: passed. The pre-commit CI job fails on main's own run at 899e9ed (run 40532180681) with mypy errors in vllm/v1/kv_cache_interface.py:156-158 and the same check-layering count; none of these files is touched here.

🤖 Generated with Claude Code

https://claude.ai/code/session_0172gtUiymppDzziDqKfy78Q

[KewaiiGamer]