Language model server developers can reduce token generation costs by collecting recent inputs rather than overwriting a large state each time. A new method from NVIDIA and Princeton University speeds up batched requests and makes it easier to roll back incorrectly predicted tokens or lower storage precision.

1. Store recent inputs instead of the full state

Hybrid models combine state space models, which summarize context in a fixed-size state, with attention layers that refer to previous tokens. State space models read and write the entire state for every token but require little computation, so memory transfer speed can become the bottleneck.

ReplaySSM represents the state as a checkpoint saved at the start of an interval plus a recent-input buffer. It is like keeping a starting balance and recent transactions separately instead of rewriting the entire ledger. The checkpoint is read for every token to compute outputs, but written only once per interval when the buffer fills. Eliminating most writes from a scheme that both reads and writes every token roughly halves the dominant state transfer volume. Output differences from the conventional method are within floating-point calculation error.

2. Reduce state transfers at large batch sizes

The evaluation covered four models: small and large models from two families. The small models used 16-bit weights, and the large models used 4-bit floating-point weights.

The two small models ran on one H100, and the two large models ran on one B300. Performance experiments used 32-bit floating-point states and 16-token input intervals.

At batch size 256, state space model kernels (units of executable code for specific operations) ran 1.46–1.80 times faster, and the full token generation stage ran 1.20–1.47 times faster. The baseline was vLLM kernels with optimized execution settings, and the kernel figures include checkpoint update costs. At batch size 1, other latency factors dominated, leaving no speed gain. Original performance results

3. Roll back incorrect predictions using buffer positions

In speculative decoding, a target model verifies predicted draft tokens together. Conventional implementations kept a state copy for each token to support rollback. ReplaySSM buffers draft inputs and treats only positions through the accepted tokens as valid, so it need not restore states for rejected tokens. Verification also computes only outputs, without creating intermediate states.

The comparison used the two large models evaluated earlier to solve math problems, with 4 draft tokens, 16-token input intervals, and 32-bit states. At batch size 512, throughput was 1.85–2.10 times that of speculative decoding in the existing vLLM framework, with accepted token lengths matched. Original speculative decoding results

4. Prevent low-precision errors from accumulating again

Low-precision states reduce the amount of data read, but rounding errors accumulate when passed into the next update. ReplaySSM separates a high-precision checkpoint for updates from a low-precision copy for outputs. The output copy is not used for the next update, so its errors do not propagate into later updates.

Experiments using standard decoding on GPQA-Diamond with 16-token input intervals used 4-bit floating-point output copies. Updates used 16-bit or 8-bit floating point with stochastic rounding, which probabilistically rounds to neighboring representable values. Three of the four models were within about 2 percentage points of the 32-bit state accuracy baseline. Original accuracy results

Accuracy was not fully preserved. Nemotron-3-Nano-4B dropped by 6.7 and 5.6 percentage points at the two update precisions, respectively. These losses are smaller than the 31.3-percentage-point drop when reads and updates shared a low-precision copy, but show the need for model-specific validation. Original per-model losses

5. Check both the public implementation and its costs

The public code is a fork of vLLM. It provides standard and speculative decoding kernels plus separate precision experiment code. Some paths have been merged into the base framework, but some kernels and precision experiment implementations remain exclusive to this repository.

Reproduction check: the evaluated models are 4B-class and 120B-class models from the Nemotron 3 and Qwen3.5 families. Speculative decoding experiments used GSM8K problems with the two 120B-class models, running each model on one B300.

Below are the reproduction commands from the original source. The first two compare standard decoding speed, and the last compares speculative decoding throughput.

python benchmarks/replayssm/e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
python benchmarks/replayssm/e2e_decode_speedup.py --model-id Qwen/Qwen3.5-4B
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4

Reducing state writes and per-draft-token copies requires maintaining an input buffer. Separating precision also adds writes of the low-precision copy during updates. Rather than eliminating storage costs, this shifts per-token state writes to interval-based updates and buffer management.

ReplaySSM's core idea of separating output computation from state updates connects reduced state transfers, optimized speculative verification, and suppression of low-precision error accumulation.