[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) by DarkSharpness · Pull Request #26788 · sgl-project/sglang

DarkSharpness · 2026-05-30T16:06:10Z

Note: after this PR, the v1 should be deprecated. v2 is now a superset of v1. (this line is written by human)

(auto-generated by claude)

Motivation

The DeepSeek-V4 DSA indexer selects the top-k KV positions per query for sparse
attention and maps the selected indices through a page table. This fused top-k +
page-table transform sits on the decode critical path, especially for long context.
This PR provides an optimized JIT kernel for it.

Feature

Runtime top-k — k = 512, 1024, and 2048 are all supported (k is a
runtime argument, so a single module serves every k).
Ragged per-batch sequence lengths.
Arbitrary (power-of-two) page size.
Optional raw (pre-transform) top-k index output.

Performance

Benchmarked on B200 against flashinfer's top_k_page_table_transform
(apples-to-apples: both do top-k and the page-table transform), k=512:

Faster in nearly all (batch, seq) configurations; up to ~1.5× at large batch /
long context.
Sustains ~3.0–3.6 TB/s (near HBM peak) for batch ≥ 2048 at seq ≥ 64K.

Correctness is validated against torch.topk across all dispatch paths
(172 unit-test cases).

CI States

Latest PR Test (Base): ❌ Run #28702891338
Latest PR Test (Extra): ✅ Run #28727693948

Replace the compile-time-K v2 (RegisterTopK/StreamingTopK/ClusterTopK across 5 headers under include/.../topk/) with the cleaner, faster cudalib design (dark @ 3c9225f) in a single impl header + kernel/dispatcher csrc. - Universal: topk is a runtime value (<= 2048), no longer a compile-time K. - Adds page-table transform (emit), per-batch ragged seq_lens, optional raw index side output, all fused into the select kernel. - Fixed cluster size = 8; one cluster per batch element via distributed shared memory (no global workspace / TMA). - Naive dispatch (heuristics intentionally deferred): cluster iff max_seq_len>=64K and batch<=128, else one block per element (trivial/register/streaming by length). - Public Python interface (plan_topk_v2, topk_transform_512_v2 signature) and the metadata/workspace plumbing are preserved for the future persistent dispatcher; the plan kernel is kept (currently unused by the naive transform). - Static shared memory < 48KB for all paths (no cudaFuncSetAttribute). Adds marker-driven bench_topk.py (jit_v1/jit_v2/flashinfer/torch) and test_topk_v2.py (all paths, identity + permutation page tables, ragged).

- Pass scores + seq_lens + page_table as input_args for the JIT providers so do_bench rotates every read tensor across cuda-graph iterations (cold L2); previously seq_lens/page_table sat in the closure and stayed L2-resident. - Exclude page_table from memory_args: only ~topk of its entries are read per row, so counting its full size would overstate the footprint (~1/64 of scores). Footprint = scores + seq_lens read + indices written.

Dynamic dispatch via the plan's cluster_threshold (replaces the flat batch<=128 rule): long items (seq>threshold) go to non-persistent clusters (one cluster per element, ~30 concurrent at occupancy 2 -- a fixed persistent pool measured ~20% slower due to its serial per-cluster loop losing inter-wave pipelining), short items to one block each. This fixes the mid-batch regressions (e.g. batch=128 seq=65536: 733 -> ~1500 GB/s; seq=131072: 1233 -> ~1970) while preserving long-context cluster throughput. Move the page-table transform out of the hot scatter/tie loops into a separate pipelined pass (write raw indices, barrier, then gather+transform), and trigger the dependent kernel (PDL) only after the full output is written -- a correctness fix vs the old inline trigger which fired before ties were emitted. Streaming bandwidth now ~3.2-3.3 TB/s for bs>=2048 at seq=65536 (k=512); longer contexts are bounded by the two-pass streaming (~2.9 TB/s, matching cudalib). All paths verified (trivial/register/streaming/cluster, identity+permutation page tables, ragged, k in {512,1024,2048}); 32 regs / 0 spill / occupancy 2.

… read) Port the finalized cudalib register kernel (dark@549130b): TopKRegister<kLocalVecs> with the full-vector load + scalar-tail-on-last-lanes loop and a collect lambda. The per-element dispatch now routes seq in (8192,16384] to Register<4> (register- resident, single read) instead of the two-pass streaming path -- e.g. b=2048 s=16384 k=512 now uses 1 read. TopKStreaming derives from TopKRegister<2>. 32 regs / 0 spill / occupancy 2; all paths still correct.

