[DSA] Add LiteTopk fused indexer top-k prefill path for SM100

[yuan-luo]

Summary

This PR is adapted from vLLM PR vllm-project/vllm#48726.

Credit to Heisenberg-Yin

Adds LiteTopk as an opt-in DSA indexer top-k path for prefill (ragged extend): fp8 MQA scoring (tcgen05 UMMA) + an online bucketed gate + a compact exact top-k, fused into one pass over KV, so the [num_q, seq_len] logits matrix is never materialized and the memory-budget chunk loop is skipped entirely.

Recall is exact by construction. With the per-row affine $b(x) = \lfloor (x - o) \cdot \delta^{-1} \rfloor$ (code: origin $= o$, inv_delta $= \delta^{-1}$) over negated scores $x = -\text{score}$ (bucket 0 = best), the scan emits exactly ${, j : b(x_j) \le \text{th} ,}$, and every refresh recomputes $\text{th}$ from counts of genuine row elements — so it only ever tightens while staying above $b(x_{(k)})$, the bucket of the true k-th best score:

$$\text{th}_0 \ge \text{th}_1 \ge \cdots \ge b\big(x_{(k)}\big) \quad\Longrightarrow\quad { j : b(x_j) \le \text{th}_t } \supseteq \text{top-}k \text{at every } t$$

A stale threshold therefore admits extra candidates but never drops one, and select trims the superset exactly.

In B200, GLM-5.2 tensors, Q=8192, topk=2048, 1.09–1.25× over dense logits + top-k at 256K–1M contexts, with 8–32 GB of transient logits memory removed.

It is opt-in (SGLANG_ENABLE_DSA_LITETOPK=1), off by default, SM100 (Blackwell) only, GLM DSA shape only (H=32, D=128), and gated to configurations it can serve exactly (num_init_tokens == num_local_tokens == 0); everything else stays on the dense path unchanged, so it adds zero risk to existing deployments.

Background: how LiteTopk replaces the dense top-k

The dense prefill path scores all KV into a [num_q, seq_len] fp32 logits buffer (8–32 GB transient at 256K–1M, forcing a memory-budget chunk loop), then runs top-k over it. LiteTopk streams KV in tiles and keeps only plausible candidates:

1. seed_prep (calibration) — per request, score a bounded KV prefix (min(8192, kv_len)) with the existing dense fp8_mqa_logits, and derive per-row bucket-space parameters (origin, inv_delta) plus an initial gate threshold = the bucket of the k-th best sample score. A sample k-th value can only be worse than the true global k-th, so the initial threshold is a valid loose bound.
2. scan — the fused kernel: UMMA-scores each KV tile, folds the per-row bucket affine into the register weights (scores accumulate directly in bucket units), and emits positions whose bucket passes the gate into a compact candidate buffer via warp-local queues. Two spare warps run a threshold-refresh daemon that walks the live bucket histogram and only tightens the gate as real counts accumulate.
3. select — exact top-k over the compact candidates (radix on float bits in bucket space), -1-padded to topk.

Modifications

New

