[LopezCastroRoberto]

(this PR description was manually written by me)

Summary

• Adds one persistent CUDA kernel that scores every visible key of a sparse-indexer prefill chunk and selects each row's top-k keys (k = 512, 1024 or 2048) without materializing the [rows, keys] logits matrix.
• One launch covers a prefill chunk of up to 16,384 query rows and up to 1M keys, including multi-request chunks. The dense path needs up to 127 call pairs (fp8_mqa_logits + top_k_per_row_prefill) for the same chunk.
• The output matches the dense path up to ties at the top-k boundary: the gated path ranks 24-bit score codes, and rows the gate cannot serve are rescanned exactly in fp32.

Acknowledgements

We'd like to acknowledge the authors of LiteTopK (arXiv:2607.11976, #48726), who were the first to fuse sparse-attention indexer scoring with top-k selection without materializing the full logits matrix, using classic sample-based threshold filtering. I open a separate PR because:

• Scope: LiteTopK fuses only single-request chunks of exactly 8192/8128 rows with 192K–1M keys. Other chunks keep dense chunking, and even a fused chunk needs a separate seed-logits launch and about 2.6 GiB of workspace. Here every eligible chunk (≤ 16384 rows, ≤ 1M keys, multi-request) is one launch with no logits buffer and a 296 MiB workspace.
• Correctness: on B300 (commit c7184685) I found an undersized shared-memory allocation (4100 memcheck errors), racecheck hazards in the candidate ring and the seed/radix-select kernels, and no recovery from candidate-buffer overflow (the mapping kernel traps). LiteTopK's selection is also approximate (99.998% recall); this PR matches dense up to boundary ties.
• Performance: the only B300 serving data in [Perf][Kernel] Fused DSA indexer Top-k kernel (LiteTopk) #48726 shows TTFT within ~1% of stock at ~200K tokens. This PR measures 1.10× at 256K and 2.05× at 1.04M (different setups, so not like-for-like).

However, most of this kernel is largely an adaptation of DeepSeek's code. This PR replaces the logits write-out with a sampled gate, candidate emission and in-kernel top-k selection.

Motivation: what materializing the scores costs

The dense prefill path writes the fp32 scores of every (query row, visible key) pair to HBM, then reads them back to select the top-k. One 16,384-row prefill chunk needs 16,384 × keys × 4 B of logits:

Keys	Logits for one 16K-row chunk	Dense calls under the 512 MiB budget	Logits HBM traffic (write + read)
16K	1 GiB	2	2 GiB
128K	8 GiB	16	16 GiB
512K	32 GiB	64	64 GiB
1M	64 GiB	127	128 GiB

This produces the following pattern:

It gets worse with context, e.g. at ISL 512K:

Two consequences:

1. The logits budget trades KV-cache capacity against indexer speed. During memory profiling, sparse_attn_indexer allocates a dummy VLLM_SPARSE_INDEXER_MAX_LOGITS_MB buffer, so the logits budget is part of the peak memory the KV cache is sized from. On GLM-5.x with an FP8 KV cache, one token costs 47.7 KB per GPU (78 MLA layers × 576 B + 21 indexer layers × 132 B; MLA KV is replicated across TP ranks). Measured at startup (TP4, B300, --max-model-len 1048576):

   Configuration	Available KV memory	KV cache tokens
   Dense, 512 MiB budget (default)	115.66 GiB	2,603,456
   Dense, 8 GiB budget	109.23 GiB	2,458,816 (−5.6%)
   Fused (this PR)	115.37 GiB	2,596,928 (−0.25%)

   At the default budget the logits buffer fits under the profiling run's other peak allocations, so it costs no KV cache. The cost is speed: rows per dense call shrink as keys grow (2²⁷ / keys), so a 128K-token prefill in 16K-token steps needs 76 dense call pairs and a 1M-token prefill needs 4,217. With the fused path (top-2048), the first four steps, whose rows see fewer than 64K keys, stay dense (21 call pairs) and every later step is one fused call: 4 at 128K, 60 at 1M. Buying fewer dense calls means raising the budget, and that memory comes straight out of the KV cache: an 8 GiB budget (one dense call per 16K-row chunk at 128K keys) costs 144,640 tokens of KV cache. One dense call per chunk at 1M keys would need 64 GiB, more than the KV cache of the entire 1M-token sequence (46.6 GiB). The fused kernel serves long-context chunks in one call each with a fixed 296 MiB workspace.

