[Spec][DSV4] perf: Fuse offline C128 speculative-draft state cleanup into a single kernel launch

[mattteochen]

Motivation

The offline C128 speculative-draft state cleanup (DeepSeekV4TokenToKVPool.clear_unaccepted_c128_draft_states) resets the ring slots written for rejected speculative tokens. It currently issues one Triton kernel launch per C128 layer. For DeepSeek-V4, which has many C128 layers, that is a long series of tiny launches on the decode-step hot path, and the per-launch overhead dominates the negligible amount of work each launch does. This PR fuses the per-layer cleanup into a single kernel launch across all C128 states.

Modifications

• Extract the row-store logic into a shared _clear_c128_draft_state Triton device function, reused by both the existing single-state kernel and the new fused kernel (no behavior change to the single-state path).
• Add _fused_clear_c128_draft_states_kernel, which clears the rejected rows across all C128 layer states in one launch. Because the per-layer kv_score buffers are separate allocations, it reads each layer's base pointer from a device-side pointer table; grid.z folds the layer index with the feature-dimension block index.
• Add C128DraftCleanup, a small validated helper that owns the pool-stable C128 state tensors, validates their shared dtype/shape/device/contiguity once, and caches the pointer table and launch geometry at construction. clear() is then a single launch with no per-call recomputation.
• Wire it into DeepSeekV4TokenToKVPool: build the fused cleanup once (CUDA + offline C128) in _init_paged_compress_states, and dispatch to it from clear_unaccepted_c128_draft_states, keeping the existing per-pool loop as the non-CUDA (HIP/NPU) fallback.

Accuracy Tests

sgl-eval run aime25 \
  --model deepseek-ai/DeepSeek-V4-Pro --api-key <api-key> \
  --n-repeats 16 --max-tokens 400000 \
  --temperature 1.0 --top-p 1.0 --thinking \
  --out-dir /sgl-workspace/logs \
  --base-url http://localhost:30000/v1
  
== aime25 ==
30 examples x 16 repeats  |  2606.8s  |  3200 tok/s  |  8.3M tokens

* pass@1[avg-of-16]  =  96.88% +/- 2.57% (SEM 0.64%)
  pass@16            =  100.00%
  majority@16        =  100.00%
  no_answer          =  1.04%
  stop_rate          =  100.00%
  truncated_rate     =  0.00%
  error_rate         =  0.00%

Speed Tests and Profiling

main:

PR:

Micro-benchmark (bench_c128_cleanup.py, fused vs. per-layer loop):

Benchmarked on NVIDIA B200 (test/registered/jit/benchmark/bench_c128_cleanup.py), comparing the existing per-layer loop (one kernel launch per C128 layer) against the new single fused launch. Latency in µs; speedup = loop / fused. State width = 1024 (half=512), ring_size=256.

dtype	states	batch	draft	launch	loop (µs)	fused (µs)	speedup
fp32	1	1	4	eager	7.17	7.14	1.00×
fp32	1	1	4	cuda_graph	1.09	1.52	0.72×
fp32	8	1	4	eager	35.84	7.17	5.00×
fp32	8	1	4	cuda_graph	11.39	1.66	6.87×
fp32	31	1	4	eager	163.26	7.14	22.9×
fp32	31	1	4	cuda_graph	43.80	1.74	25.2×
fp32	31	8	4	eager	161.95	7.20	22.5×
fp32	31	8	4	cuda_graph	47.85	2.23	21.4×
fp32	31	32	4	eager	161.38	11.26	14.3×
fp32	31	32	4	cuda_graph	50.96	5.16	9.9×
fp32	31	1	8	eager	162.98	7.17	22.7×
fp32	31	1	8	cuda_graph	46.51	1.84	25.2×
fp32	64	1	4	eager	512.03	7.17	71.43×
fp32	64	1	4	cuda_graph	91.44	1.78	51.32×
bf16	1	1	4	eager	7.14	7.17	1.00×
bf16	1	1	4	cuda_graph	1.47	1.52	0.97×
bf16	8	1	4	eager	35.84	7.17	5.00×
bf16	8	1	4	cuda_graph	11.47	1.62	7.09×
bf16	31	1	4	eager	165.86	7.17	23.1×
bf16	31	1	4	cuda_graph	45.69	1.74	26.2×
bf16	31	8	4	eager	166.11	7.17	23.2×
bf16	31	8	4	cuda_graph	47.64	2.15	22.2×
bf16	31	32	4	eager	164.48	11.26	14.6×
bf16	31	32	4	cuda_graph	50.75	5.12	9.9×
bf16	31	1	8	eager	164.19	7.17	22.9×
bf16	31	1	8	cuda_graph	45.49	1.82	25.0×
bf16	64	1	4	eager	516.80	7.17	72.10×
bf16	64	1	4	cuda_graph	93.41	1.78	52.42×

Takeaways. The loop cost scales linearly with the C128 layer count (launch overhead); the fused kernel stays ~flat. At the real DSV4 decode geometry (31 C128 layers, states=31) the fused path is ~23× faster eager and ~25× faster under CUDA graph, growing to ~50–70× at 64 layers. At states=1 the two are at parity (a single launch either way; the fused pointer-table indirection costs a sub-µs overhead under CUDA graph) — expected, since fusion targets the many-layer case.

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 #31469920652
Latest PR Test (Extra): ❌ Run #31469920480

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.