[Perf][DSA] Unify q8kv8 sparse-prefill topk_length source with the bf16 path; delete backscan kernel

[zkyue]

Motivation

The SM90 q8kv8 (FP8) sparse-prefill path currently derives its per-row valid-topk
early-exit count with a dedicated Triton backscan kernel, gated behind
SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH (default OFF) — so upstream users get no
early-exit at all on this path today.

The bf16 sparse path already solved this without any extra kernel or flag: #31128 passes
metadata.dsa_cache_seqlens_int32 (seqlens clipped to index_topk, computed once per
batch) to flash_mla_sparse_fwd unconditionally. This PR gives the q8kv8 path the same
treatment:

• pass metadata.dsa_cache_seqlens_int32 as topk_length unconditionally (same
  shape-mismatch fallback guard as the bf16 path);
• delete the backscan Triton kernel, its wrapper, and the env flag;
• adapt the backscan sweep test into a metadata-equivalence test.

The early-exit win itself was measured by the path's author in #31888: kernel time
−25.9% @ s_q=437, +2.1% → +6.5% tput at adoption shapes (credit @JackChuang — this PR
only changes where topk_length comes from and makes it the default). Part of the
DeepSeek V4 perf effort (#33636, Indexer & top-k; roadmap #23602).

Net: −3 lines (+107/−110), one fewer kernel launch per layer-call, one fewer env knob, and the
q8kv8 path stops being the only sparse path without the early-exit by default.

Why the metadata source is safe (and equivalent)

• The topk indices of a row with seqlen < index_topk are tail-padded with -1;
  dsa_cache_seqlens_int32 = seqlens.clamp(max=index_topk) can only meet or exceed the
  position of the last valid index, and -1 entries inside the consumed range are
  masked out in-kernel (their KV loads are predicated off) either way. Verified bitwise
  below.
• All-pad rows (DP/CP padding: pad_dsa_cache_seqlens pads with zeros) hand the kernel
  topk_length = 0 where the backscan clamped to 1. Verified on the kernel source and on
  hardware: num_topk_blocks = 0 skips all three warp-specialized loops symmetrically
  (no barrier is arrived at outside the loops), and both 0 and 1 produce identical
  sentinel outputs (out = 0, max_logits = -inf, lse = +inf).
• topk_length is a device tensor consumed in-kernel; host-side control flow stays
  identical across DP ranks (the only host branch is on the tensor's shape, which is
  rank-uniform metadata — same guard the bf16 path has carried since [Perf][DSA] Pass topk_length to flash_mla_sparse_fwd in the sparse attention path #31128).

Evidence (H100, three-arm)

Arms: (a) topk_length=None (today's default), (b) backscan-derived (the deleted
opt-in path), (c) metadata-derived (this PR).

Source equivalence, (b) vs (c): bitwise identical — out, max_logits, lse all
torch.equal across 24 case-runs (8 shape/pattern combos × 3 seeds; patterns: full rows,
ragged, short (≤ topk/4), all-pad-row mixes, plus 2% in-range -1 sprinkle so the
backscan count is strictly smaller than the metadata value on some rows; shapes up to
s_q=2048, topk=2048, s_kv=65536, h ∈ {64,128}, d_qk ∈ {512,576}).

Early-exit numerics, (c) vs (a) — same property the opt-in path already had:
max_logits/lse bitwise identical; out differs by at most 1 bf16 ULP on ~0.01%
of elements, only in rows where block iterations are actually skipped, fully
deterministic (each arm is run-to-run bitwise stable, and a full-width topk_length
tensor is bitwise identical to None). The bitwise sweeps in #31888 pass because their
length patterns keep the iteration count unchanged; a trailing all-masked block pair
perturbs the fp32 O-accumulator by ≤1 ULP of the bf16 output. Flagging this explicitly
rather than repeating the "bit-exact vs None" claim.

Timing (caliber: torch.profiler CUDA kernel-time sum per call, kernel only —
wrapper/derivation excluded; backscan's own derivation kernel listed separately;
s_q=4096, topk=2048, s_kv=65536, h=64, d_qk=576, 30 iters after warmup):

row pattern	(a) none	(b) backscan (+derive)	(c) metadata	(c) vs (a)
full (no pads)	1.861 ms	1.848 ms (+0.006)	1.848 ms	1.007×
half (topk/2)	1.927 ms	1.006 ms (+0.012)	1.005 ms	1.917×
short (≤ topk/4)	1.949 ms	0.415 ms (+0.017)	0.415 ms	4.693×

(c) preserves the full early-exit win of (b) while deleting the per-layer derivation
kernel.

Note on the wrapper's range validation

sparse_mla_q8kv8_prefill_fwd validates 0 <= topk_length <= topk with two
torch.any(...).item() calls — two device→host syncs per layer-call (~70 µs wall
measured on an otherwise-idle stream). This PR leaves that validation untouched: it is
deliberate, tested API surface
(test_q8kv8_sparse_prefill_rejects_invalid_topk_length_bounds), and the DSv4 runtime
backend (#32327, deepseek_v4_backend.py) already passes topk_length through this
same validated wrapper unconditionally. At fully-packed rows (nothing to skip) the sync
overhead is a small wall-time cost (~70 µs against an unchanged 1.85 ms kernel) while
any pad fraction wins 1.9–4.7× kernel time; making the range check sync-free is a
separable wrapper-level discussion affecting both callers, not part of this diff.

Testing

• Both registered q8kv8 test files pass on H100 (SM90): 41/41
  (test_sparse_mla_q8kv8_prefill_sm90.py 27, test_q8kv8_sparse_prefill_backend.py 14).
  • test_q8kv8_topk_length_backscan → test_q8kv8_topk_length_metadata_equivalence:
    asserts bitwise equality of metadata-derived vs exact index-derived counts on
    ragged + all-pad + in-range--1 patterns, and the all-pad sentinel outputs
    (out=0, max_logits=-inf, lse=+inf) with topk_length=0.
• pre-commit run green on all touched files.
• No remaining references to the deleted kernel/wrapper/flag
  (grep -rn "backscan\|SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH\|q8kv8_topk_length" python/ test/ docs/).

---

CI States

Latest PR Test (Base): ❌ Run #32830440901
Latest PR Test (Extra): ❌ Run #32830440631
Latest PR Test (AMD ROCm 7.2): ❌ Run #32830440815