ReplaySSM: Cache SSM Inputs, Not State

Summary (Overview)

• Core innovation: ReplaySSM represents the recurrent state of State-Space Models (SSMs) as a checkpoint plus a bounded buffer of recent inputs, rather than eagerly materializing the state at every decoding step. This separates output computation from checkpoint updates.
• Key insight: The recurrence can be unrolled so that the current state equals a decayed checkpoint plus a sum over recent inputs, enabling state writes only on periodic "flush" steps rather than every step.
• Performance gains: At serving batch size 256, ReplaySSM accelerates standard decoding by 1.20–1.47× over a tuned vLLM baseline across four hybrid models, and delivers ~1.8× the throughput of vLLM's speculative decoding on 120B-class models.
• Three challenges addressed: (1) Large state I/O with low arithmetic intensity, (2) expensive rollback in speculative decoding, and (3) propagation of quantization error through the recurrence.
• Low-precision extension: The separation enables different precisions for output computation vs. recurrent updates, reducing KL divergence by 13–36× at long generation lengths compared to shared low-precision checkpoints.

Introduction and Theoretical Foundation

Background

Long chain-of-thought reasoning makes decoding a substantial cost in LLM inference. While Transformer attention requires a KV cache whose size grows with sequence length, SSMs (Mamba-2, Gated DeltaNet, Kimi Delta Attention) compress past information into a fixed-size recurrent state. Hybrid architectures interleaving SSM and attention layers are widely adopted in production models including Nemotron-3, Qwen3.5, Kimi-K3, and GLM-5.3-Flash.

Three Challenges of Conventional SSM Decoding

Challenge 1: Large state I/O with little compute. Each decoding step loads the full state, applies a decay and rank-1 update, and stores the state back. This is memory-bound with low arithmetic intensity. For Nemotron-3-Super-120B-A12B at batch size 256 on one B300, SSM state updates remain the largest latency component up to 100K tokens because SSM layers have five times as many layers as attention, and each FP32 state is 4 MiB per sequence.

Challenge 2: No cheap rollback. Speculative decoding requires restoring the state to the last accepted token when draft tokens are rejected. SSMs have no token history to point into, so vLLM stores a state snapshot per draft token, adding $O(T)$ state writes per step and capping concurrency.

Challenge 3: Quantization error propagation. In low-precision decoding, the quantized state both produces the output and feeds the next update, so rounding error re-enters the recurrence at every step and accumulates with sequence length.

Key Theoretical Insight

For Mamba-2, the recurrence can be unrolled. Let $S_{\mathrm{ckpt}}$ be the state before a window of $h$ inputs indexed locally by $j = 1, \dots, h$:

$$S_{t} = \alpha_{t} S_{t-1} + v_{t} k_{t}^{\top} = \gamma_{h} S_{\mathrm{ckpt}} + \sum_{j=1}^{h} \gamma_{j+1:h} v_{j} k_{j}^{\top} \tag{4}$$

where $\gamma_{i:j} = \prod_{\ell=i}^{j} \alpha_{\ell}$ is the cumulative decay. The state can be reconstructed from a checkpoint and recent inputs rather than maintained eagerly.

Methodology

Core Algorithm: ReplaySSM Standard Decoding

The state is represented as:

• Checkpoint $S_{\mathrm{ckpt}}$: stored in HBM, written only on flush steps
• Replay buffer $\mathcal{B}_{1:h}$: bounded buffer of recent SSM inputs $\{(\alpha_j, v_j, k_j)\}$

The core operation reconstructs the state:

$$\mathrm{REPLAY}(S_{\mathrm{ckpt}}, \mathcal{B}_{1:h}) = \gamma_{h} S_{\mathrm{ckpt}} + \sum_{j=1}^{h} \gamma_{j+1:h} v_{j} k_{j}^{\top} \tag{5}$$

Algorithm 1 (Standard Decoding):

1. Append current inputs to buffer
2. Reconstruct state via REPLAY
3. Compute output $o = Sq$
4. If buffer full ($h = L$): flush — write checkpoint to HBM, clear buffer

Over an $L$-step window: both methods perform $L$ full-state reads, but conventional decoding performs $L$ full-state writes while ReplaySSM performs only 1. Window size $L = 16$ is used throughout.

Speculative Decoding Extension

Algorithm 2 extends standard decoding to verify $T$ tokens per step:

• Append all $T$ draft inputs to buffer first
• Compute outputs without materializing intermediate states via OUTPUTONLY evaluation
• Commit only accepted prefix by moving buffer pointer
• Flush only committed records

