ReplaySSM: Cache SSM Inputs, Not State

Ze-Wei Liou; Harris Nover; Amey Naik; Andrii Skliar; Igor Shovkun; Ido Shahaf; Tri Dao

cs.LGOct 9, 2026

ReplaySSM: Cache SSM Inputs, Not State

Authors: Ze-Wei Liou, Harris Nover, Amey Naik, Andrii Skliar, Igor Shovkun, Ido Shahaf, Tri Dao

Organizations: Princeton University · NVIDIA

Abstract

State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models. This eager state management causes large state I/O with low arithmetic intensity, leaves speculative decoding no cheap rollback, and feeds state quantization error into subsequent recurrent updates. We introduce ReplaySSM, which unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs. This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic. Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states. The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates. We implement ReplaySSM in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, it accelerates standard decoding by $1.20$ -- $1.47\times$ over a tuned vLLM baseline across four hybrid models, and delivers about $1.8\times$ the throughput of vLLM's speculative decoding on the two 120B-class models. Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.

Figures & tables

Figure 1: ReplaySSM overview. Top: conventional decoding loads and stores the full state at every step. ReplaySSM unrolls the recurrence to separate output computation from checkpoint updates: it caches recent SSM inputs and folds them into the checkpoint periodically, avoiding full-state writes on most steps. Bottom: three challenges of conventional decoding, and how ReplaySSM addresses each. (1) SSM layers can bottleneck hybrid models up to 100K tokens (Nemotron-3-Super-120B-A12B-NVFP4, batch size 256, one B300); ReplaySSM roughly halves their state traffic. (2) State rollback, expensive to support in conventional speculative decoding, becomes a buffer-pointer move. (3) A low-precision checkpoint serves outputs while a high-precision one carries the recurrence, so rounding error no longer re-enters the recurrence.

Checkpoint configuration	MMLU-Pro	GPQA-Diamond	Math-L5
Conventional FP32 (reference)	$\mathbf{77.6\%}$	$\mathbf{73.7\%}$	$\mathbf{92.3\%}$
	( $67.2$ / $81.4$ / $79.0$ / $82.9$ )	( $56.7$ / $79.3$ / $76.3$ / $82.7$ )	( $89.5$ / $95.5$ / $93.7$ / $90.7$ )
Shared FP16+SR	$\mathbf{-0.09}$	$\mathbf{-2.02}$	$\mathbf{0.00}$
	( $-0.3$ / $+0.3$ / $-0.5$ / $+0.1$ )	( $-2.2$ / $-3.4$ / $-1.5$ / $-1.0$ )	( $-1.1$ / $+0.2$ / $-0.3$ / $+1.2$ )
Shared FP8+SR	$\mathbf{-0.24}$	$\mathbf{-1.05}$	$\mathbf{+0.01}$
	( $-0.6$ / $+0.2$ / $-0.4$ / $-0.1$ )	( $-1.0$ / $+0.0$ / $-0.5$ / $-2.7$ )	( $-1.0$ / $+0.7$ / $-0.5$ / $+0.9$ )

Table 2: Downstream accuracy. Bold: mean change from FP32 over the four models, in percentage points (FP32 row: accuracy). Gray: per-model values, Nemotron-3-Nano-4B / Nemotron-3-Super-120B-A12B-NVFP4 / Qwen3.5-4B / Qwen3.5-122B-A10B-NVFP4. Red: losses

\geq 5

points.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Method	R/W	Model (GB)	Nsight (GB)	Note
vLLM snapshots	$1\mathrm{R}+T\mathrm{W}$	$3.2$ – $9.7$	$3.22$ – $9.73$	grows with $T$
The Mamba in the Llama	$1\mathrm{R}+2\mathrm{W}$	$3.2$	$3.23$ – $3.45$	one fused launch
Snakes and Ladders, Activation Replay	$3\mathrm{R}+3\mathrm{W}$	$6.4$	$6.37$ – $6.55$	replay, copy, verify
Snakes and Ladders, Joint Attainment	$3\mathrm{R}+4\mathrm{W}$	$7.5$	$7.46$ – $7.70$	two rows, then a gather
STree, algorithm	$2\mathrm{R}+1\mathrm{W}$	$3.2$	$3.22$ – $3.35$	replay, read-only verify
STree, released path	$5\mathrm{R}+4\mathrm{W}$	$9.7$	$9.67$ – $9.79$	two extra copies