• python/sglang/jit_kernel/csrc/dsa_litetopk/dsa_indexer_kernels.cuh — the warp-specialized scan kernel (vendored 1:1 from vLLM PR #48726).
• python/sglang/jit_kernel/csrc/dsa_litetopk/dsa_indexer.cuh — seed_prep / select kernels, TMA helpers, shape config (vendored).
• python/sglang/jit_kernel/csrc/dsa_litetopk/entry.cuh — the only sglang-facing layer (TVM-FFI launchers); mirrors the upstream launcher argument order so future syncs stay mechanical.
• python/sglang/jit_kernel/csrc/dsa_litetopk/vendor_deep_gemm/ — minimal vendored DeepGEMM headers (tcgen05 / TMA / UMMA PTX), with the DeepSeek license preserved.
• python/sglang/jit_kernel/dsa/litetopk.py — Python orchestration: per-request calibration loop, scratch allocation, the three primitive calls.
• test/registered/jit/test_dsa_litetopk.py — recall tests vs an fp32 torch reference (details below).

Modified

• python/sglang/srt/layers/attention/dsa/dsa_indexer.py — use_dsa_litetopk gate in Indexer.__init__ + _get_topk_ragged_litetopk branch in _get_topk_ragged.
• python/sglang/srt/environ.py — SGLANG_ENABLE_DSA_LITETOPK (EnvBool, default False).

sglang deviations

• Per-request gate calibration. Upstream calibrates every row from a single shared kv[:sample_len] prefix of the gathered buffer — in a multi-request ragged batch that samples positions a row can never attend, which can over-tighten the gate. Here each request samples its own KV prefix with per-row causal ke, so thresholds are causally valid by construction.
• -1 index padding (upstream pads 0), matching the dense ragged transform contract (fast_topk_transform_ragged_fused) so downstream sparse attention sees the same invalid-slot marker. Output coordinates are identical to the dense path: gathered-KV absolute positions.
• Real SM count for the KV-split heuristic instead of the hardcoded 148 (the kernel template keeps NUM_SMS for 1:1 vendoring).
• Flat-score bucket-span floor (second commit). With all-identical K rows the sample span degenerates to ~0, the original $10^{-20}$ floor makes inv_delta explode, and cross-kernel fp32 accumulation-order noise then spans whole buckets — the gate dropped every candidate. Floor the span at $\sim \text{mag}/256$ (one bucket $\geq \text{mag} \cdot 2^{-16}$, ~16× above the noise) plus an absolute $10^{-6}$ floor. Coarser buckets only loosen the gate, so the fix is recall-safe by construction. Upstream has the same latent issue (only validated on real tensors); found by the adversarial test below.
• Float4 row-alignment padding (third commit). seed_prep reads sample-logits rows with 16 B float4 loads; the per-request sample width min(sample_len, kv_len) is arbitrary, and any width not a multiple of 4 misaligns every odd row base — a device-side fault upstream never hits (its shared prefix has a fixed width). The wrapper pads the width to a multiple of 4 with -inf (already skipped by every pass's isfinite guard), and the entry layer fails loudly on misaligned widths instead of faulting.

Implementation notes

• Recall safety is monotonic end to end: the initial threshold is a loose upper bound from a causally-valid sample; every refresh recomputes from counts of genuine row elements and only tightens; the gate consumes a one-window-stale prefetched threshold (staleness admits extras, never drops); select is exact over the surviving candidates.
• The scan keeps the DeepGEMM 2.5 scoring loop (register-held weights, row-pair TMEM loads, early UMMA release) with non-persistent KV-split scheduling (blockIdx.x = q-block, blockIdx.y = KV window), which keeps all SMs busy on the small-Q chunks long-context prefill actually produces.
• Ragged final q-blocks are handled by forcing padded rows to an empty KV range; TMA out-of-bounds reads zero-fill and are dropped by the per-row range checks.
• The JIT build pins sm_100a and vendors only the DeepGEMM headers the kernel needs; nothing is added to the AOT sgl-kernel wheel.

Validation

Hardware: NVIDIA B200 (SM100), CUDA 13.

Correctness — recall vs an independent fp32 torch reference

test/registered/jit/test_dsa_litetopk.py, registered on the 4-gpu-b200 CI runner (base-b-kernel-unit). The reference recomputes

$$\text{score}[r, j] ;=; \sum_{h} w[r, h] \cdot \mathrm{relu}!\big(q[r, h, :] \cdot k[j, :]\big) \cdot s[j], \qquad j \in [, ks[r],\ ke[r] ,)$$

in fp32 (with $s[j]$ the per-position dequant scale kv_scale); checks are tie-tolerant exactness: every selected index must score $\geq v_k - \varepsilon$, and every index scoring $> v_k + \varepsilon$ must be selected (with $v_k$ the k-th best reference score, $\varepsilon = 10^{-3} \cdot \max(|v_k|, 1)$), plus count / -1 padding / no-duplicate / causal-validity assertions.

• [(32, 32768)] @ topk=2048 — single long request.
• [(16, 8192), (16, 24576), (8, 4096)] @ topk=2048 — multi-request ragged.
• [(32, 32768)] @ topk=512.
• [(8, 1024)] @ topk=2048 — short rows: valid < topk → -1 padding.
• [(9, 4099), (5, 1023)] @ topk=512 — odd sample widths (float4 alignment regression) + ragged final q-block padding.
• Flat-score adversarial ([(8, 16384)], all-identical K rows): every position ties at the k-th value — worst case for the threshold-bucket boundary and candidate-buffer pressure. Caught the bucket-span bug fixed in the second commit (0 candidates selected pre-fix; count + validity + uniqueness green post-fix).

Performance — upstream kernel-level measurement

From the vLLM PR #48726 evaluation (B200, GLM-5.2 tensors, Q=8192, topk=2048), LiteTopk vs dense fp8_mqa_logits + top-k: 1.09× / 1.17× / 1.25× at 256K / 512K / 1M context, with the 8–32 GB transient logits allocation removed. sglang-side e2e numbers will accompany the PR that turns the path on by default.

Notes

• Exact top-k SET, unspecified order: tie-breaking at the k-th value follows atomic arrival order, and the within-row output order is unsorted (the dense radix path is also unordered). Not valid for deterministic mode.
• The exact-set guarantee holds while the candidate buffer bound (cand_cap $= \max(4 \cdot \text{topk},\ 16384)$ per row) is not exceeded; the calibration sample must also be $\geq$ topk (default $8192 \geq 2048$).
• Prefill (ragged extend) only; decode keeps the existing paged path. GLM DSA shape only (H=32, D=128); fp8 index-K cache only.
• The dense path is untouched and remains the default; the env flag flips a per-Indexer boolean, nothing else.

---

CI States

Latest PR Test (Base): ❌ Run #33949270684
Latest PR Test (Extra): ❌ Run #33949270583
Latest PR Test (AMD ROCm 7.2): ❌ Run #33949270696

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Heisenberg-Yin]

Dear Sir,

This is the first author of LiteTopk, Ziqi Yin, I am further working on improve LiteTopk. You can add my wechat ntu_yinziqi or my email ziqi003@e.ntu.edu.sg for further discussion.

Best
Ziqi

[Heisenberg-Yin]

Dear Luo Yuan,

We have updated the code to the latest version in the GitHub (https://github.com/Heisenberg-Yin/LiteTopK) , including an updated implementation of LiteTopK and a newly proposed LiteDSA, which achieves a 1.34–1.4× end-to-end prefill speedup on GLM-5.2. The paper on GitHub has also been updated accordingly.

The updated accelerate ratio is

	256K	512K	768K	1M
ours (ms)	11.74–11.93	21.92–22.04	31.80–31.94	41.79–41.95
vs official	1.13x	1.21–1.22x	1.25–1.28x	1.27–1.28x

Best Regards,
Ziqi

[yuan-luo]

@yuan-luo Please mark the source of litetopk and add @Heisenberg-Yin as co-author

Also for the next time when migrating any codes, it's required that the source must be marked. Otherwise the future PRs will be questioned. Thanks for understanding

@Fridge003 Thanks for reminding, I added @Heisenberg-Yin as co-author in all the commits and marked source of LiteTopK.

outdated

[Heisenberg-Yin]

Dear all,

@Fridge003

I’m very glad to have @yuan-luo helping me upstream this method into SGLang. Recently, I’ve been working with him to integrate support for DeepSeek v4 Flash. My latest implementation achieves a 1.4× speedup for GLM-5.2 1M prefill and a 1.25× speedup for DeepSeek v4 Flash on vLLM.

Also thanks for your attention to my work and this pr.

Best Regards,
Ziqi

[Heisenberg-Yin]

Dear all,

I have tested the latest performance on sglang.

Performance Results

Model	Raw	LiteTopK	LiteTopK Speedup	LiteDSA	LiteDSA Speedup
GLM-5.2 TP8+EP8	134.717s	86.004s	1.566×	77.557s	1.737×
DeepSeek-V4 TP8+EP1	68.931s	49.105s	1.404×	42.761s	1.612×

Note: LiteDSA Speedup refers to the overall speedup of the combined LiteTopK + LiteDSA configuration relative to the raw baseline.

It will be my great honor to work with @yuan-luo to integrate these techniques into SGLang.

Best Regards,
Ziqi