[yuan-luo]

Motivation

The fused KDA chain-verify kernel (fused_kda_conv_gating_verify, gated by SGLANG_OPT_FUSED_KDA_VERIFY) and ReplaySSM spec-verify (--enable-linear-replayssm + spec ring) were mutually exclusive: _can_run_fused_chain_verify returned False whenever replayssm_rawv was set, so enabling ReplaySSM silently dropped the verify hot path back to the unfused reference pair (transpose-copy + causal_conv1d_update + transpose-copy + recurrence). We could have the ~1.3x fused verify kernel or the ReplaySSM memory win (no intermediate_ssm snapshot buffer, the dominant spec scratch), but not both.

On Blackwell the CuTe DSL dspark MTP kernel already composes ring-write with a fused verify, but that path is SM100-only with a fixed Kimi-K3 shape contract. This PR brings the same composition to the Triton fused chain-verify, which is the only fused verify available on Hopper and on non-dspark chain speculation.

Beyond unblocking the combination, ring mode removes the fused kernel's dominant store: the per-step [K, V] fp32 state snapshot (64 KB per head-step) is replaced by the ring records (raw v + raw k + gate + beta, ~1 KB per head-step), on a kernel that is latency-bound.

Technical principle

Per draft step $t$ the fused kernel computes the conv activations, the gates, and one delta-rule update:

$$\hat{x}_t = \mathrm{SiLU}(w \ast x_{t-3:t}) \rightarrow (\hat{q}_t,\ \hat{k}_t,\ \hat{v}_t), \qquad g_t = \ell\cdot\sigma\big(e^{A}(a_t + b_{dt})\big), \qquad \beta_t = \sigma(b_t)$$

$$\tilde{h} = \mathrm{diag}(e^{g_t}),h_{t-1}, \qquad u_t = \beta_t,(\hat{v}_t - \tilde{h}^{\top} k_t), \qquad h_t = \tilde{h} + k_t,u_t^{\top}, \qquad o_t = h_t^{\top} q_t$$

In ring mode the kernel stores $(\hat{k}_t,\ \hat{v}_t,\ g_t,\ \beta_t)$ into the per-slot ring at exactly the point where $\hat{k}_t$ is still pre-l2norm and $\hat{v}_t$ still pre-delta ($g_t$, $\beta_t$ just formed), instead of snapshotting $h_t$. On commit, the accepted prefix $a \le T$ is folded back into the persistent checkpoint in place:

$$h_0 \leftarrow \mathrm{Fold}\big(h_0;\ {(\hat{k}_s, \hat{v}_s, g_s, \beta_s)}_{s=1}^{a}\big)$$

The fold kernel is a bit-clone of the recurrence (same decay-delta-rank-1 op order, num_warps=1), so replaying the ring reproduces the committed state bit-identically — which is why the stored operands must be the pre-mutation values. Write traffic per head-step drops from $4KV = 65536$ B (snapshot) to $2V + 2K + 4K + 4 \approx 1028$ B (ring), about $64\times$, and the pool-level intermediate_ssm buffer ($\mathrm{slots} \times (\gamma{+}1) \times HV \times K \times V$ fp32) disappears entirely.

Modifications

python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py — add CACHE_RING to the kernel: each draft step stashes its pre-l2norm k, pre-delta v, in-kernel gate b_g, and b_beta into the per-slot ReplaySSM rings, at the exact point the unfused CACHE_RING block in fused_sigmoid_gating_recurrent.py does (after beta is formed, before l2norm mutates b_k and the delta rule mutates b_v), so the commit-time exact fold replays identical bytes. The ring slot reuses the already-loaded h0_idx; rawk dedups via the existing is_qk_owner (one writer per k-head); g/beta write once per v-head at i_v == 0; the t < MAX_CACHE_LEN guard mirrors the reference kernel's absorb-overflow protection. With CACHE_RING=False (the default) everything folds away — the existing paths compile to the same code as before. The wrapper takes the same cache_ring / replayssm_* parameter names as the unfused kernel so the backend passes one kwargs dict through both, and validates ring shapes/dtypes/inner strides.

python/sglang/srt/layers/attention/linear/kda_backend.py — hoist the ring_kwargs construction ahead of the fused branch and share it between the fused chain-verify call and the unfused dispatcher call. _can_run_fused_chain_verify no longer bails on replayssm_rawv; under ReplaySSM it skips the intermediate_state_cache checks (the buffer is None there by design) and instead validates the four per-layer ring views via the new _replayssm_ring_ok helper (ndim, shape, dtype, packed inner strides, L >= draft_token_num, device). Any mismatch falls back to the unfused path, which handles more layouts.

test/registered/kernels/test_fused_kda_conv_recurrent_verify.py — new test_replayssm_ring_matches_unfused parameterized over the existing 8 cases plus 3 GQA (HV != H) cases. Ring mode runs both kernels with intermediate_states_buffer=None (exactly the production ReplaySSM configuration) and asserts the output, conv state, conv windows, and all four rings are bitwise equal between fused and unfused. The rings share a garbage-filled template cloned per side, so a full-tensor compare also catches out-of-slot scribbles and negative-slot writes. The GQA cases pin the per-k-head rawk vs per-v-head g/beta writer split — the one mistake that would silently corrupt another head's ring. Ring values are elementwise (conv FMA chain, gate, sigmoid) and sit upstream of every tl.sum, so bitwise equality holds at the production num_warps=4.

