ReplaySSM Rethinks State Storage to Accelerate Hybrid LLM Decoding and Inference
NVIDIA and Princeton separate output computation from recurrent state updates to double throughput.
1Hybrid Architecture Decoding Bottleneck
Hybrid LLMs interleave state-space model layers like Mamba-2, GDN, and KDA with traditional attention. While SSMs use fixed-size states to bypass growing KV caches, serving batch sizes make SSM layers the primary latency bottleneck up to 100K context tokens due to memory bandwidth limits.
2Core Mechanism of ReplaySSM
Developed jointly by NVIDIA and Princeton's Dao AI Lab, ReplaySSM decouples output generation from recurrent state updates. Instead of loading and writing full state matrices every step, it caches recent token inputs in a buffer and computes outputs directly against a frozen state checkpoint.
3Two-Phase Step Execution Model
During non-flush steps, the engine appends incoming input records to a local buffer and evaluates outputs without updating persistent memory. Once the buffer reaches window length L (typically 16), a flush step merges buffered tokens into the checkpoint and writes to HBM exactly once.
4Halving State Memory Traffic
Standard decoding requires continuous state reading and writing for every single token. ReplaySSM reads the checkpoint per step but writes only compact input vectors. This cuts dominant state traffic roughly in half, addressing bandwidth saturation on high-throughput serving systems.
5Standard Decoding Performance Gains
Evaluations on NVIDIA Blackwell Ultra B300 GPUs at batch size 256 show substantial kernel acceleration. Across Nemotron 3 and Qwen3.5 architectures, SSM decoding kernels run 1.46x to 1.80x faster, resulting in 1.20x to 1.47x end-to-end decode step speedups over tuned vLLM baselines.
6Eliminating Speculative Snapshots
Speculative decoding in SSMs previously required taking expensive memory snapshots after every draft token to allow rollbacks. ReplaySSM eliminates these snapshots entirely by appending draft inputs directly to the buffer, rolling back seamlessly by moving a buffer pointer without restoring states.
7Output-Only Verification Math
During verification, ReplaySSM evaluates draft tokens without materializing intermediate state matrices. By restructuring vector products into chunkwise matrix multiplications, draft verification runs efficiently on GPU Tensor Cores, sharing key-query projections across value heads to reduce arithmetic overhead and save on-chip memory.
8Doubling Speculative Serving Throughput
On 120B-class MoE models executing GSM8K prompts with four draft tokens, ReplaySSM delivers 1.85x to 2.10x higher serving throughput at batch size 512. The gains stem directly from removing snapshot memory writes and utilizing compute-dense Tensor Core verification kernels.
9Dual-Precision Quantization Strategy
Quantizing recurrent SSM states usually causes rapid compounding errors because rounded outputs loop back into subsequent updates. ReplaySSM solves this by maintaining a high-precision checkpoint for the recurrence update path alongside a low-precision copy dedicated strictly to non-flush output reads.
10Mitigating Long-Context Error Drift
On teacher-forced PG-19 tests out to 128K tokens on Qwen3.5-4B, an FP32 recurrence with NVFP4 reads slashed KL divergence from 0.136 down to 0.0038. Isolate rounding errors do not compound across recurrence steps, maintaining baseline output stability across extended generation sequences.
11Benchmark Accuracy Retention
Under rigorous GPQA-Diamond evaluations, standard shared NVFP4 checkpoints degrade accuracy by up to 31 percentage points. Decoupled precision using FP16 or FP8 recurrence alongside NVFP4 reads recovered these drops, keeping three out of four evaluated models within 2 percentage points of FP32.
12Ecosystem Integration and Deployment
The technique is integrated into major inference runtimes including TensorRT-LLM, vLLM, and FlashInfer, while powering production deployments in GLM-5.3-Flash and Kimi-K3. The open-source code enables broad hardware adoption, establishing a scalable paradigm for next-generation hybrid model inference.