Specialize the per-element kernel on the host-known max_seq_len bucket (kLevel): 0 (<=8K) compiles trivial+register<2> only, 1 (<=16K) adds register<4>, 2 (>16K) adds streaming. Smaller buckets emit a leaner kernel -- e.g. the decode bucket (kLevel 0) drops to 24960 B smem (no spill reservation) vs 45440 B for the full kernel -- giving ~1-3% on short-context decode (b=8192 s=8192 k=512: 1523 -> 1563 GB/s). Cluster path unchanged (separate kernel).

…adata Replace the per-element cluster launch ({batch_size, 8} blocks, most of which exit in the short/mixed case) with a FIXED persistent pool of kNumPersistentClusters (30 = 15*2, occupancy 2) clusters that loop to fetch the plan's compacted long-item work list (stride gridDim.x). The all-short case now wastes only the pool in the first wave (~0.2us) instead of batch_size*8 cluster blocks; PDL overlaps it; all-long small-batch finishes in the first wave. Measured (B200, k=512): pool=30 beats pool=15 across the small-batch cluster regime (b=30 s=262144: 1372 vs 932 GB/s) -- 15 clusters under-fill the GPU. Short-context fallback inside use_cluster verified: width=131072/seq=4096/b=128 = 5.13us vs 4.91us non-cluster baseline (~0.2us pool no-op), all correct. Plan now compacts items (batch_id, seq_len) into the work list. Metadata shrunk to (batch+1, 2) int32 / 8B rows (dropped reserved padding); the pool reads seq_len from the list (no seq_lens re-read). TopKCluster::forward is loopable (no in-forward PDL/barrier; the kernel does PDL once + cluster.sync per item). 32 reg / 0 spill / occupancy 2. pytest 37 + stress 0 errors.

Set each candidate's max_batch_size to kNumPersistentClusters (30). The old caps (16 at the 64K floor) left batch~17-30 small-batch long-context items on the streaming path even though clustering them (one pool wave) is faster. Measured (k=512): batch=30 seq=73728 511->596, seq=98304 577->768 GB/s; batch=64 still streams (correct -- >pool would serial-loop). bs>=2048 streaming target and the short-context fallback are unaffected. Pool size 30 (2 waves) beats 15 across the small-batch cluster regime. pytest 37 + stress 0 errors.

Fuse the cluster/non-cluster split for small batch into one kernel and tune the cluster path; the bs>=2048 streaming target is unchanged (3.1-3.6 TB/s). - Small-batch kernel: batch<=30 (= kNumPersistentClusters) runs one 8-block cluster per element, each rank writing its result into the elected rank's smem via map_shared_rank, then that rank transforms -- short items use one elected block, long items the full cluster. 30<batch<=512 keeps the persistent pool + main kernel (kLevel=3) epilogue. Output stays on-chip (smem scratch -> regs -> one global write), no raw round-trip. - occ2 for the fused kernel (was occ1): one cluster per element is GPC-bound at occ1 (~8 clusters/wave, batch>8 serializes); occ2 fits ~30 (=15*2). batch 16-30 cluster regime -28..-31% (seq 98304: 18.2->13.0us; 131072 b30: 20.1->14.2us). - Batch-aware cluster floor: 32K for batch<=8, else 64K. Measured fused crossover ~32K @ b<=8 vs ~56K @ b16-30; b<=8 @ 64K 12.8->9.1us, no b16-30 regression. Threaded through a runtime TopKLaunchParams.cluster_floor. - Seq-dependent plan caps (was a flat cap = pool size): the pool runs N long items in ceil(N/30) waves and longer seq is worth more waves, so the cap grows with seq {65536:30, 98304:48, 131072:60, 196608:80, 262144:112, 393216:128}. Fixes under-clustering -- half-long-b64 49->32us, uniform[8k,256k]-b64 45->35us, 524288 -b128 now clusters (126 vs 133us stream). auto tracks min(cluster,stream), 0 regress. - MaxSmem<Ts...> helper: static smem sized+aligned for the max of several impl Smem types (Streaming vs Cluster differ), replacing an unguarded largest-pick. - Fix dropped top-k index on the persistent path: the cluster loop was missing the inter-item barrier forward() requires, so a cluster reusing its radix smem before all 8 ranks finished the prior item gave an occasional 511/512. forward()'s cross-rank DSMEM phases are internally cluster.sync()'d, so a conditional __syncthreads() (block-local, between items only) suffices and is cheaper than a cluster-wide barrier. Repro 0/300 (all-long) + 0/200x5 (few-long). - Hardened test to 162 cases: every template at its boundaries (8192/8193, 16384/16385, floor+-1, scalar tail), both dispatch shapes, per-row page tables, k in {512,1024,2048}. pytest 162 + randomized stress + race (now deterministic) + dropped-index repros all clean.