benchmark/kernels/bench_kda_verify_sweep.py (new) — in-graph sweep over the four dispatchable variants ({unfused, fused} x {snapshots, ring}), CUDA-graph timed (these are ~10us kernels; bare launches would measure launch overhead). Includes a built-in fused-vs-unfused ring bitwise guard before timing, a --sweep-bv mode to re-tune KDA_VERIFY_BLOCK_V (BV=4 was tuned with snapshot writes on, so ring mode may move the optimum), --lower-bound for the safe-gate (ling-v3) vs standard-gate shapes, and --hv-heads for GQA. Imports only sglang.kernels.* so it runs on boxes where the sglang.srt import chain is unavailable.

Correctness

The acceptance criterion is deliberately stronger than "matches the current fused path": under ReplaySSM the committed state is produced by the fold kernel replaying the ring, so the invariant that matters is that fused+ring and unfused+ring fill bit-identical rings (then the unchanged fold produces bit-identical committed states). The new test asserts exactly that, plus bitwise-equal outputs, conv states, and conv windows.

Performance

fused+ring vs unfused+ring is the end-to-end win this PR delivers under ReplaySSM (previously forced unfused); fused+ring vs fused isolates the snapshot-write removal. Measured on H20-3e (sm90), ling-v3 TP4 shape (H = HV = 8, K = V = 128, T = 4, safe gate), us/layer, CUDA-graph timed, repeat-run variance ~1-2%:

python3 benchmark/kernels/bench_kda_verify_sweep.py --lower-bound -5.0   # ling-v3 TP4 shape
python3 benchmark/kernels/bench_kda_verify_sweep.py --lower-bound -5.0 --sweep-bv --batch-sizes 1 4 16 64

H20-3e (sm90):

B	unfused	unfused+ring	fused (BV=4)	fused+ring (BV=4)	fused+ring vs unfused+ring	fused+ring (best BV)
1	11.47	9.56	8.17	7.52	1.27x	6.56 (BV=16)
2	12.33	9.73	12.73	12.22	0.80x	—
4	14.43	10.18	19.53	19.06	0.53x	7.72 (BV=32)
16	23.82	15.06	71.62	67.36	0.22x	17.31 (BV=32)
64	79.38	51.49	280.06	289.93	0.18x	65.68 (BV=32)

B300 (sm103):

B	unfused	unfused+ring	fused (BV=4)	fused+ring (BV=4)	fused+ring vs unfused+ring	fused+ring (best BV)
1	14.33	10.66	6.14	7.14	1.49x	6.76 (BV=16)
2	14.13	10.65	7.58	8.58	1.24x	—
4	15.17	11.06	12.28	13.93	0.79x	7.37 (BV=16)
16	17.83	12.09	36.25	41.23	0.29x	10.45 (BV=32)
64	53.11	28.05	172.81	153.47	0.18x	34.19 (BV=32)

Takeaways on this card: with BV re-tuned, fused+ring beats unfused+ring by 1.44x at B=1 and 1.30x at B=4, and beats fused (snapshots) by up to 2.0x at B=16 / 1.75x at B=64 — the snapshot store is exactly what dominates at larger batch, and ring mode removes it. The crossover to the unfused pair sits around B~8-16 even at the best BV, confirming the batch-aware dispatch follow-up; and the optimal BV shifts larger once the snapshot write is gone (BV=16-32 vs the current BV=4), confirming the BV re-tune follow-up. B200 (production target) numbers pending — the fixed BV=4 constant was tuned there and the crossover point will differ.

Follow-ups

• Re-tune KDA_VERIFY_BLOCK_V on B200 with --sweep-bv once ring mode lands; the current constant was measured with the snapshot store dominating the write traffic.
• The tree path (retrieve_*) and ragged verify layouts keep the unfused kernels, unchanged.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a reviewer when merging.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

---

CI States

Latest PR Test (Base): ❌ Run #34865490140
Latest PR Test (Extra): ❌ Run #34865489076
Latest PR Test (AMD ROCm 10): ❌ Run #34865489373

[BBuf]

This enables fused+ring for all eligible batches, but the runtime still uses BV=4. In the H20 results above, B=4/16/64 go from 10.18/14.99/51.46 us to 19.06/67.38/290.05 us. Can we keep the unfused path for the batch/architecture combinations that regress? The best-BV sweep isn't wired into dispatch, and B=16/64 still regress even with that tuning.

[yuan-luo]

Agree. The best-BV sweep isn't used by the runtime. I've restricted the new fused+ring path to SM90 with B=1.

[yuan-luo]

I re-swept the hopper and blackwell, revised the gate accordingly.

[BBuf]

Could we add a backend test that runs verify -> commit -> verify with both partial and full acceptance? This test compares outputs and ring contents, but doesn't exercise commit/conv rollback or check the next round's output. Please compare the committed SSM state, conv state, and next verify output against the unfused path.

[yuan-luo]

Added backend-level verify -> commit -> verify tests covering partial and full acceptance, multiple layers, and GQA.