Output-only verification avoids materializing states:

$$\text{OUTPUTONLY}(S_{\text{ckpt}}, \mathcal{B}_{1:h+T}, Q) = S_{\text{ckpt}} Q D + V \big(W \odot (K^{\top} Q)\big) \tag{6}$$

where $D$ is diagonal with $D_{ss} = \gamma_{h+s}$, $\odot$ is elementwise product, and $W$ holds cumulative decays under a causal mask.

The buffer is implemented as a circular buffer with capacity $L + T$, so dropping the committed prefix advances the ring origin instead of copying data.

Low-Precision State Extension

Algorithm 3 keeps two checkpoints:

• $S_{\mathrm{ckpt}}^{HP}$: high-precision checkpoint for recurrent updates
• $S_{\mathrm{ckpt}}^{LP}$: low-precision copy for non-flush output computation

Non-flush steps read the low-precision copy; flush steps read the high-precision checkpoint, reconstruct the committed state, and cast it independently into both formats. This breaks the feedback loop where rounding error re-enters the recurrence.

For GDN, corrections $u_t$ cached on non-flush steps come from the low-precision checkpoint, so the algorithm also caches $v$ and recomputes corrections from the high-precision state at flush time.

Empirical Validation / Results

Experimental Setup

4B models run on one H100; 120B-class models on one B300. All performance experiments store the recurrent state in FP32.

Standard Decoding Results

At $B = 256$ and $L = 16$, ReplaySSM speeds up:

• SSM kernel: 1.46–1.80×
• Whole decode step: 1.20–1.47×

Speedup grows with batch size as the conventional kernel becomes memory-bound. At $B = 1$, the kernel is latency-bound and ReplaySSM is within a few percent of baseline.

Speculative Decoding Results

At $L = 16$, $T = 4$, with MTP-head drafting on GSM8K:

• At $B = 512$: 1.85× and 2.10× the throughput of conventional speculation; 2.16× and 1.90× that of standard decoding
• Conventional speculation stops gaining beyond $B = 128$ due to snapshot storage overhead
• Verification kernel: ReplaySSM verifies 8 tokens for 0.89–1.21× of a standard step (flushes included), vs. 4.66× for conventional at $T = 8$

Low-Precision State Results

KL divergence from FP32-state decoding (Qwen3.5-4B, PG-19, teacher forcing):

Downstream accuracy (mean change from FP32 over four models, percentage points):

Loss concentrates in long generations. On GPQA-Diamond (median 61K tokens in top quartile), shared NVFP4+SR loses 20.9–28.2 points in the longest quartiles, while separate NVFP4/FP8+SR keeps gaps at 2.0–3.4 points.

Theoretical and Practical Implications

Theoretical Contributions

1. Reformulation of state management: Shows that eager state materialization is unnecessary; the state can be represented as a checkpoint plus input buffer, separating output computation from state updates.

2. Output-only evaluation: The factorization $o = (vk^{\top})q = v(k^{\top}q)$ avoids materializing the $d_v \times d_k$ state, reducing FLOPs and enabling efficient verification.

3. Breaking quantization feedback: Separate precisions for output computation and recurrent updates eliminate the accumulation of rounding error through the recurrence, with the remaining error being only the read error of the low-precision copy.

Practical Implications

• Production adoption: ReplaySSM is adopted in large-scale production models and integrated into LLM serving systems
• Memory traffic reduction: Roughly halves dominant state memory traffic, critical for memory-bound SSM kernels at serving batch sizes
• Speculative decoding viability: Makes speculative decoding practical for SSMs at serving batch sizes, where it previously provided little benefit
• Low-precision serving: Enables aggressive state quantization (NVFP4 reads) with minimal accuracy loss when paired with FP8/FP16 recurrence checkpoints

Conclusion

ReplaySSM unrolls the SSM recurrence to represent the state as a checkpoint plus a buffer of recent inputs, separating output computation from checkpoint updates. This yields three benefits:

1. Standard decoding: State writes occur only on flush steps, roughly halving dominant state traffic (1.20–1.47× end-to-end speedup)
2. Speculative decoding: Rollback becomes a buffer-pointer move with output-only verification (~1.8× throughput of conventional speculation)
3. Low-precision states: Separate read/recurrence precisions keep rounding error out of the recurrence (13–36× KL reduction at long generations)

Future direction: Every step still reads the checkpoint. Since each checkpoint enters steps scaled by a per-head cumulative decay, heads with negligible decay could skip reads on non-flush steps, turning flushes into pure writes — potentially further reducing memory traffic.