…nsform) The flashinfer/torch baselines did plain top-k with NO page-table transform, so the public benchmark unfairly penalized the JIT kernels (which do top-k AND a page transform). Make all four providers do top-k + a page transform: - flashinfer: flashinfer.top_k_page_table_transform (its fused topk + transform). Its API is page_size=1, so it gathers through a per-token (batch, seq) src_page_table; row_to_batch=None keeps it on the B200 cluster path. - torch: torch.topk followed by a page_size=1 gather. - jit_v1 / jit_v2: unchanged (production page_size=64). The page-table tensor is still excluded from the reported footprint (only ~top-k of its entries are read per row). With the transform now counted for all providers (k=512, seq>=128K): jit_v2 wins every batch except the bs=32/64 valley -- where flashinfer's adaptive 8/4-way cluster split beats our fused->persistent transition -- and wins large batch by ~1.5x (bs=512 seq=256K: 177 vs 275 us).

…mit_invalid Cleanup, no functional change (pytest 162 + randomized stress 0 errors): - Remove all references to the source/upstream repo from topk.cuh's docstring and the kWarpSize comment. - Remove the dead global workspace: topk_transform_512_v2 allocated a per-call (bs, 2 + 2*1024) int32 buffer that the kernel ignored (it uses distributed shared memory, no global scratch). Dropped from the Python wrapper and the C++ transform() signature (was a [[maybe_unused]] param). - Remove the always-null raw_out side output from TopKProblem / transform_output: v2 is dispatched only when raw indices aren't needed (indexer gates on raw_indices is None), so it was never set. v1's out_raw_indices is untouched. - Remove the unused emit_invalid() (orphaned when TopKTrivial was replaced by trivial_transform). Integration with the public branch is unchanged: the indexer/metadata call sites and the public wrapper signatures (topk_transform_512_v2, plan_topk_v2) match main; only the internal workspace arg was removed, which was never part of the Python signature.

Add an optional raw-index output to topk_transform_512_v2, mirroring v1's out_raw_indices: a tvm::ffi::Optional<tvm::ffi::TensorView> in C++ and a torch.Tensor | None (default None) in Python. When provided, raw_out[i, j] is the pre-page-transform position selected for slot j (or -1 for padding), written alongside the page-transformed page_indices in the same transform_output pass. Plumbed through TopKLaunchParams.raw_indices -> TopKProblem.raw_out, so it is filled on every path (trivial / register / streaming / fused small-batch cluster / persistent epilogue). nullptr (None) is the default and costs nothing. test_topk_v2_raw_indices verifies out[j] == page_to_indices(raw[j]) per slot (and the -1 padding aligns) across all paths. pytest 172 passed.

Rename the device-side implementation header (impl classes + helpers) and update the include in topk_v2.cuh. No code change; pytest 172 passed.

gemini-code-assist

Code Review

This pull request refactors and optimizes the DeepSeek-V4 top-k transform kernels, introducing a universal JIT module where topk is a runtime argument and consolidating the implementations into a unified topk_impl.cuh header. It also adds benchmark and correctness test suites. The review feedback highlights critical synchronization issues with Program Dependency Link (PDL) triggers, specifically premature triggers in topk_persistent_cluster_kernel and topk_main_kernel that cause race conditions, and a missing trigger in topk_small_batch_kernel that can lead to deadlocks. Additionally, the reviewer noted a type inconsistency for seq_lens in topk_plan that should be updated to int32_t* for consistency.

