fix(dsv4): avoid fused TopK v2 on Hopper

[Emilio597]

Motivation

DeepSeek-V4 decode can produce incomplete DSA TopK v2 output on Hopper when a small mixed batch crosses the 32K C4 cluster floor. In an 8x H200 TP8 deployment, this surfaced as a reproducible downstream CUDA illegal-memory access when one long request crossed raw token 131072 while seven short requests remained active.

The failure is in the current fused small-batch DSMEM path, not in Marlin, Mooncake, FP8 KV cache, or the general persistent-cluster implementation.

Root cause

Calling topk_transform_512_v2 directly on H200 with an identity page table gave the following result:

C4 seq_len:       32767  32768  32769  33000  2000  12000  16385  1000
unique/512 slots:   512    512    115    207    512    512    512   512

For batch_size <= 15, the small-batch cluster floor is 32768. The two rows just above that floor enter Cluster::forward in topk_small_batch_kernel and return incomplete output on SM90. Register and streaming rows in the same launch are correct.

Additional isolation:

• Moving the PDL wait before the first seq_lens read did not repair the two failing rows.
• The existing persistent-cluster path was correct: batch=31, seq_len=65537 returned 512 unique valid indices for all 31 rows.
• Disabling TopK v2 made the full workload stable, and restoring the original v2 dispatch reproduced the crash at the same C4 boundary.

The direct kernel corruption is deterministic. The link from stale/duplicate TopK slots to the downstream sparse-attention illegal access is supported by the full-model boundary A/B.

Changes

• Keep the fused small-batch path on SM100+.
• On Hopper, route the same shapes through the existing persistent-cluster plus main-kernel path.
• Add a regression test with mixed row lengths immediately below, at, and above the 32768 small-batch cluster floor.

Blackwell dispatch is unchanged.

Validation

Tested environment:

GPU: 8x NVIDIA H200 (SM90a)
Model: DeepSeek-V4-Flash-0731
SGLang base: 7adf2f4a9a4389d1c021a2caf128cf7e2f4adb35
PyTorch: 2.11.0
CUDA toolkit: 13.1
TP: 8
KV cache: FP8 E4M3
MoE runner: Marlin
Context length: 400000
Decode CUDA Graph: enabled

Results with the Hopper fallback:

• Direct mixed-boundary kernel check: all 8 rows returned 512 unique indices, all in range.
• Deterministic boundary reproducer (1 x 130800 -> 10000 plus 7 x 2000 -> 10000): all requests completed; 80,000 output tokens; no CUDA error.
• Full stability gate (40 x 200000 -> 40000, concurrency 40): 40/40 requests completed; 8,000,000 input and 1,600,000 output tokens; 1566.01 output tok/s; HTTP health remained 200.
• Persistent-path regression (1 x 298933 -> 10000 plus 7 x 9731 -> 10000): all requests completed, including the C4 > 65536 row.
• The same full stability gate with Mooncake HiCache write-through: 40/40 completed; actual Mooncake SSD read/write counters increased; no CUDA error.

The deployed validation branch also contained a PDL-wait placement cleanup inside the fused kernel. A PDL-only A/B did not fix the corruption, and the fused kernel is not launched on SM90 after this dispatch change. This PR intentionally excludes that unrelated cleanup and contains only the minimal Hopper fallback.

The tested topk_v2.cuh and regression-test base blobs are unchanged on current main as of 2026-08-06.

Related work

• [JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) #26788 introduced the current TopK v2 implementation.
• DSv4 topk_v2 fused-cluster kernel crashes on B300/SM100 with DP-attn + DeepEP #25574 / fix(dsv4 topk_v2): honor cluster contract in fused kernel SMALL/TRIVIAL branches #25575 concern the pre-[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) #26788 fused kernel and code that no longer exists.
• Fix DSV4 topk v2 fused PDL trigger #25785 adds a PDL trigger for the older implementation; it does not address this output corruption.
• [BugFix][DSA] Harden top-k v1/v2 kernels against negative padded seq_lens #31123 handles negative padded seq_lens, a different failure class.
• [DeepSeek-V4] Fix nvcc 13 crash building the topk_v2 kernel #32910 fixes the CUDA 13 compiler crash when building this kernel, separate from the runtime correctness issue here.

Checklist

• The change is scoped to the affected architecture and dispatch path.
• A focused regression test is included.
• git diff --check passes.
• Full-model correctness and long-running serving validation are reported above.
• CI (this PR).

---

CI States

Latest PR Test (Base): ❌ Run #31081928822
Latest PR Test (Extra): ❌ Run #31081928039

[DarkSharpness]

Generated by Claude.

B200 (sm_100a) runtime result: the corruption does not reproduce on Blackwell, so keeping the fused path enabled on SM100 — as this PR does — does not ship a silently-wrong Blackwell. This also closes the "sm_100a is compile-verified only" gap left open in #34167.

Since the trigger is the CUDA toolkit that builds the JIT module rather than the GPU (see #34167), every affected nvcc minor was measured separately, each with a fresh TVM_FFI_CACHE_DIR. Hardware: B200 (sm_100a), single GPU, torch 2.13.0+cu130. Base: unfixed main @ 71043b9.

build	nvcc	11 fused/ctrl shapes	randomized stress	CUDA-graph replays	result
main (unfixed)	13.1.115	0 bad	0/100	0/100	clean
main (unfixed)	13.2.78	0 bad	0/300	0/200	clean
main (unfixed)	13.3.73	0 bad	0/100	0/100	clean
#34167	13.2.78	0 bad	0/300	0/200	clean

All three toolkits are cells that are 157/157 bad on H200, so this is an architecture split, not an out-of-range toolchain.

Three checks so this negative is not a false negative:

1. Right compiler — the generated build.ninja shows the intended nvcc and -gencode=arch=compute_100a,code=sm_100a.
2. Unfixed source — topk_impl.cuh:779 still carries const auto cur_out = is_primary ? problem.out : smem->tmp_out;.
3. The affected kernel really launched — profiling each case shows all 9 fused shapes dispatching to topk_small_batch_kernel<true>, with the two controls landing in topk_main_kernel<true, 2> (streaming) and topk_persistent_cluster_kernel + topk_main_kernel<true, 3> (persistent pool), as intended. Shapes covered include the reported 1 x 130800 + 7 x 2000, the 32K boundary mix, 8 x 32769, 15 x 49152, 1 x 65537, 8 x 98304, 30 x 131072, and k in {512, 1024, 2048}.

Consistent with codegen: compiling the same unfixed source with one nvcc (13.2) for both targets, topk_small_batch_kernel lowers differently — sm_90a keeps the merged pointer entirely in block-local shared stores (1 generic LD / 44 generic ST / 234 STS), while sm_100a uses generic LD/ST (31 / 65 / 276), and generic addressing is correct for both the DSMEM arm and the CTA-local arm. Suggestive rather than conclusive, but it matches an sm_90a-only mis-lowering.

test/registered/kernels/ops/attention/test_topk_v2.py also passes on B200 with #34167 applied (244 passed).

[DarkSharpness]

Thanks for the PR. This should already be fixed by #34167. Feel free to reopen if that PR does not work in your case.