Appendix

Table 3: SSM-state traffic per layer per verification step at

B=256

, where one state is

4.2

MB and the batch holds

1.07

GB. Model: reads and writes times

1.07

GB. Nsight: DRAM bytes at base clock with a cold L2, summed over the launches of one step, for

T=2

8

. The two state copies of the released STree path are memory copies rather than kernels, so their bytes are taken from the model.

Figure 7: Per-step speculative kernel latency against prior methods on Nemotron-3-Super-120B-A12B (Mamba-2;

B=256

L=16

, FP32 state, one B300), normalized to conventional standard decoding (

T=1

). STree is measured as its algorithm, replay plus read-only tree scan without the released path’s two state copies, and every kernel runs at its best swept launch configuration. STree does not depend on acceptance. Each band starts at the all-rejected cost (

a=1

) and ends at the most expensive outcome,

a=T

for ReplaySSM, whose bars include flushes, and

a=T-1

for The Mamba in the Llama, whose replay is skipped when every draft is accepted.

Symbol	Meaning (Mamba-2)	GDN
$S_{\mathrm{ckpt}}$	Checkpoint state in HBM; the recurrent state at the end of the last flushed window.	same
$S_{\mathrm{ckpt}}^{\mathrm{HP}}$ , $S_{\mathrm{ckpt}}^{\mathrm{LP}}$	High- and low-precision copies of the checkpoint. HP carries the recurrence across flushes; LP serves non-flush output computation.	same
$S$	Working state on chip; written to HBM only at a flush, as the new checkpoint.	same
$x=(\alpha,v,k)$ , $q$	SSM inputs of one token: decay, value, key, and query (Eq. 1 ).	Inputs are $(\alpha,\beta,v,k)$ with write strength $\beta$ (Eq. 2 ). Records store the correction $u$ in place of $v$ ; two-precision records are $(\alpha,\beta,u,v,k)$ .
$u$ , $U_{c}$ , $U_{d}$	–	Correction $\beta(v-\alpha Sk)$ (Eq. 3 ); committed and draft corrections as matrix columns (Eq. 10 ).
$\mathcal{B}$ , $\mathcal{B}_{i:j}$	Replay buffer and its records at positions $i$ to $j$ ; position $1$ is the oldest record since the last flush.	same

Appendix

Table 4: Symbols in Algorithms 4 – 6 .

Model	Weights	SSM / Attn. layers	Key heads	Value heads	$d_{k}$	$d_{v}$
Nemotron-3-Nano-4B	BF16	21 / 4	8	96	128	80
Nemotron-3-Super-120B-A12B	NVFP4	40 / 8	8	128	128	64
Qwen3.5-4B	BF16	24 / 8	16	32	128	128
Qwen3.5-122B-A10B	NVFP4	36 / 12	16	64	128	128

Appendix

Table 5: Layer counts and SSM-layer shapes. Each key head serves the value heads of its group.

Model	Weights	$B\!=\!1$	$B\!=\!8$	$B\!=\!32$	$B\!=\!128$	$B\!=\!256$	$B\!=\!512$
Nemotron-3-Nano-4B	BF16	$0.98\times$	$1.03\times$	$1.16\times$	$1.38\times$	$1.47\times$	$1.51\times$
Nemotron-3-Super-120B-A12B	NVFP4	$0.96\times$	$1.01\times$	$1.06\times$	$1.24\times$	$1.42\times$	$1.72\times^{\dagger}$
Qwen3.5-4B	BF16	$0.99\times$	$1.03\times$	$1.08\times$	$1.21\times$	$1.27\times$	$1.30\times$
Qwen3.5-122B-A10B	NVFP4	$0.96\times$	$1.03\times$	$1.09\times$	$1.16\times$	$1.20\times$	$1.22\times^{\dagger}$

Appendix

