[DSA] Fix top-k v2 dropping non-primary ranks' output on CUDA 13.1+ (root cause for #33835)

[DarkSharpness]

Motivation

DeepSeek-V4 DSA decode can return a corrupt top-k from the fused small-batch
cluster kernel, which downstream sparse attention then dereferences as garbage KV
indices (the illegal memory access reported in #33835).

The trigger is the CUDA toolkit version, not the GPU architecture: identical
source and identical hardware (H200, sm_90a) are correct when built with nvcc
12.9/13.0 and wrong when built with 13.1/13.2/13.3.

Root cause

TopKCluster::forward picked its phase-3 scatter destination up front:

const auto cur_out = is_primary ? problem.out : smem->tmp_out;
...
cur_out[pos] = idx;

problem.out can be a shared::cluster (DSMEM) alias of the elected rank's
buffer, while smem->tmp_out is shared::cta. Merging both into one pointer
variable makes cicc 13.1+ mis-lower the block-local arm on sm_90a and silently
drop every non-primary rank's staged output — tmp_out stays zero, and phase
3.5 then faithfully copies zeros to perfectly correct DSMEM addresses.

Slot-diffing one row against a CUDA 12.9 build shows exactly one contiguous bad
run:

correct slots: 129/512      bad-slot runs: [(61, 443, 383)]

Slots 0–60 are the primary's own scatter and 444–511 are the handle_tie tail;
the 383 lost slots are exactly the seven peers' contributions. Read-back
instrumentation confirms the mechanism — tmp_out[0] == 0 on all seven peers
under 13.3 versus valid indices (8161, 10064, 16264, …) under 12.9, while
count_gt reports 34–64 candidates found.

Verified not the cause (each tested and ruled out): enable_smem_spilling,
__restrict__ on TopKProblem::out, the store's address space (forcing a
generic st.u32 at the emit changes nothing), the load side (volatile loads),
and memory ordering (explicit barrier.cluster.arrive.release /
wait.acquire, __threadfence()). Pointer values, start_write/num_write,
and the store addresses were all dumped and are correct in both toolchains —
only the data was wrong.

Changes

1. Keep the two scatter destinations in separate code paths so neither
pointer ever carries two address spaces, with a comment so nobody merges them
back. This is the actual fix.

2. Drop the peer_problem copy from #32910, stating the block-local pointer
at the read-back site instead. That __builtin_assume is load-bearing, not an
optimization: removing it reproduces the #32830 cicc segfault on 13.1/13.2/13.3.

3. Make top-k v1 take a runtime topk (second commit, separable). v1 baked
topk in via -DSGL_TOPK, building one module per k; since the macro fed a
constexpr rather than a template parameter, the modules exported identically
mangled symbols, and setup_kernel_smem_once's function-local static is emitted
as STB_GNU_UNIQUE — merged across every loaded object regardless of
RTLD_LOCAL. The module used second therefore skipped its
cudaFuncSetAttribute opt-in and failed to launch with 64 KB of dynamic shared
memory (CUDA error: invalid argument). bench_topk.py hits this because it
sweeps k in {512, 1024} in one process; a server only ever uses one k.

Scope correction vs #33835

#33835 attributes this to Hopper and to rows just above the 32K small-batch
floor. Measured here, every fused small-batch cluster shape is affected
(batch <= 30 and seq_len > cluster_floor), including 65537 and 98304 which
that PR reports as fine — the floor only makes the path reachable at
batch <= 15. The persistent-pool path was never affected because it stages
output in global memory. With the root cause fixed, no arch gate is needed and
the fused path keeps working on Hopper.

The boundary regression test added by #33835 is still worth keeping; it covers a
real gap (batch <= 15 with seq_len in (32768, 65536]).

Validation

H200 (sm_90a), 157-row suite = the report's shapes + a boundary sweep +
register/streaming/persistent controls, k in {512, 1024, 2048}:

nvcc	before	after
12.9	0 bad	0 bad
13.0	0 bad	0 bad
13.1	157 bad	0 bad
13.2	157 bad	0 bad
13.3	157 bad	0 bad

Those are all five stable nvcc minors >= 12.9. Additionally: 1500-iteration
randomized stress and 500 CUDA-graph replays clean; builds for sm_90a and
sm_100a on every toolchain; v1's runtime-topk path passes 80 cases spanning
k in {1, 2, 7, 31, 64, 100, 255, 256, 257, 300, 511, 512, 513, 777, 1023, 1024}
against trivial / boundary / radix shapes.

No performance regression (CUDA 12.9, where both old and new are correct, so
the comparison is at equal correctness):

case	main	this PR
register 8x8192	3.882	3.884
streaming 100x65536	14.276	14.261
fused-cluster 8x33000	8.387	8.230
fused-cluster 8x98304	9.431	9.309
fused-cluster 30x131072	13.805	13.564
1 long + 7 short (report shape)	9.740	9.575
persistent 128x131072	42.010	42.025

Worst case +0.1% (noise); the fused cluster path gets 1.3–2.9% faster.

Caveats

• sm_100a is compile-verified only — no Blackwell was available to run on.
  Runtime confirmation there would be welcome.
• CI builds cu130, which is a clean cell, so CI cannot currently catch this
  class of bug. Worth considering a toolkit-version guard or moving a build to
  13.1+.
• The remaining exported-weak-symbol hazard in v1 (the kernel host stub itself)
  is not addressed here. Compiling JIT modules with hidden visibility would be
  the general fix; it matters because that half fails silently rather than
  erroring.

Checklist

• Root cause identified and the disproven hypotheses documented
• Verified across every stable nvcc minor >= 12.9, before and after
• Performance measured at equal correctness; no regression
• Blackwell (sm_100a) runtime validation
• CI

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #31310567974
Latest PR Test (Extra): ❌ Run #31310567856

[BBuf]

It‘s a UB fix, LGTM.

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/31310567974/job/93237778346?pr=34167