SSM decoding stops writing state every step: caching recent inputs halves state traffic and matches vLLM speculation
Synopsis
State-space model decoding previously read and wrote the full recurrent state at every step, making state traffic the memory-bandwidth bottleneck and leaving speculative decoding without cheap rollback. ReplaySSM represents the state as a checkpoint plus a buffer of recent inputs, writes the checkpoint only when the buffer fills, and computes outputs directly on most steps. Implemented in vLLM for Mamba-2 and Gated DeltaNet, it speeds up standard decoding over a tuned baseline at serving batch size 256, reaches about the throughput of vLLM's speculative decoding on two 120B-class models, and separate state precisions reduce output KL divergence and downstream accuracy losses.
Interpretation
State management in state-space model decoding moves from writing the full recurrent state back every step to representing the state as a checkpoint plus a buffer of recent inputs, which separates output computation from checkpoint updates. Decoding previously followed the recurrence step by step, loading the state, updating it, computing the output from the updated state and storing it back, so state traffic was the memory-bandwidth bottleneck and speculative decoding had to keep a state snapshot per draft token. Implemented in vLLM for Mamba-2 and Gated DeltaNet, it speeds up standard decoding over a tuned vLLM baseline at serving batch size 256 across four hybrid models.
A non-flush step appends the current inputs and computes the output from the checkpoint and the buffer without storing a state, and only a full buffer folds the buffered inputs into the checkpoint. Checkpoint writes are amortized over a window, roughly halving the dominant state memory traffic, while outputs match conventional decoding up to floating-point differences. A fixed window size gives the lowest amortized kernel latency for all four models, matching the traffic analysis prediction.
Speculative rollback becomes a buffer-pointer move and verification needs only outputs without materializing intermediate states, while a low-precision checkpoint serves outputs and a higher-precision one carries the recurrence. vLLM previously stored a state snapshot per draft token, slowing an already memory-bound kernel and capping concurrency, and a low-precision state fed rounding error into later updates that accumulated with generation length. On two 120B-class models at batch size 256 throughput approaches vLLM's speculative decoding, and NVFP4 reads with FP16 or FP8 recurrence cut shared-NVFP4 downstream losses from 13.05, 5.88 and 10.92 points to 2.36, 0.15, 0.10 and 1.77, 0.40, 0.38 points.
Perspective
The report is for inference-systems and kernel engineers and covers Mamba-2 and Gated DeltaNet in vLLM under standard decoding, speculative decoding and low-precision state; performance experiments use an FP32 state, and low-precision experiments lower only state precision while weights keep their released precision. A next direction is letting heads whose decay has become negligible skip the checkpoint read on non-flush steps, turning their flush into a pure write.
Every step still reads the checkpoint, and that read is the next target. In the low-precision results, separate precisions keep the low-precision copy's rounding error out of the recurrence, but rounding in the high-precision path remains and non-flush outputs still depend on the approximate copy; the accuracy loss of a shared low-precision checkpoint appears with long outputs, and on RULER-128K, which has long inputs and short outputs, every configuration stays within noise of FP32.