Table 6: End-to-end per-step speedup over tuned conventional vLLM standard decoding (

L=16

, FP32 state). † Inflated by the B300 power cap, which throttles the conventional kernel inside the full model but not ReplaySSM; the kernel speedups in Table 7 are unaffected.

Figure 8: Standard decoding across batch size on Nemotron-3-Nano-4B (

L=16

, FP32 state, one H100). Left: DRAM bandwidth of the conventional kernel as a fraction of peak. Middle and right: isolated kernel and end-to-end per-step speedup over tuned conventional vLLM.

Figure 9: Amortized kernel speedup over tuned conventional vLLM across replay-window sizes, amortized over one replay cycle (

B=256

, FP32 state).

L=16

is best in all four settings. The vertical axis starts at

0.5

Table 13

	Nemotron-3-Super-120B-A12B-NVFP4		Qwen3.5-122B-A10B-NVFP4
	Conventional	ReplaySSM	Conventional	ReplaySSM
$T\!=\!2$	$1.95$	$1.95$	$1.95$	$1.95$
$T\!=\!4$	$3.43$	$3.44$	$3.56$	$3.57$
$T\!=\!6$	$4.24$	$4.20$	$4.73$	$4.76$
$T\!=\!8$	$4.40$	$4.44$	$5.52$	$5.53$

Appendix

Table 10: Mean acceptance length at

B=256

Figure 10: Speculative kernel latency on Qwen3.5-122B-A10B (GDN;

B=256

L=16

, FP32 state), normalized to conventional standard decoding (

T=1

). ReplaySSM bands span all drafts rejected to all drafts accepted and include flush costs.

Figure 11: Non-flush kernel speedup over an FP32 state by read-checkpoint precision (

B=256

L=16

, one B300).

Figure 12: Output KL from conventional FP32-state decoding under teacher forcing on PG-19 (

L=16

), the three models outside Figure 6 , on one shared vertical axis. Nemotron-3-Nano-4B stops at 32K (see text). Arrows mark the gap between a shared NVFP4 checkpoint and NVFP4 reads with FP32 recurrence at the last plotted

N

Figure 13: Output KL across replay-window sizes on Qwen3.5-4B.

L=1

requantizes the state at every step, as conventional decoding does, and exists only for a shared checkpoint. Separate configurations use FP32 recurrence. The grey band is the FP32-state ReplaySSM floor over

L=1

–

64

		Generation length $N$
Model	Precision	1K	8K	32K	128K
Nemotron-3-Nano-4B	FP8	0.59	1.3	6.6	–
	NVFP4	0.51	0.67	0.90	–
Nemotron-3-Super-120B-A12B	FP8	1.1	1.4	2.0	2.2
	NVFP4	0.44	0.45	1.4	0.95
Qwen3.5-4B	FP8	0.76	1.2	2.2	div.

Appendix

Table 11: Output KL of a shared checkpoint with round-to-nearest divided by that with stochastic rounding, at matched precision and

L=16

. Above one, RTN is worse. div.: the RTN state became non-finite before this

N

. Nemotron-3-Nano-4B is read to 32K (Appendix E.4.3 ).

	MMLU-Pro		GPQA-Diamond		Math-L5
Configuration	$\Delta$ (pp)	Tokens (K)	$\Delta$ (pp)	Tokens (K)	$\Delta$ (pp)	Tokens (K)
Nemotron-3-Nano-4B
FP32 accuracy (%)	$67.2$	2.9 / 6.9	$56.7$	10.3 / 26.7	$89.5$	2.1 / 4.5
Shared FP16+SR	$-0.3$	2.9 / 6.8	$-2.2$	10.7 / 29.9	$-1.1$	2.2 / 4.9
Shared FP8+SR	$-0.6$	3.0 / 6.9	$-1.0$	10.7 / 29.0	$-1.0$	2.2 / 4.9
Shared NVFP4+SR	$-21.0$	2.1 / 3.6	$-31.3$	4.3 / 5.0	$-41.7$	1.7 / 3.2

Appendix

Table 12: Downstream accuracy for every model, configuration, and benchmark.

\Delta