- Move enable_smem_spilling into device:: (utils.cuh) so every radix kernel shares one helper. - Add LaunchKernel.config({.use_pdl, .cluster_dim}).launch(...) and use it for the cluster launches in place of .enable_cluster().enable_pdl()(...). - Rename CombinedTopKKernel -> TopKKernel; expand the TensorMatcher checks to one field per line with shape comments. - Make the persistent-cluster loop barrier unconditional (#pragma unroll 1 + always __syncthreads()); the trailing barrier is harmless at loop exit. - Add a file docstring crediting the vllm and flashinfer reference implementations. pytest 172 passed.

…loor 8 -> 15) The fused small-batch path launches one 8-block cluster per batch element; on B200 exactly 15 such clusters fit in a single wave (occupancy 2). A fine batch sweep shows a sharp +25-40% latency step from batch 15 -> 16 as the 2nd wave starts. For batch <= 15 (one wave) the kernel is latency-bound, so the 8-way split beats streaming from a much lower seq. A direct force-cluster vs force-stream measurement puts the crossover at ~36-40K for batch <= 15 (vs ~64K for batch >= 16). Raising the low-floor threshold from 8 to 15 routes batch 9-15 to the 32K floor: batch 15 @ seq 64K: 12.9us -> 9.5us (-26%) batch 12 @ seq 56K: 11.9us -> 9.4us (-21%) with no regression at/below batch 8 or at batch >= 16 (kept on the 64K floor, which the same sweep confirms: clustering there loses by ~2us at 36-49K). pytest 172 passed.

…baselines Register bench_topk.py to base-b-kernel-benchmark-1-gpu-large (est 120s) so run_suite.py collects the new JIT benchmark (an unregistered bench file is a CI gap). Benchmark jit_v2 against jit_v1, flashinfer and torch. flashinfer and torch gather per token (page_size=1), so hand them a (batch, seq) page table rather than the page-size-64 one the JIT kernels use -- this fixes a torch gather out-of-bounds; jit_v1 skips k=2048 (unsupported).

Line-wrapping only (black 26.1.0, matching the pre-commit pin); no logic change. Fixes the Lint job.

Add test_topk_v2_output_indices: checks the optional raw (pre-transform) index output directly against torch.topk across every dispatch template/boundary, isolating the top-k selection from the page-table transform (the existing tests validate the transformed page output and invert it back). Uses a permuted page table so raw != out, catching any leak of transformed indices into the raw buffer. pytest 244 passed.

The v2 top-k test and benchmark registered CI suites (register_cuda_ci) but lived inside the importable sglang package (python/sglang/jit_kernel/...). The check-no-registered-tests-in-package lint hook rejects that: registered tests/benchmarks must live under test/registered/ so they are collected by run_suite.py and not shipped in the wheel. Move (pure rename, contents unchanged): - python/sglang/jit_kernel/tests/deepseek_v4/test_topk_v2.py -> test/registered/jit/deepseek_v4/test_topk_v2.py - python/sglang/jit_kernel/benchmark/bench_topk.py -> test/registered/jit/benchmark/bench_topk.py Both use absolute sglang.* imports, so the relocation needs no import changes. check_no_registered_tests_in_package.py and check_registered_tests.py both pass locally. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MQgtB3k9vngYs9MyAu7BSo

…tration The main merge tightened the registered-test validator: CUDA tests must use the modern stage=/runner_config= form instead of the legacy suite="..." name (a legacy suite matches no PR-test workflow, so the test silently never runs). Migrate the two topk_v2 files to match their siblings (test_c4_v2.py, bench_dsv4_fp4_indexer.py): - test_topk_v2.py: suite="base-b-kernel-unit-1-gpu-large" -> stage="base-b-kernel-unit", runner_config="1-gpu-large" - bench_topk.py: suite="base-b-kernel-benchmark-1-gpu-large" -> stage="base-b-kernel-benchmark", runner_config="1-gpu-large" check_registered_tests.py and check_no_registered_tests_in_package.py both pass locally. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MQgtB3k9vngYs9MyAu7BSo

DarkSharpness · 2026-07-03T07:32:54Z

/rerun-failed-ci

Fridge003 · 2026-07-04T22:08:40Z

/rerun-failed-ci

Fridge003 · 2026-07-05T02:35:38Z

Extra test run https://github.com/sgl-project/sglang/actions/runs/28702891231?pr=26788

