ReplaySSM caches SSM inputs instead of state, cutting hybrid-model decode traffic in half — 1.46–1.80× kernels, 1.85–2.10× speculative throughput

NVIDIA and Princeton's Dao AI Lab published ReplaySSM, a decoding scheme for hybrid SSM/attention models that stops loading and storing the full recurrent state every token, replacing it with a rarely-written checkpoint plus a small buffer of recent inputs. Across four Nemotron 3 and Qwen3.5 models it reports 1.46–1.80× faster SSM decode kernels at batch size 256, and 1.85–2.10× higher speculative-decoding throughput on two 120B-class MoE models at batch 512. The NVIDIA blog says the kernels are already integrated into TensorRT-LLM, vLLM, and FlashInfer.

What shipped

Hybrid models interleave a few attention layers with many state-space layers — Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3, and GLM-5.3-Flash all do this, using Mamba-2, Gated DeltaNet (GDN), or Kimi Delta Attention (KDA) as the SSM layer. The SSM's fixed-size state was meant to make decoding cheaper than an ever-growing KV cache. At serving batch sizes it does the opposite: on Nemotron-3-Super-NVFP4 at batch 256, the SSM layers are the largest component of the decode step up to 100K tokens of context.

The reasons are arithmetic. The model has 40 SSM layers against 8 attention layers, and one layer's FP32 state is 4 MiB per sequence — the size of an attention layer's FP8 KV cache for 8K tokens. Every decoded token loads the full state from HBM, applies a rank-1 update, and stores it back, so the kernel is bandwidth-bound rather than compute-bound. The state does two jobs at once — producing this token's output and becoming the next token's state — and the standard implementation fuses both into every memory pass.

ReplaySSM unrolls the recurrence over a window. For the Mamba-2 update Sₜ = αₜSₜ₋₁ + vₜkₜᵀ, the state h tokens back can be written as γ_h·S_ckpt plus a sum over the buffered (α, v, k) records. On non-flush steps the kernel appends the current inputs to the buffer and computes the output from the checkpoint plus buffer — no state store. When the 16-record buffer fills (L = 16 beat every other window size tested across all four models), a flush step folds the records into the checkpoint, writes it to HBM once, and clears the buffer. Conventional decoding reads the state L times and writes it L times per window; ReplaySSM writes once. Outputs match conventional decoding up to floating-point differences.

This also unblocks speculative decoding, which barely helps SSMs at serving batch sizes today. Because an SSM state update has no easy inverse, serving engines like vLLM work around rollback by storing a state snapshot after every draft token — one read and T writes per step on an already memory-bound path, with the snapshots capping how many requests fit in GPU memory. With ReplaySSM, drafts are just appended to the input buffer, verification runs, and the accepted prefix is committed by advancing a pointer; records past the pointer are ignored and overwritten, exactly like a KV cache.

There is a second, subtler benefit. A verification step that produces only the output — never the next state — can be rewritten so no intermediate state is formed at all: O = S_ckpt·Q·D + V(W ⊙ (KᵀQ)), with D holding cumulative decay of draft positions and W a causally masked decay matrix. The KᵀQ product is shared across all value heads in a group, the operands are the small cached keys/values and draft queries, and everything is bulk matrix work except a small T×T triangular solve on the GDN delta-rule path — the same chunkwise computation prefill kernels already run, with the buffer plus drafts as the chunk.

The evaluation covers four models: Nemotron-3-Nano-4B (Mamba-2, BF16, one H100) and Nemotron-3-Super-120B-A12B (Mamba-2, NVFP4, one Blackwell Ultra B300), plus Qwen3.5-4B (GDN, BF16, one H100) and Qwen3.5-122B-A10B (GDN, NVFP4, one B300). At batch size 256 the ReplaySSM SSM kernels run 1.46–1.80× faster than vLLM's tuned kernels, translating to a 1.20–1.47× faster whole decode step; at batch 1 there is no speedup, because other latencies dominate. For speculative decoding with each model's MTP head, T = 4, on GSM8K prompts at batch 512, throughput is 1.85–2.10× that of vLLM's snapshot-based speculative path per the blog's summary; an independent overview of the paper pins the Nemotron model's point at 1.85×.

