[yuan-luo]

Motivation

To address #28511.

The GDN (Gated DeltaNet) decode kernel reads the full recurrent state S [HV, V, K] from HBM and writes it back every decode step. On the H20‑3e this packed decode is 70–79% memory-bandwidth-bound at batch ≥ 64 — the per-step state read+write traffic dominates.

ReplaySSM (Dao AI Lab, https://dao-lab.ai/blog/2026/replayssm/) removes the per-step state write: keep a small per-slot ring buffer of the last L steps' (d, k, g), read the output via the history route, and flush the full state only every L steps. Per-step state traffic drops from read+write ( ~8dn ) to read-only ( ~4dn ) → roughly halved.

This PR ports the ReplaySSM decode path into SGLang behind a flag, fully integrated with CUDA graph and the radix prefix cache. It covers both linear-attention gate granularities with one gate-type-parametrized kernel: GDN (per-head scalar gate, radix + CUDA graph) and KDA (Kimi Delta Attention, per-K-channel gate).

Modifications

Adds --enable-linear-replayssm (with --linear-replayssm-cache-len L, default 16). When enabled, single-token linear-attention decode runs a buffered output-only kernel (fused_recurrent_linear_replayssm_decode, IS_KDA constexpr) instead of the packed decode — fused_recurrent_gated_delta_rule_packed_decode for GDN or fused_recurrent_kda_packed_decode for KDA. The GDN math (per-head scalar gate $\alpha$):

$$ \begin{aligned} \alpha &= \exp(g) && \text{per-head decay} \\ d_{\mathrm{cur}} &= \beta (v - \alpha (Sk)) && \text{GDN corrected delta-rule vector} \\ o &= \alpha (Sq) + d_{\mathrm{cur}} (k^{\top} q) && \text{output-only, } k^{\top} q \text{ factored first} \\ S_{\mathrm{new}} &= \alpha S + d_{\mathrm{cur}} k^{\top} && \text{persisted only on a flush (every } L \text{ steps)} \end{aligned} $$

where $S \in \mathbb{R}^{V \times K}$; $Sk$ and $Sq$ are matrix-vector contractions over the key dimension; and $q$ is already L2-normed (optional) and scaled.

The $[V, K]$ state is reconstructed from the checkpoint $S_0$ + ring in K-tiles and read with $q$ / $k$ immediately, so it is never materialized to HBM on a non-flush step. The reconstruction tl.dot runs on tensor cores (TF32 for fp32, bf16 tensor cores for bf16) — required so the $L$-deep reconstruction stays hidden under the reduced memory traffic.

Design

This is a memory-traffic optimization for a memory-bound decode kernel, so several decisions trade extra (cheap, tensor-core) compute for less HBM traffic, and several more are about doing that safely under CUDA graph + the radix prefix cache. The non-obvious ones:

1. Output-only reconstruction, never materialize the state

On a non-flush step the buffered state before the current token is

$$S = A S_0 + \sum_{j} w_j d_j k_j^{\top}$$

with total decay $A = \prod_i \alpha_i$ and per-position replay decay $w_j = \prod_{i>j} \alpha_i$, both built from the cached per-step gates $g_j$ via a cumsum (no per-step state needed). $S$ is reconstructed in K-tiles and each tile is immediately contracted with $q$ and $k$ into the scalars $S q$ / $S k$, so the full $[V,K]$ tile never lands in HBM. At $L=1$ the ring is always empty and the kernel reduces algebraically to fused_recurrent_gated_delta_rule_packed_decode (used as the bit-exact correctness oracle).

2. Tensor cores are mandatory, not optional (the main perf trap)

The reconstruction is $O(L\cdot K\cdot V)$ compute — $L\times$ the baseline's rank-1 update. That only stays hidden under the (halved) memory traffic if the tl.dot runs on tensor cores. We hit this directly: forcing an IEEE fp32 dot (e.g. to pass an over-tight unit-test tolerance) drops tensor cores and makes the kernel ~10–20× slower — turning a 1.5× win into a 15× regression. The fix is to keep the dot on tensor cores (TF32 for fp32 inputs, bf16 tensor cores for bf16) and set tensor-core-realistic test tolerances instead. TF32 (~4e-4) / bf16 (~1e-3) reconstruction error is benign end-to-end (GSM8K parity below); GDN's per-step decay $\alpha<1$ attenuates old reconstruction error so it does not accumulate across flushes over a long decode.

3. write_pos is a per-slot decode-position counter, not derived from seq_len

The ring cursor is tracked per physical mamba slot (in MambaPool), not computed from seq_len. Deriving it from seq_len % L would be wrong: seq_len includes the prefill length, so a slot's first decode step would read a non-empty ring of zeros. Instead:

• Reset to 0 at slot allocation — at that moment prefill has written the post-prefill state into temporal[slot], which is exactly the correct checkpoint $S_0$ with an empty ring.
• Advanced once per forward (in the attention metadata build), not per layer — write_pos is the decode position, shared by all GDN layers in a step. Advance rule: next = 0 if (write_pos == L-1 or force_flush) else write_pos + 1.

The prefill path is untouched; it writes the full state to temporal[slot], which becomes the decode checkpoint.

4. CUDA-graph safety

The decode forward is captured once and replayed, so write_pos cannot be a fresh tensor per step. It is a static per-bs buffer captured by the graph and refreshed in-place each replay from the persistent per-slot counter, mirroring how mamba_cache_indices is handled (_replay_metadata, outside the captured region). The per-slot advance also runs there, guarded by not in_capture so warmup/capture never mutates the persistent counter. The flush decision is a device-side branch (is_flush = write_pos == L-1), so a single captured graph handles every step regardless of where in the ring cycle each row is.

5. Radix coordination — flush exactly when the cache snapshots

The radix prefix cache snapshots the working temporal[slot] at seq_lens % mamba_track_interval == 0 boundaries for prefix reuse. For ReplaySSM temporal[slot] is the checkpoint, which is stale between flushes. So the decode kernel must flush on exactly those steps. force_flush is computed from the identical condition and the same post-increment seq_lens_cpu the radix cache uses (schedule_batch.py), so there is no off-by-one — the ring is folded into temporal[slot] on the same step the snapshot reads it. The kernel ORs force_flush into is_flush while cache_valid still uses the true write_pos (only committed ring entries are read, even when flushing mid-buffer).

On a prefix hit, MambaPool.copy_from clones the (now always fully-flushed) snapshot into the request's private slot; we reset that slot's write_pos to 0 so the destination ring starts empty.

6. flag-OFF is byte-identical

When the flag is off no ring tensors are allocated, the metadata fields stay None, and the kernel dispatch falls through to the existing packed decode — zero change to the current path.

7. One kernel for GDN (scalar gate) and KDA (per-K gate)

GDN and KDA are the same delta-rule decode modulo gate granularity: GDN decays the whole state by a per-head scalar $\alpha=\exp(g)$; KDA decays it per key channel, $S' = S,\mathrm{Diag}(a) + d,k^{\top}$ with $a=\exp(g)\in\mathbb{R}^{K}$. A single IS_KDA constexpr selects the gate path (mirroring the baseline fused_recurrent_*_packed_decode split), and the GDN path is bit-for-bit the original kernel (no regression — see Validation).

The non-obvious part is that ReplaySSM accumulates the gate into the reconstruction, so per-K gating touches more than the one baseline line:

• ring g_cache becomes per-K: […, L] (GDN) → […, L, K] (KDA);
• total / replay decay become per-channel cumsum over the buffer ($A[c]=\exp(\sum_j g_j[c])$, $w_j[c]=\exp(\sum_{i>j} g_i[c])$);
• the replay decay folds onto the cached keys (per-K) instead of onto $d$ (scalar), and the current-token $\alpha$ folds into $q$/$k$ instead of a scalar post-multiply.

Everything else — the rank-1 reconstruction skeleton, output-only readout, K-tiling, ring / write_pos / flush, and the entire runtime integration (memory pool ring, CUDA graph, COW) — is gate-agnostic and shared. MambaPool sizes g_cache from cache_params.is_kda (Mamba2 → scalar, KimiLinear → per-K). KDA ships without radix coordination for now (force_flush left off for KDA; KimiLinear has the radix cache disabled upstream anyway), so the KDA ring flushes only at the natural write_pos == L-1 wrap.

Notes investigated

• finish-donate: cache_finished_req donates the working temporal at the request's full (non-L-aligned) length. We checked whether this exposes a stale ReplaySSM checkpoint on exact full-prefix reuse; in practice it is TF32-level only and GSM8K parity (below) shows no correctness impact.
• Roofline: on H20-3e the packed GDN decode is 70–79% HBM-bandwidth-bound at batch ≥ 64; ReplaySSM cuts state traffic to ~0.53×, for a ~1.7× kernel ceiling at high batch.

Usage

python -m sglang.launch_server --model-path <Qwen3.5-GDN-model> \
  --linear-attn-decode-backend triton \
  --enable-linear-replayssm --linear-replayssm-cache-len 16
# works with the radix cache and CUDA graph both ON (defaults)

Implementation

Landed as reviewable slices (one commit each):

Slice	What
kernel	fused_recurrent_linear_replayssm_decode Triton kernel (IS_KDA constexpr) + wrapper + unit test (GDN + KDA) + microbench
1a	MambaPool ring (d/k/g + per-slot write_pos); GDN backend decode dispatch behind the flag (eager)
1b	CUDA-graph-safe write_pos (static per-bs buffer refreshed/advanced outside the captured region)
2a	kernel force_flush input (`is_flush
2b	radix integration: force_flush = the radix track mask (seq_lens_cpu % mamba_track_interval == 0, identical to schedule_batch.py), so the ring flushes on exactly the steps the radix cache snapshots temporal[slot]; write_pos reset on COW (MambaPool.copy_from)
KDA	generalize the kernel to a per-K gate (IS_KDA constexpr; GDN path byte-identical); MambaPool sizes per-K g_cache from cache_params.is_kda; KDA backend + kda_triton decode dispatch threading; flag generalized --enable-gdn-replayssm → --enable-linear-replayssm

write_pos is a per-SLOT decode-position counter (reset to 0 at slot alloc and at COW), advanced once per forward; flag-OFF is byte-identical to the current path.

Validation (H20-3e)

GDN correctness — Qwen3.5-35B-A3B (GDN-hybrid)

• Kernel unit test (test_linear_replayssm_decode.py): GDN fp32 bit-exact at L=1, tensor-core-realistic tolerances for L>1, forced-flush mid-buffer — all pass.
• GSM8K (200 q): flag-on 0.880 vs baseline 0.860 — within noise, no regression, re-confirmed after the GDN+KDA kernel unification. (The earlier GDN-only integration slices validated at 0.900 vs 0.850 with radix + CUDA graph ON.)
• Greedy decode is byte-identical to baseline until the first near-tie token, then differs only via benign TF32 word-flips (both continuations correct) — not corruption.

KDA correctness — per-K gate

• Kernel unit test: KDA passes against fused_recurrent_kda_packed_decode at the same tensor-core tolerances (L=1 and L>1), forced-flush included — confirms the per-K reconstruction math.
• e2e plumbing verified live: server starts with the flag, MambaPool allocates per-K g_cache [.,L,K] (g≈k in size, vs GDN's scalar g≈k/128 — confirms is_kda routing), decode runs end-to-end without crash.

Performance

• Kernel microbench: 1.2–1.5× vs packed decode at batch ≥ 16 (L=8); slightly slower at batch 1 (overhead-bound, not an operating point).
• E2E (128 concurrency, 256/512 in/out): ~2.3% TPOT improvement. The end-to-end gain is model-dependent — small here because Qwen3.5-35B-A3B is MoE-heavy and the GDN linear layers are a diluted fraction of decode; a GDN-dense model would show more.

Scope / follow-ups

• GDN: fully integrated (radix + CUDA graph, GSM8K parity). KDA: kernel + decode dispatch wired and kernel-validated; radix coordination for KDA is a follow-up (force_flush currently off for KDA — it flushes only at the natural L-1 wrap).
• KV-pool budgeting: mamba_cache_per_req does not yet include the ReplaySSM ring bytes, so the KV pool can be over-sized when the rings are large (KDA / big L) → OOM at default mem_fraction_static. Workaround: lower --mem-fraction-static. Proper fix (count ring bytes in mamba_cache_per_req) is a follow-up.
• Requires --linear-attn-decode-backend triton and mamba_scheduler_strategy=no_buffer (the default). The extra_buffer/ping-pong donation path is not yet covered and is guarded out.
• Kernel is ported from vLLM fused_recurrent_replayssm.py (ReplaySSM); the vllm::persistent-style math is preserved for diff-friendly upstream sync.

Credit

Algorithm: ReplaySSM, Dao AI Lab (2026).

---

CI States

Latest PR Test (Base): ❌ Run #28182074297
Latest PR Test (Extra): ✅ Run #28182074008

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yuan-luo]

Working on KDA/GDN unified solution.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yuan-luo]

Oh, another quick read, it seems --enable-linear-replayssm is guarded to require the Triton decode backend and mamba_scheduler_strategy=no_buffer, which makes sense. Do we also need to guard or wire this through for disaggregated decode mode? For example: the disaggregated decode-side pool (HybridMambaDecodeReqToTokenPool) does not seem to receive these args

@kaixih Thanks for the comment. PD disaggregated decode-side pool HybridMambaDecodeReqToTokenPool doesn't receive these args. I agree that we need this guard. Added a guard in this PR, and I'll follow up a more decent PR to update HybridMambaDecodeReqToTokenPool accordingly.

[yuan-luo]

Implementing the extra_buffer mode would be very helpful, as it consumes a lot of memory!

@Swipe4057 Agree, I'll do it in a follow-up PR.

[YAMY1234]

Similar method for the flashinfer support: #27725 cc @Jane1019-Libra @nvpohanh @ameynaik-hub

[kaixih]

Similar method for the flashinfer support: #27725 cc @Jane1019-Libra @nvpohanh @ameynaik-hub

I feel your mentioned one is closer to #28695. Since the flashinfer is currently updating to follow this replaySSM style, should we merge the efforts? How do you think @YAMY1234 @yuan-luo ?

[yuan-luo]

Similar method for the flashinfer support: #27725 cc @Jane1019-Libra @nvpohanh @ameynaik-hub

I feel your mentioned one is closer to #28695. Since the flashinfer is currently updating to follow this replaySSM style, should we merge the efforts? How do you think @YAMY1234 @yuan-luo ?

@kaixih Sure, we can discuss it in #28695.

[yyq0210]

Thanks for this work. The unified kernel works correctly for both gates, but I'm seeing GDN and KDA behave very differently and wanted to check if this is expected.

Setup: amortized decode microbench — write_pos cycles 0→…→L-1→0 so timed steps include ring reconstruction and the 1-in-L flush (not just the cheapest non-flush step). Each variant timed vs its own packed baseline with an identical launch loop. HV=32, H=16, K=V=128, bf16.

GDN (scalar gate, g_cache=[slots,HV,L]): win at batch ≥ 64, ~1.2–1.37× amortized.

KDA (per-K gate, g_cache=[slots,HV,L,K]): slower than packed everywhere:

B	L=8 best / amort	L=16 best / amort
16	0.85× / 0.74×	0.85× / 0.78×
64	0.88× / 0.77×	0.88× / 0.82×
256	0.93× / 0.82×	0.93× / 0.87×

The KDA g_cache is K× larger and the L-deep reconstruction redoes per-K decay folding each step, so the ring's saved state traffic doesn't offset the added compute vs an already-efficient packed baseline. On a KDA-dense model (Kimi-Linear-48B, 20 KDA layers) I also see a consistent ~3.6% TPOT regression with --enable-linear-replayssm.

Questions:

1. Were the reported speedups measured on GDN only, or was KDA benchmarked too? If so, which config?
2. Is ReplaySSM intended for KDA, or is it primarily a GDN optimization? Would gating it to GDN make sense?

[yuan-luo]

Thanks for this work. The unified kernel works correctly for both gates, but I'm seeing GDN and KDA behave very differently and wanted to check if this is expected.

Setup: amortized decode microbench — write_pos cycles 0→…→L-1→0 so timed steps include ring reconstruction and the 1-in-L flush (not just the cheapest non-flush step). Each variant timed vs its own packed baseline with an identical launch loop. HV=32, H=16, K=V=128, bf16.

GDN (scalar gate, g_cache=[slots,HV,L]): win at batch ≥ 64, ~1.2–1.37× amortized.

KDA (per-K gate, g_cache=[slots,HV,L,K]): slower than packed everywhere:

B L=8 best / amort L=16 best / amort
16 0.85× / 0.74× 0.85× / 0.78×
64 0.88× / 0.77× 0.88× / 0.82×
256 0.93× / 0.82× 0.93× / 0.87×
The KDA g_cache is K× larger and the L-deep reconstruction redoes per-K decay folding each step, so the ring's saved state traffic doesn't offset the added compute vs an already-efficient packed baseline. On a KDA-dense model (Kimi-Linear-48B, 20 KDA layers) I also see a consistent ~3.6% TPOT regression with --enable-linear-replayssm.

Questions:

1. Were the reported speedups measured on GDN only, or was KDA benchmarked too? If so, which config?
2. Is ReplaySSM intended for KDA, or is it primarily a GDN optimization? Would gating it to GDN make sense?

@yyq0210 Thanks for the insight comment. Your numbers match my verification, KDA loses performance, GDN wins at batch >= 64. KDA's ReplaySSM win is on the spec-decode memory side (#27658), not decode. Rather than gate it off, I documented it as following: the flag help + a dispatch comment now flag that KDA decode regresses and isn't recommended. Keeping the kernel wired. Since KDA hit is compute-bound, we'd better revisit in Blackwell.