ReplaySSM: Cache SSM Inputs, Not State

What it claims

ReplaySSM targets a bottleneck most engineers would not expect in hybrid models. For Nemotron-3-Super-120B-A12B-NVFP4 at batch size 256 on one B300, SSM state updates, not attention, are the largest per-step latency component up to about 100K tokens of decode context. Two facts drive this:

• The model has five times as many SSM layers as attention layers.
• Each SSM layer's FP32 state is 4 MiB per sequence, the same size as an attention layer's FP8 KV cache at 8K tokens.

A conventional step both loads and stores that state, while attention only loads its KV cache. So the SSM layers move more bytes until roughly 80K tokens. The measured crossover sits nearer 100K, which the authors attribute to attention kernels being better optimized.

How it works

• Standard decoding. Each step appends the token's inputs (α, v, k) to a buffer. It rebuilds the state on chip as the decayed checkpoint plus a sum over the buffered inputs, then computes the output. The checkpoint is written to HBM only when the buffer reaches the window size L. Reads are unchanged: there are still L full-state reads per window. The saving is on writes, which drop from L to 1, and that is why traffic only "roughly halves".
• Window size. All experiments use L = 16. A larger L spreads each checkpoint write over more steps but makes every step load and replay more buffered records.
• Gated DeltaNet (GDN). The update is rewritten into Mamba-2's outer-product form using a correction term u_t = β_t(v_t − α_t S_{t−1} k_t). Each buffer record stores u_j in place of v_j. Computing u_t still requires the previous state.
• Speculative decoding. All T draft records are appended, and outputs are computed directly from the checkpoint plus the buffer, without materializing intermediate states. Accepting a draft prefix advances a pointer. Rejected records are simply overwritten later. Flushes fold in only already-committed records, so they can run before acceptance is known. The buffer is circular with capacity L+T.
• Two precisions. Non-flush steps read a low-precision checkpoint. Flush steps read a high-precision one and write both, so low-precision rounding error never feeds the recurrence.

flowchart TD
  A[Append T draft records] --> B{h+T > L?}
  B -- no --> C[Load LP checkpoint]
  B -- yes --> D[Load HP checkpoint]
  C --> E[Output-only verify]
  D --> E
  D --> F[Fold committed records, store HP + LP]
  E --> G[Advance pointer by accepted length]

Evidence

All four models run at batch size 256, L = 16, with FP32 state, against a tuned vLLM baseline. Kernel figures include flush costs.

The mapping of bars to models is read from the extracted figure text. It is consistent with the abstract's 1.20–1.47× end-to-end range.

The paper also says prior rollback methods (Wang et al. 2024; Wu et al. 2024, 2026) are at most a few percent faster than vLLM's per-draft snapshots at serving batch sizes. This is why it compares only against vLLM's speculative decoding.

How to use it

Code is at https://github.com/Johnny-Liou/ReplaySSM. It is implemented inside vLLM for Mamba-2 and GDN layers. The fetched text gives no install commands, licence or supported vLLM version.

Limits

• Outputs match conventional decoding only up to floating-point differences.
• Kimi Delta Attention is named as compatible in principle but not implemented.
• The fetched text cuts off in Section 5. The setup behind the ~1.8× speculative-decoding throughput (drafter, draft length, acceptance rates) and the KL-divergence and accuracy numbers for two-precision state are not visible here.
• The text shows results only at batch size 256 on a single GPU. Gains at small batch sizes, where state traffic matters less, are not shown.
• The GDN end-to-end gains (1.20–1.27×) are smaller than Mamba-2's (1.42–1.47×).