is the change from that model’s FP32 reference in percentage points, and the FP32 row gives its absolute accuracy. Tokens are the mean and 90th percentile of output length in thousands.

Configuration	Nemotron-3- Nano-4B	Nemotron-3-Super- 120B-A12B-NVFP4	Qwen3.5-4B	Qwen3.5-122B- A10B-NVFP4
FP32 score	$90.4$	$96.9$	$95.2$	$96.8$
Shared FP16+SR	$+0.1$	$-0.1$	$+0.1$	$-0.1$
Shared FP8+SR	$0.0$	$-0.1$	$+0.1$	$0.0$
Shared NVFP4+SR	$-0.3$	$0.0$	$0.0$	$+0.1$
Separate FP8 / FP16+SR	$+0.1$	$-0.1$	$+0.2$	$+0.1$
Separate NVFP4 / FP16+SR	$+0.2$	$-0.1$	$+0.1$	$+0.1$

Appendix

Table 13: RULER-128K score for every model and configuration. The FP32 row gives the absolute score and the other rows the change from it in percentage points.

Figure 14: Accuracy by quartile of FP32 reference generation length, all models and benchmarks. Bins are defined by the FP32 run. Within each benchmark row the three larger models share one vertical range; the Math-L5 row starts at 50%.

Explore similar work

SketchSSM: Write to the Full State, Read from a Compact Sketch
Sep 27, 2026Omin Kwon, JoongWon Shin, Minseo Kim +3LLM Inference Acceleration Linear Attention
Seoul National University · UC Berkeley · KAIST
Just Let Linear States Forget the Distant Past: Prefix Caching via Suffix Replay for Hybrid LLMs
Sep 27, 2026Yirui Liu, Ruoling Qi, Xuaner Wu +6Long-Context Language Model Inference LLM Inference Acceleration
Institute of Artificial Intelligence, China Telecom (TeleAI) China · Shanghai Jiao Tong University · Individual Researcher +1
Faster LLM Inference via Sequential Monte Carlo
Apr 17, 2026Yahya Emara, Mauricio Barba da Costa, Chi-Chih Chang +4Speculative Decoding LLM Inference Acceleration
Cornell University · Makora · MIT +1

Sep 27, 2026, cs.LGJ/K move · Enter open · S save

SketchSSM: Write to the Full State, Read from a Compact Sketch

Omin Kwon, JoongWon Shin, Minseo Kim, Kurt Keutzer+2

Seoul National University · UC Berkeley · KAIST

LLM Inference Acceleration Linear Attention Hybrid Attention

Open paper AskPDF arXiv

Hybrid-attention models replace most softmax attention layers with linear attention, reducing KV-cache growth and enabling larger decode batches where recurrent-state access becomes a major bottleneck. ReplaySSM amortizes state updates by buffering keys and values, but each new query still requires a full-state read even though the state remains unchanged between state updates. We observe that low-rank state-weighted query approximation accurately preserves state-read outputs. Although future queries are unknown, the basis vectors used to approximate them can be fixed offline. Based on this observation, we introduce SketchSSM, which preserves full-state updates while approximating reads. At each state update, SketchSSM reads the full state once to precompute outputs for these basis vectors, storing them in a compact sketch. Each subsequent decode step combines the sketch vectors with query-dependent coefficients to reconstruct the output without a full-state read. Across four Mamba-2-, GDN-, and KDA-based models, SketchSSM at a mean sketch rank of 8 reduces state-access traffic by approximately 10

\times

while matching the average accuracy of the FP32 full-state baseline across four decode benchmarks, and preserves recall on four RULER retrieval tasks. At this rank on one NVIDIA B300, linear-attention kernel speedups over the Standard vLLM baseline reach 7.30

\times

, 5.02

\times

, and 5.24

\times

for Mamba-2, GDN, and KDA, respectively, with up to 2.77

\times

higher decode throughput on Nemotron 3 Super.

Abstract

Explore similar work

SketchSSM: Write to the Full State, Read from a Compact Sketch

Just Let Linear States Forget the Distant Past: Prefix Caching via Suffix Replay for Hybrid LLMs

Faster LLM Inference via Sequential Monte Carlo