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

[zkyue]

Motivation

Related to #25574 (topk_transform_512_v2 illegal memory access on SM100 under DP-attention). We set out to reproduce and fix that issue on B200 (SM100) and found the following:

1. The code quoted in DSv4 topk_v2 fused-cluster kernel crashes on B300/SM100 with DP-attn + DeepEP #25574 no longer exists. [JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) #26788 rewrote topk_v2.cuh; the old topk_fused_transform SMALL/TRIVIAL early-return branches are gone, and the rewritten fused small-batch kernel already has cluster-uniform fall-through control flow. Note the branch condition is per-blockIdx.x and all 8 blocks of a cluster share blockIdx.x, so even in the old kernel every block of a cluster took the same branch — there was never a divergent cluster barrier (compute-sanitizer --tool synccheck on the pre-[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) #26788 kernel reports 0 errors, and 24k+ stress iterations over mixed TRIVIAL/SMALL/LARGE batches, eager + CUDA-graph replay, never faulted on B200). The B300-specific crash could not be reproduced at kernel level on B200, and its proposed patch (fix(dsv4 topk_v2): honor cluster contract in fused kernel SMALL/TRIVIAL branches #25575) targets the deleted code.

2. What does reproduce — deterministically, at current main — is the negative-seq_lens crash class. Negative per-row lengths (DP-padded / idle-companion rows; [DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 #30378 observed -4 from GLM 5.2 MTP draft-extend metadata, and DP-attention idle rows are the same class) are read through unsigned conversions in both top-k kernels and become ~4e9-token rows:

   • v1 (topk_v1.cuh): illegal memory access for any negative length (notably, the SGLANG_OPT_USE_TOPK_V2=0 fallback is equally exposed);
   • v2 (topk_v2.cuh): illegal memory access (e.g. -1048576, on every dispatch shape: fused small-batch cluster, persistent pool, streaming) or — worse — a silently garbage top-k row (-1/-4, where the unsigned chunk arithmetic happens to wrap to a zero-length chunk and handle_tie pads the row with bogus indices).

   [DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 #30378 clamped two DSA triton producers caller-side, but the kernels stay exposed to every other caller — e.g. the DSv4 DP-attention path (dsv4/metadata.py feeds c4_seq_lens to plan_topk_v2 unclamped).

Modifications

Per review feedback, negative lengths are handled with signed comparisons (no arithmetic clamp) so a negative row takes the documented trivial all-(-1) "no tokens" path -- every non-negative candidate index compares >= -1, so none qualifies:

• topk_v1.cuh: seq_len stays int32_t; the signed trivial dispatch (seq_len <= (int32_t)kTopK) routes negative rows into naive_transform, whose fill comparison is signed; the radix path casts to uint32_t only after the > kTopK guard proves the value positive;
• topk_v2.cuh: a signed is_trivial() ((int32_t)seq_len <= (int32_t)topk) used by the main and fused small-batch kernels; trivial_transform's per-element comparison is signed;
• topk_v2.cuh topk_plan (parameter type fixed to const int32_t* to match the tensor dtype): the threshold-count loop compares signed against the candidate thresholds (all int32-exact); the pool-compaction predicate proves a row positive (signed) before the unsigned threshold compare, so negative rows are never routed to the persistent pool as ~4e9-token items while a full-uint32 static_cluster_threshold keeps its meaning.

Python docstrings updated (producers should still clamp to 0 themselves; the kernel no longer turns a producer bug into an illegal address or silent corruption). Cost is unchanged -- paired timing below.

New tests: test_topk_v2_negative_lengths (8 configs x k in {512, 2048}, one config per dispatch shape: level 0/1/2, fused small-batch cluster with -1/-4/-1048576, persistent pool), test_topk_v2_plan_negative_lengths (pins the plan kernel's counting loop via a candidate-cap-boundary distribution and the plan metadata, plus a full-uint32 static-threshold boundary check), and test_topk_v1_negative_lengths. Output buffers are poison-prefilled; each test asserts negative rows come back all-(-1) (both out and raw) and non-negative rows in the same batch still match torch.topk.

Accuracy Tests

On B200 (CUDA 13.3), before this patch (current main):

• [40000]*16 + [-1048576, 400] (streaming), [131072, 9000, 9000, -1048576] (fused cluster), [131072]*20 + [400]*10 + [-1048576] (persistent pool): CUDA error: an illegal memory access was encountered — deterministic;
• [131072, 9000, 9000, -4] / [..., -1]: padded row returns 512 garbage indices (silent corruption);
• v1 with [9000, 9000, 9000, -4]: illegal memory access;
• on the pre-[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) #26788 kernel (the code DSv4 topk_v2 fused-cluster kernel crashes on B300/SM100 with DP-attn + DeepEP #25574 was filed against), [131072, 9000, 9000, -4] reproduces the exact reported failure signature (illegal address surfacing at the fused-kernel launch).

After this patch:

• all of the above complete, negative rows = all (-1);
• full registered suite test/registered/jit/deepseek_v4/test_topk_v2.py: 262 passed (244 pre-existing + 18 new);
• stress: 14 scenarios (branch mixes, batch-15 boundary, zero rows, negative rows) x eager + CUDA-graph replay: pass;
• compute-sanitizer memcheck + synccheck on negative-row and mixed-branch cluster scenarios: 0 errors.

Speed Tests and Profiling

Paired A/B timing (median of 5x50 launches, 4 interleaved rounds, B200), main vs this PR, us/launch:

shape (batch x seq)	dispatch	main	this PR
128 x 8192	level 0 (Register2)	5.41-5.72	5.24-5.45
64 x 16384	level 1 (Register4)	6.19-6.37	6.20-6.24
100 x 65536	level 2 (Streaming)	13.85-14.16	13.84-13.86
8 x 131072	fused small-batch cluster	10.82-10.96	10.80-10.83
64 x 131072	persistent pool + main	23.95-24.06	24.07-24.10
4 x 9000	decode-like level 1	5.07-5.09	5.04-5.09

All deltas within run-to-run noise on a shared machine.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #29328239541
Latest PR Test (Extra): ❌ Run #29328239210

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[DarkSharpness]

seq_len < 0 is sick. In this case I would recommend fix the upstream (NOTE that, not only topk kernel consumers seq_len, other kernels may suffer from the same bug). Another workaround with 0 overhead is to change the trivial path into signed integer comparison (fill all indices with -1, since non-negative number >= -1)

[zkyue]

Thanks — agreed, and adopted in the latest push.

Signed comparison adopted: no arithmetic clamp remains anywhere; negative rows are routed by comparisons alone:

• topk_v1.cuh: seq_len stays int32_t; the trivial dispatch (seq_len <= (int32_t)kTopK) routes negative rows into naive_transform, whose fill comparison is now signed — every slot gets -1, exactly as you describe (a non-negative index never qualifies against a negative length). The radix path casts back to uint32_t only after the > kTopK guard has proven the value positive.
• topk_v2.cuh transform: a small is_trivial() helper ((int32_t)seq_len <= (int32_t)topk) used by both topk_main_kernel and topk_small_batch_kernel; trivial_transform's per-element comparison is signed as well. The raw bit pattern is preserved in TopKProblem::seq_len; no path reachable by a negative row touches the length before the signed dispatch catches it.
• topk_v2.cuh plan: the threshold-count loop compares signed against the candidate thresholds (all int32-exact constants). The one place a lone signed compare was NOT sufficient is the pool-compaction predicate: static_cluster_threshold is a full uint32_t, so an int32 cast of the threshold would misroute for values > INT32_MAX. It therefore proves the row positive first and keeps the threshold compare unsigned — sl > 0 && (uint32_t)sl > cluster_threshold — which is bit-identical to the previous behavior for positive rows and never routes a negative row. Covered by a new static-threshold boundary assertion in the plan test.

Re-validated the full battery on B200: 262/262 registered tests (the 18 negative-length regressions stay green, incl. the plan-metadata pin), the illegal-address/garbage repros, CUDA-graph replay, compute-sanitizer memcheck + synccheck (0 errors), and paired timing on all six dispatch shapes (unchanged within noise).

On fixing upstream: fully agreed the producer is the right long-term fix — a padded row should never carry a negative length in the first place. #30378 clamped two GLM DSA triton producers (triton_ops/pad.py, triton_ops/dsa_metadata.py), but the DSv4 DP-attention path still feeds c4_seq_lens into plan_topk_v2 unclamped (dsv4/metadata.py), which is why I would keep the kernel-side signed handling as defense-in-depth: it turns a producer bug from an IMA — or worse, a silently wrong top-k row (the -4 case wraps into a garbage row rather than crashing) — into the defined all-(-1) output. Happy to file the DSv4 producer-side fix as a follow-up PR if you want it.

On other consumers: your warning matches what we saw while auditing. The same unsigned-length pattern exists elsewhere in the DSv4 JIT family, e.g. paged_mqa_metadata.cuh:108 hands the seq_lens buffer over as uint32_t*, and common.cuh:60 / c128_online_v2.cuh:544 do const uint32_t seq_len = seq_lens[...] (immediately followed by prefix_len = seq_len - extend_len, which can also wrap on inconsistent inputs). Can run a follow-up sweep over those consumers as well if useful.

[zkyue]

Closing this one out. Since it was opened, the concrete hazards in this kernel have been addressed directly upstream — #30645 (tie overflow / inf scores → invalid indices) and #34167 (CUDA 13.1+ non-primary-rank output drop) — and #24654's SGLANG_DSA_TOPK_BROADCAST covers the cross-TP tie-divergence observation that motivated part of this hardening. We also don't currently have a confirmed producer of negative padded seq_lens on the CUDA path at main, so a defensive guard without a live repro doesn't meet the bar for merging.

Closing to keep the queue clean; happy to revive with a concrete reproducer if one shows up.