A third result addresses quantization. A quantized SSM state is its own next input, so its rounding error re-enters the recurrence every step — unlike a KV cache, which is read but never written back into itself. Qwen3.5-122B-A10B-NVFP4, for instance, already ships NVFP4 weights with an FP8 KV cache and an FP32 SSM state. ReplaySSM decouples the two uses: a high-precision checkpoint drives the recurrence while a low-precision copy serves the reads, so the read copy's error never propagates into later updates. The blog reports three of the four evaluated models staying within about 2 percentage points of FP32 accuracy on GPQA-Diamond.

The timeline is as notable as the technique. The core idea appeared as a Dao AI Lab blog post in June 2026; the paper landed on arXiv on October 9; the NVIDIA blog post of October 12 claims integration into TensorRT-LLM, vLLM, and FlashInfer plus production use in GLM-5.3-Flash and Kimi-K3. The integration trail is independently visible, though partial: the vLLM RFC #47572 and draft PR describe the buffer-length flag and ring semantics (with only the Mamba-2 standard-decode kernel merged upstream so far, per the repo's adoption table); a FlashInfer PR by a paper author implements the same ring-buffer cache contract; and SGLang is porting it too. The Apache-2.0 repo is a vLLM fork with Mamba-2 and GDN kernels for standard and speculative decoding.

Why it matters

ReplaySSM attacks the bottleneck that dominates hybrid-model serving, and it is being adopted, not merely proposed. If you serve Nemotron 3, Qwen3.5-3.8, GLM-5.3-Flash, or Kimi-K3 at meaningful batch sizes, the SSM state traffic and the snapshot-based speculative path are costing you real throughput today; the fix is a cache-layout change plus a reformulated equation, not a new model.

The speculative-decoding result is the most consequential piece. Speculation underperformed on hybrids not because of the drafter but because rollback on a non-invertible recurrence forced snapshotting and serial verification. ReplaySSM's checkpoint-plus-buffer representation turns rollback into a pointer advance and the serial per-token recurrence into chunkwise GEMMs. That is the same structural trick prefill already exploits, and it is why the verification path gets fast.

The deeper lesson is the reframing itself: cache the inputs, not the state. For any recurrence with a decay structure, the state is reconstructible from a checkpoint plus recent inputs, which converts a read-modify-write memory pattern into an append log with cheap replay. It is the same reasoning behind replay buffers in other domains, and it generalizes cleanly to GDN and the delta-rule update in KDA.

Caveats

All performance numbers above are reported by the authors and NVIDIA, not independently benchmarked. The open-source artifacts exist, but the paper's own tables are the only source for most figures.

The win is strictly conditional on being memory-bound. ReplaySSM provides no speedup at batch 1, and vLLM's own RFC table shows whole-model speedups of 0.98–0.99× at batch 1. The SGLang port is more explicit: the rewrite buys nothing when the verify kernel runs below roughly 30% of peak HBM bandwidth and only fully pays off around 70%+. ReplaySSM is also not bit-exact — reconstruction is mathematically lossless, but tensor-core reassociation differs from the sequential baseline at about 1e-4 (TF32) and 1e-3 (bf16), and forcing IEEE exactness makes the kernel 10–20× slower.

The accuracy claim is thinner than the throughput claim. The blog's own auto-generated summary cites a single benchmark, GPQA-Diamond, with three of four models within about 2 points of FP32, and the fourth model's delta was not visible in any text we could access (the paper abstract also reports output KL divergence). Quantization error is reduced, not eliminated.

Coverage is narrower than the buzz suggests. Only Mamba-2 and GDN were evaluated — KDA was not — and all runs were single-GPU with no tensor parallelism. Only L = 16 results appear in the main evaluation (the blog says it beat the other window sizes tried on all four models), and per-request latency was not measured. The production claims for GLM-5.3-Flash and Kimi-K3 rest on the authors' README and a Z.ai blog snippet; the cited Kimi technical-report section is a PDF we could not independently confirm.

Sources

• ReplaySSM: Cache SSM Inputs, Not State — arXiv 2610.13778
• ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding — NVIDIA Technical Blog
• AlphaXiv overview of the ReplaySSM paper
• Dao AI Lab blog, June 2026 (earlier ReplaySSM write-up)
• ReplaySSM code (Apache-2.0)
• vLLM RFC #47572
• FlashInfer PR #3975
• SGLang RFC #28511