NVIDIA Releases ReplaySSM to Minimize SSM State Writes

NVIDIA and Princeton University's Dao AI Lab have released the ReplaySSM technique. This technique addresses the bottleneck of saving state to HBM at every step in SSM layers of hybrid models by representing the state as a checkpoint and a recent input buffer.

Core Mechanism

• Non-flush steps: Add inputs to the buffer and compute only the output, skipping state saving.
• Flush steps: Reflect buffer contents into the checkpoint and save to HBM only once.
• Speculative decoding rollback: Handled by moving ring buffer pointers, eliminating state snapshot saving overhead.

Performance Improvements

Against the vLLM tuning baseline with FP32 state and length 16, the following performance improvements were achieved.

• Standard decoding: SSM kernel 1.46~1.80x and overall decode step 1.20~1.47x speedup at batch size 256.
• Speculative decoding: Nemotron-3-Super 1.85x and Qwen3.5-122B 2.10x throughput improvement at batch size 512.
• Verification latency: While vLLM's T=8 latency is 4.66x, ReplaySSM maintains 0.89~1.21x levels.

Quantization and Accuracy

Separated low-precision (NVFP4) for reading and high-precision (FP32) for recursion to prevent error accumulation. Using NVFP4 state reduced GPQA-Diamond accuracy loss to 1.77~2.36 points. The Mamba-2 non-flush kernel is 3.8x faster than FP32 when reading NVFP4.

Implementation and Limitations

• vLLM v0.31.0: Supports only Mamba-2 standard decoding (--use-replayssm).
• ReplaySSM GitHub: GDN path and low-precision kernels exist only in the research fork.
• Application conditions: Gains increase with larger batch sizes, but it becomes slower for batch size 1 local inference.