2. The logits round trip is pure overhead. Only k of the keys scores per row are needed. The fused kernel writes only a gated candidate set, about 1K–4.5K candidates per row depending on k (median on the captures below) at 6 B each, i.e. about 0.1–0.45 GB per 16K-row chunk, against 16–128 GiB of logits traffic at 128K–1M keys.

What this PR does

Kernel (csrc/fused_mqa_topk/). One CTA per SM; each CTA claims 8-row blocks from a global counter (heaviest first) and keeps one continuous Q/KV/TMEM pipeline across blocks:

• Scoring: the FP8 MQA-logits pipeline adapted from DeepGEMM (TMA loads, UMMA into TMEM, warp-specialized producers). Since relu(s) = (s + |s|) / 2, the weighted sum of the linear halves is a single query (the heads folded with their weights) times the key; it runs on the tensor core as extra UMMA rows, so the math warps only accumulate the |s| terms.
• Sampled gate: a strided sample of the block's keys gives each row a threshold that keeps a few thousand candidates. The sample histogram is binned relative to each row's maximum over the first sample tile, so rows whose scores share a large common offset still get a tight gate.
• Candidate emission: keys above the gate are buffered per warp and drained to a per-row candidate slab by two drain warps while the scan continues.
• Selection: one warp per row selects the top-k from its candidates in shared memory (a sampled bracket, then histogram levels in registers), right after the block's scan, before its slot is reused.
• Exact rescan: a row whose candidate count leaves [k, 16384], a row with weights unsafe to scale, or a block whose candidates exceed the selection scratch is rescanned exactly in fp32 (the dense scoring pipeline with a streaming top-k merge), so the output never depends on the gate being right.

Memory

• Fused workspace: a pool of 2 eight-row slots per SM, fixed size, independent of rows, keys and context. It is 296 MiB on a 148-SM B300, allocated once per device during the profiling run, so it is accounted against the KV cache: −6,528 tokens (−0.25%) versus dense at the default budget (table above).
• The profiling run still reserves the dense logits budget: chunks below the dispatch threshold take the dense path.

Limitations of this PR -- TODOs

Anything outside these bounds takes the existing dense path, unchanged.

• Hardware: SM100 family only (i.e., B200/GB200, B300/GB300).
• Models: 32 index heads × 128 dims (e.g., GLM-5.2/5.3, GLM-5.3 kpool, DeepSeek-V4.1-Flash), top-k 512/1024/2048, FP8 indexer cache.
  • Not supported yet: DeepSeek-V3.2 (64 index heads), MXFP4 indexer cache, and the DeepSeek-V4.1 two-level candidate selection layers (below).
• Parallelism: TP, DP (with or without EP) and PCP (with or without DCP) are supported. DCP without PCP is not: each rank holds only a shard of the KV cache, so it keeps the dense path. GLM-5.3-Flash's kpool indexer does not support PCP + DCP itself. Only TP4 is validated end to end.
• Sizes: up to 1M keys and 16,384 rows per call. The row cap is not a kernel-limitation: larger chunks are not validated end to end yet, so are excluded. Only chunks whose longest row sees ≥ 32K keys (≥ 64K for top-2048) use the kernel.

DeepSeek-V4.1 two-level candidate selection. DeepSeek-V4.1-Flash has 8 indexer layers (index_source_layer_ids 2, 8, 14, 20, 24, 28, 32, 36). Only layers 2, 8 and 14 use the kernel at this point, which is why its e2e gain is smaller than GLM-5.x (every GLM-5.x indexer layer uses the kernel). See figure below.

