Hybrid-model decoding has a new bottleneck: stop writing the SSM state every step, and speculative throughput reaches 2.10x
Synopsis
Hybrid models interleave state-space layers with a few attention layers, and the state-space layers use a fixed-size state instead of a key-value cache that grows with context, yet at serving batch sizes they become the latency bottleneck and complicate speculative decoding and quantization. ReplaySSM represents the state as a checkpoint plus a buffer of recent inputs, separating output computation from state updates so most steps append inputs without writing state back. Across four hybrid models, the kernel runs 1.46x to 1.80x faster and the decode step 1.20x to 1.47x faster at batch size 256, and two 120B-class models reach 1.85x and 2.10x the speculative decoding throughput of vLLM at batch size 512.
Interpretation
The decoding bottleneck in hybrid models has moved from attention layers to state-space layers, because each step reads the whole state and writes it back. State-space layers were designed for fast decoding, using a fixed-size state in place of a key-value cache that grows with context length; at serving batch sizes the state-space layers, not attention, are the largest component of the decode step up to 100K tokens of context. On Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU, the state-space layers are the largest component of the decode step; the model has five times as many state-space layers as attention layers, and one layer's FP32 state is 4 MiB per sequence while an FP8 KV entry is 0.5 KiB per token.
Representing the state as a checkpoint plus a buffer of recent inputs lets most decode steps append a few input records instead of writing the state back. Conventional decoding does both jobs at once: it updates the state, reads the output from it, and stores it for the next step, so over a window the state is read L times and written L times. ReplaySSM was implemented in vLLM and compared with vLLM's own kernels at their best launch configuration with a window of L=16; at batch size 256 the kernel runs 1.46x to 1.80x faster and the decode step 1.20x to 1.47x faster, while at batch size 1 there is no speedup because other latency sources dominate.
Speculative decoding no longer needs a state snapshot per draft token, because the accepted prefix is committed by advancing the buffer pointer. Serving engines such as vLLM previously stored a state snapshot after every draft token and picked the one matching the accepted prefix, which is one read and T writes per step on a memory-bound path and caps how many requests fit in GPU memory. On GSM8K prompts with T=4 at batch size 512, the two 120B-class models reach 1.85x and 2.10x the speculative serving throughput of vLLM's speculative decoding with matched acceptance lengths; vLLM's kernel at T=8 is almost five times the latency of its standard decoding step, while ReplaySSM verifies eight tokens in about the time of one standard step.
The recurrence checkpoint and the copy read for outputs can use different precisions, so the low-precision copy's rounding error never re-enters the recurrence. A shared low-precision checkpoint is written back into the recurrence at every step, so its rounding error accumulates with generation length; a KV cache has no such loop and therefore tolerates lower precision. On PG-19 text under teacher forcing out to 128K generated tokens, Qwen3.5-4B with NVFP4 reads and an FP32 recurrence cuts the KL divergence from 0.136 to 0.0038; on GPQA-Diamond a shared NVFP4 checkpoint loses 3 to 31 percentage points, while NVFP4 reads with an FP16 or FP8 recurrence keep three of four models within about 2 percentage points of FP32.
Perspective
The report covers four hybrid models, Nemotron 3 with Mamba-2 layers and Qwen3.5 with GDN layers, from 4B dense models on one NVIDIA H100 to 120B-class MoE models on one NVIDIA Blackwell Ultra B300, with performance runs keeping the state in FP32 and a window of L=16, and speculative decoding tests using each model's multi-token prediction head for drafts. No speedup appears at batch size 1, where other latency sources dominate; the gain grows with batch size because the state-space share of the step grows. It is aimed at readers working on hybrid-model inference serving.
A shared low-precision checkpoint accumulates error with generation length; with separate precisions the remaining read error does not grow with length, but the authors note that better quantization of the read copy, such as separate scales along the state and channel dimensions, could cut that remaining error further. On GPQA-Diamond three of four models stay within about 2 percentage points of FP32 under the separate-precision configuration, while Nemotron-3-Nano-4B still trails by 6.7 and 5.6 points. An FP16 recurrence with stochastic rounding tracks FP32 closely, so a half-precision recurrence checkpoint is enough.