…form (runtime k <= 2048) (sgl-project#26788) Co-authored-by: Claude <noreply@anthropic.com>

claude added 12 commits May 30, 2026 09:02

refactor(dsv4 topk v2): rename impl header topk.cuh -> topk_impl.cuh

2aca0de

Rename the device-side implementation header (impl classes + helpers) and update the include in topk_v2.cuh. No code change; pytest 172 passed.

DarkSharpness requested review from BBuf, HydraQYH, celve and yuan-luo as code owners May 30, 2026 16:06

github-actions Bot added the jit-kernel label May 30, 2026

gemini-code-assist Bot reviewed May 30, 2026

View reviewed changes

claude added 3 commits May 30, 2026 11:25

DarkSharpness added the run-ci CI: run the baseline test suite on this PR label May 30, 2026

claude added 2 commits May 30, 2026 11:38

style(dsv4 topk v2): black-format test_topk_v2.py

68271b1

Line-wrapping only (black 26.1.0, matching the pre-commit pin); no logic change. Fixes the Lint job.

DarkSharpness changed the title ~~[Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048)~~ [JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) May 31, 2026

Fridge003 added the high priority label Jun 22, 2026

Fridge003 and others added 4 commits June 22, 2026 01:23

Merge branch 'main' into jit_topk_v2_perf

02135dd

Merge remote-tracking branch 'origin/main' into pr-26788

6dc19b9

root and others added 2 commits July 3, 2026 04:34

Merge remote-tracking branch 'origin/main' into pr-26788

c6a2eec

Merge branch 'main' into jit_topk_v2_perf

2653abe

Fridge003 added run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fastfail labels Jul 4, 2026

Fridge003 added the release-highlight Candidate PR for release note highlight label Jul 5, 2026

Fridge003 approved these changes Jul 6, 2026

View reviewed changes

Fridge003 merged commit c016c6f into sgl-project:main Jul 6, 2026
465 of 583 checks passed

amd-bot mentioned this pull request Jul 6, 2026

[CI Monitor] Daily Report - 2026-07-06 bingxche/sglang-ci-bot#127

Open

DarkSharpness mentioned this pull request Jul 6, 2026

[DSA] Fold page-table into fused top-k v2 (decode): drop page_size=1 expansion #30274

Merged

amd-bot mentioned this pull request Jul 7, 2026

[CI Monitor] Daily Report - 2026-07-07 bingxche/sglang-ci-bot#128

Open

b8zhong mentioned this pull request Jul 7, 2026

DO NOT MERGE: revert c016c6f355 to test gb300 glm5-nvfp4 NaN regression #30418

Closed

1 task

amd-bot mentioned this pull request Jul 14, 2026

[CI Monitor] Daily Report - 2026-07-14 bingxche/sglang-ci-bot#135

Open

zkyue mentioned this pull request Jul 14, 2026

[BugFix][DSA] Harden top-k v1/v2 kernels against negative padded seq_lens #31123

Closed

5 tasks done

amd-bot mentioned this pull request Jul 15, 2026

[CI Monitor] Daily Report - 2026-07-15 bingxche/sglang-ci-bot#136

Open

yctseng0211 mentioned this pull request Jul 15, 2026

[AMD] DeepSeek-V3.2: disable top-k v2 on rocm (unblock stage-c AMD PR Test) #31291

Closed

5 tasks

This was referenced Jul 16, 2026

[CI Monitor] Daily Report - 2026-07-16 bingxche/sglang-ci-bot#137

Open

[CI Monitor] Daily Report - 2026-07-17 bingxche/sglang-ci-bot#138

Open

weireweire mentioned this pull request Aug 5, 2026

[DSV4] Support raw-index output in TopK v2 #33672

Merged

Emilio597 mentioned this pull request Aug 6, 2026

fix(dsv4): avoid fused TopK v2 on Hopper #33835

Closed

4 of 5 tasks

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table trans…

5eeb9e4

…form (runtime k <= 2048) (sgl-project#26788) Co-authored-by: Claude <noreply@anthropic.com>

dahaobenhao mentioned this pull request Sep 5, 2026

[Feature] Tuning / per-GPU config for DSv4 top-k v2 cluster launch plan (kClusterFloor / kNumPersistentClusters / kCandidates) #37150

Open

Navigation Menu

Repository navigation

[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) - #26788

[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048)#26788
Fridge003 merged 23 commits into
sgl-project:mainfrom
DarkSharpness:jit_topk_v2_perf

DarkSharpness commented May 30, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DarkSharpness commented Jul 3, 2026

Uh oh!

Fridge003 commented Jul 4, 2026

Uh oh!

Fridge003 commented Jul 5, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

DarkSharpness commented May 30, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Feature

Performance

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DarkSharpness commented Jul 3, 2026

Uh oh!

Fridge003 commented Jul 4, 2026

Uh oh!

Fridge003 commented Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

DarkSharpness commented May 30, 2026 •

edited by github-actions Bot

Loading

Fridge003 commented Jul 5, 2026 •

edited

Loading