• Layer 20 (candidate source) also writes, per row, the top 2,048 blocks of 8 keys ranked by their max score. The kernel would need a second output and a second selection over block maxima, which rank far deeper than its 512-key gate.
• Layers 24–36 (consumers) mask their scores to those blocks before the top-512. The kernel would need the block list as an input.
• indexer_sparse_logits (opt-in) already speeds up the consumers by scoring only the candidate blocks, but it requires the MXFP4 indexer cache. [Perf] Enable DeepGEMM sparse MQA logits on the FP8 indexer cache, DSv4.1 #62011 addresses this.

Performance

Kernel. B300, real GLM-5.x indexer inputs captured from serving (MRCR 2- and 8-needle and random prompts, layers 0, 9 and 19), one 16,384-row prefill chunk, top-2048, median of do_bench_cudagraph(rep=200) over captures. Dense is chunked by the 512 MiB budget exactly as in vLLM and captured in one CUDA graph.

Keys	Captures	Dense (ms)	Fused (ms)	Speedup, median (range)	Fused exact top-k	Dense exact top-k
32K	9	2.62	2.43	1.05× (1.00–1.21)	9/9	0/9
64K	9	6.12	4.71	1.29× (1.22–1.36)	9/9	0/9
128K	9	13.0	8.76	1.49× (1.41–1.53)	9/9	0/9
256K	9	27.8	16.9	1.63× (1.58–1.66)	9/9	0/9
448K	9	62.8	29.0	2.20× (2.11–2.22)	9/9	0/9
512K	6	78.4	32.8	2.40× (2.34–2.42)	6/6	0/6
944K	3	246	58.9	4.18× (4.13–4.21)	3/3	0/3
1M	3	280	62.0	4.52× (4.38–4.55)	3/3	0/3

"Exact top-k": the selected keys' scores equal the fp32 top-k scores. The current dense path rounds its logits to bf16, so it is not exact at the boundary. See #61149 (comment)

End to end. TP4, B300, FP8 KV cache, prefix caching off. Dense is the upstream base of this PR.

Up to 32K every chunk stays dense (below the dispatch threshold of 64K visible keys). At 64K only the last chunk of each prompt reaches it. At 1.04M, TTFT drops by 711 s (55%).

The two Flash models gain less in the same concurrency sweep, for different reasons:

• GLM-5.3-Flash: every indexer layer uses the kernel, but each one gains less. All 45 layers compute the indexer and all take the kernel. Its kpool indexer scores pools of 4 tokens and selects 512 of them, so a 1M-token prompt is 256K keys at top-512 per layer. There the kernel is 1.7× faster than dense, against 4.4× at 1M keys. The kernel's advantage grows with keys, because dense needs more logits windows and each one re-reads every key. The kernel also only starts at 128K tokens (the 32K-key top-512 threshold).
• DeepSeek-V4.1-Flash: only some indexer layers use the kernel at this point (see detailed explanation above in Limitations/TODOs). 8 of its 40 layers compute the indexer (the rest reuse one). Only layers 2, 8 and 14 take the kernel, at half-length keys (ratio-2 compressed, 512K keys for 1M tokens) and top-512. Extending this kernel to other layers is a TODO for follow-up PRs.
• GLM-5.2/5.3, by comparison: 21 of 78 layers compute the indexer, all 21 take the kernel, at 1M keys and top-2048.

Accuracy

MRCR shows no accuracy difference: fused and dense agree within noise in every context range. Measured before the latest refactors, which keep the output exact up to ties (real-capture validation below).

Context	Samples	Fused	Dense	Difference
0–8K	103	0.394	0.335	+0.059 (+0.9σ)
8K–16K	98	0.220	0.148	+0.072 (+1.5σ)
16K–32K	101	0.416	0.380	+0.036 (+0.5σ)
32K–64K	98	0.495	0.484	+0.011 (+0.2σ)
64K–128K	99	0.789	0.713	+0.076 (+1.4σ)
128K–256K	101	0.768	0.791	−0.022 (−0.5σ)
256K–512K	100	0.724	0.779	−0.054 (−1.1σ)
512K–1M	97	0.683	0.683	0.000
All	797	0.561 ± 0.016	0.539 ± 0.016	+0.022 (+1.0σ)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94835 for commit ccd96c6671be.

