Can caching SSM inputs instead of state speed up hybrid model decoding?

ReplaySSM caches recent SSM inputs instead of writing state each step, roughly halving state traffic and speeding decoding 1.20–1.47×.

Direct answer

Yes, at serving batch sizes. ReplaySSM replaces per-step state writes with a checkpoint plus a buffer of recent SSM inputs, roughly halving dominant state memory traffic and speeding standard decoding by 1.20–1.47× over a tuned vLLM baseline at batch size 256 across four hybrid models.

1sources cited

This article was generated with WisPaper-powered search and paper analysis.

A fixed-size state still moves every step

State-space models compress past context into a recurrent state whose size does not grow with sequence length. That saves the key-value cache traffic that grows with context in attention, but it does not make decoding free. Conventional SSM decoding loads the full state, applies a decay and a rank-1 update, computes the output from the updated state, and stores it back, every step. The paper calls this eager state management: large state I/O with low arithmetic intensity.

The cost shows up at serving batch sizes. In a hybrid model that interleaves SSM and attention layers, the paper reports that SSM updates can bottleneck decoding up to 100K tokens even though attention cost grows with output context length. The measurement cited is Nemotron-3-Super-120B-A12B-NVFP4 at batch size 256 on one B300. Two further problems follow from the same design: speculative decoding has no cheap way to roll back a rejected suffix, and a quantized state feeds its rounding error into later recurrent updates.

ReplaySSM stores a checkpoint and recent inputs

ReplaySSM unrolls the recurrence so the state is represented as a checkpoint plus a buffer of recent SSM inputs. Output computation reads the checkpoint and the buffered inputs; it does not need to write a new state. The checkpoint is updated only when the buffer fills, folding the buffered records into it. Between flushes, a step reads the checkpoint once and writes nothing.

The traffic accounting explains the gain. A non-flush step costs one checkpoint read (1R). A flush step costs one read plus one write (1R + 1W) and occurs roughly once per L/α steps, where L is the buffer window and α relates to acceptance in speculative decoding. The paper reports that this roughly halves the dominant state memory traffic. The comparison is against conventional decoding, which reads and writes the full state at every step.

Speculative rollback becomes a pointer move

Speculative decoding drafts several tokens, verifies them together, and must restore the state to the last accepted token when a suffix is rejected. With an eagerly updated state, that rollback is expensive. With ReplaySSM, recent inputs sit in a buffer, so rollback is a buffer-pointer operation, and verification can be done on outputs only without materializing intermediate states.

The flush rule changes under speculation because only some drafted tokens are committed. The ring buffer has capacity L + T for T drafted records per step, and the implementation flushes exactly when the committed count h plus T exceeds L. The decision depends only on h and T, both known before the step runs, so it never waits on sampling. The paper notes this is an early flush: it may fire when partial acceptance would have left room, so the flush interval is roughly L/α steps rather than a fixed L, and higher acceptance flushes more often.

Separate precisions keep rounding out of the recurrence

The same separation lets a low-precision checkpoint serve output computation on most steps while a high-precision checkpoint serves recurrent updates. Each flush starts from the high-precision checkpoint and writes a fresh high-precision checkpoint; the low-precision copy is derived from it and never read back into the recurrence, so its rounding error does not accumulate. Non-flush outputs still depend on the approximate copy, and rounding error in the high-precision path remains.

The paper's analysis predicts how error scales with decay and window size, not its magnitude, and assumes independent errors with fixed bias and variance. The authors therefore measure error directly as KL divergence between next-token distributions of low-precision and FP32-state decoding. In the reported results, separate precisions reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint, and the discussion states separate precisions reduce the output KL of an NVFP4 checkpoint by 13–36× at the longest measured generation length.

Speedups appear at large batches, not at batch one

The implementation runs in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, ReplaySSM accelerates standard decoding by 1.20–1.47× over a tuned vLLM baseline across four hybrid models: Nemotron-3-Nano-4B, Nemotron-3-Super-120B-A12B, Qwen3.5-4B, and Qwen3.5-122B-A10B. The baseline is vLLM's kernel at the best launch configuration found by a sweep, not its shipped default.

The batch-size dependence is the key boundary. At batch size 1 the decode-step speedup is below parity, 0.96–0.99×, because the conventional kernel reaches only 21.6% of peak DRAM bandwidth on H100 and 9.3% on B300, so cutting its memory traffic saves nothing. From about batch size 128 the baseline saturates at 84–86% of peak and kernel speedup levels off at 1.4–1.8×, while end-to-end speedup keeps growing because SSM layers take a larger share of the decode step. The end-to-end gain is also largest at short and medium sequence lengths; at long sequences attention dominates the step because its KV-cache traffic grows linearly.

Speculative decoding gains about 1.8× on 120B models

On the two 120B-class models, ReplaySSM delivers about 1.8× the throughput of vLLM's speculative decoding. Those experiments use each model's multi-token-prediction head as the drafter. This is a different comparison from the 1.20–1.47× standard-decoding figure: one measures throughput against vLLM's speculative decoding, the other measures per-step speedup against tuned conventional standard decoding at batch size 256.

The accuracy results carry their own conditions. The low-precision experiments use emulated quantization, and FP32 refers to the SSM state only; weights keep their released precision and activations are BF16. Accuracy is reported in bins defined by the FP32 run, and because FP32 accuracy itself falls by 21 to 50 points across bins on MMLU-Pro and GPQA-Diamond, the quantity to read is the gap between a configuration and FP32 within a bin, not absolute accuracy. Wherever the shared NVFP4 checkpoint loses accuracy, its gap to FP32 grows with length; on the three larger models it matches FP32 in short bins and loses in the top bin, by up to 28.2 points on Nemotron-3-Super-120B-A12B GPQA-Diamond. NVFP4 reads with FP8+SR recurrence stay within the FP32 interval in 41 of 48 bins.

The checkpoint read is the next bottleneck

ReplaySSM removes the state write from non-flush steps but still reads the checkpoint at every step, because output computation needs it. The paper identifies that read as the next target. The checkpoint enters each step scaled by a per-head cumulative decay, so a head whose decay has become negligible can skip the read on non-flush steps and turn its flush into a pure write. The paper outlines this direction rather than measuring it.

For a reader deciding whether the result applies, the paper's own split is the useful one: the gain is a serving-batch phenomenon, largest at short and medium sequence lengths, and it depends on the conventional kernel being memory-bound. At batch size 1, or at long sequences where attention dominates, the same mechanism does not deliver the reported speedup. The separate-precision benefit is a distinct claim, measured as KL divergence and downstream accuracy rather than latency, and it should be read against the FP32 gap within each length bin.

Sources used in this answer

1

ReplaySSM: Cache SSM Inputs, Not State

ReplaySSM represents the SSM state as a checkpoint plus a buffer of recent inputs, roughly halving dominant state memory traffic and accelerating standard decoding by 1.20–1.47× over a tuned vLLM baseline at batch size 256 across four hybrid models.