ℹ️ This PR is 94 commits behind upstream main. Merged upstream main into this branch as ccd96c6671be before starting CI. Pull the branch before pushing more changes.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/94835

[Heisenberg-Yin]

Hi @LopezCastroRoberto, I put together a three-way diagram to illustrate the relationship between the dense path, LiteTopK, and this PR, to help to understand this PR.

The first two panels are preserved from SGLang PR #32094. The added third panel highlights my understanding of this PR: sampled threshold gating produces candidates, and top-k selection runs inside the persistent kernel as each 8-row group completes, before reusing the slot for the next group.

Sharing it here in case it is useful for the PR description.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98755 for commit c3d2ce273f4d.

ℹ️ This PR is 13 commits behind upstream main. Merged upstream main into this branch as c3d2ce273f4d before starting CI. Pull the branch before pushing more changes.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98776 for commit 516c9ef23600.

ℹ️ This PR is 3 commits behind upstream main. Merged upstream main into this branch as 516c9ef23600 before starting CI. Pull the branch before pushing more changes.

[Pb314314]

@LopezCastroRoberto, adding a concrete memory example to my earlier decode follow-up: the motivation is the remaining BF16 logits intermediate at high concurrency and large speculative K, beyond removing the FP32 copy.

We have complementary PRs addressing redundant FP32 logits materialization: #62172 for cooperative top-k and #64138 for persistent top-k. Direct BF16 consumption still materializes the dense BF16 logits tensor, which contributes to retained CUDA graph pool memory.

For a more representative capacity example, consider the DeepSeek-V4.1 dense indexer path with 1M logits columns, rather than V4's 4×-compressed width. The dense BF16 intermediate alone requires approximately 2 × num_requests × (K + 1) × logits_width bytes for uniform verification, before padding:

Requests	Speculative K	Query rows	BF16 logits alone
512	7	4096	8 GiB
1024	7	8192	16 GiB

These are tensor-capacity calculations, not measured end-to-end allocations or fusion savings. The 8192-row example assumes capture coverage is extended to that size; #63406's current default uniform-decode token cap is 4096. The model-derived 1M-column dense-path capacity is documented in #64138; it does not describe candidate-consuming sparse-logits layers. The memory analysis in #63406 provides the measured graph-memory context. Importantly, this BF16 intermediate remains even after removing all redundant FP32 copies.

I'm preparing an SM100 decode row-chunking PR, which is not pushed yet, to bound this allocation. Chunking may trade memory for additional scheduling/scoring/top-k launches and smaller execution batches; capturing those launches in a CUDA graph does not eliminate their GPU work.

A decode extension of fused scoring + top-k could potentially avoid materializing the full BF16 matrix and offer a better memory/latency trade-off. The actual graph-pool savings and latency benefit still need matched decode benchmarks, including the fused implementation's own workspace. High-concurrency, large-K verification would be a useful follow-up case to evaluate.

AI assistance was used to prepare this comment.

[Gin-Sin]

Nice work. The single-launch design and the exact rescan fallback are a clean improvement over the dense path.

One suggestion, as a reader: since the motivation (avoiding the materialized [rows, keys] logits) and the sampled-gate idea come from LiteTopK (#48726, arXiv:2607.11976), it might help to say so in the Summary, e.g. "Builds on and extends LiteTopK (#48726) …", rather than only in Related Work at the end. A short source comment in csrc/fused_mqa_topk/candidates.cuh, whose candidate-record layout follows LiteTopK's, would also help future readers trace the design. It's a small change, but it makes the lineage clear to anyone skimming the top of the description.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98901 for commit a61a51741fd9.

[khluu]

approved for buildkite changes