BuildNot yet confirmed elsewhere1 publisher2 min readPublished
NVIDIA and Princeton cache SSM inputs in place of state to speed hybrid-model decoding
NVIDIA and Princeton's Dao AI Lab cache SSM inputs instead of state, targeting a decode step where SSM layers are the largest component up to 100K tokens. ReplaySSM is already integrated into TensorRT-LLM, vLLM and FlashInfer, NVIDIA says.
The Engineer · Build desk
Key figures: one SSM layer's FP32 state equals the KV cache of 8,192 tokens. That FP32 state is 4 MiB per sequence. An FP8 KV entry is 0.5 KiB per token. SSM layers are the largest component of the decode step up to 100K tokens, at batch 256 on one B300 GPU.
| Figure | Value | Note | Claim |
|---|---|---|---|
| One SSM layer's FP32 state, as KV cache | 8,192 tokens | 19 | |
| FP32 SSM state, one layer | 4 MiB/sequence | 10 | |
| FP8 KV entry | 0.5 KiB per token | 11 | |
| Context where SSM layers are largest decode part, up to | 100,000 tokens | batch 256 on one B300 GPU | 8 |
What happened
- The affected family is large: Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3 and GLM-5.3-Flash all interleave Mamba-2, GDN or KDA layers with a few attention layers.
- An SSM decode step loads and stores the full state with little arithmetic, so the kernel is bound by HBM bandwidth, not compute.
- For speculative decoding, vLLM stores a snapshot after every draft token, which costs one read and T writes per step and caps how many requests fit in GPU memory.
- A quantized SSM state feeds its rounding error back into the recurrence each step, so Qwen3.5-122B-A10B-NVFP4 pairs 4-bit weights with an FP8 KV cache but keeps an FP32 SSM state.
- On most steps, which the post calls non-flush steps, ReplaySSM appends the current inputs to a buffer and computes the output from the checkpoint and buffer without storing a state.
Why it matters
- contradiction SSM layers exist to keep decoding cheap with a fixed-size state, yet at serving batch sizes that fixed state is what makes them the latency bottleneck.
- constraint The 100K-token result comes from a five-to-one layer mix, a 4 MiB FP32 state and batch 256 on one B300. A model with fewer SSM layers, a smaller state or a small batch needs its own measurement before anyone assumes the same split.
- decision Anyone who judged speculative decoding not worth running on a hybrid model at serving batch sizes has a reason to re-measure it on an engine build that carries ReplaySSM.
- cost Engines that cache hybrid-model state now hold a buffer of input records beside each checkpoint, so the memory accounting that caps concurrent requests changes with it.
The post traces all three problems to one design choice. Conventional decoding updates the state, reads the output from it and stores it for the next step in one pass [6]. ReplaySSM separates output computation from state updates [4].
The size figures show why the store hurts. NVIDIA puts one layer's FP32 state at 4 MiB per sequence [10] and an FP8 KV entry at 0.5 KiB per token [11]. That makes one SSM layer's state equal to the KV cache of 8,192 tokens [19]. In the Nemotron example the model has five times as many SSM layers as attention layers [9], so each attention layer comes with SSM state worth 40,960 tokens of KV [20]. Each step both loads and stores that state, so the traffic is equal to reading a KV cache of about 81,920 tokens [21].
That estimate assumes the 0.5 KiB figure applies per attention layer. It sits close to NVIDIA's measurement on Nemotron-3-Super-NVFP4 at batch size 256 on one Blackwell Ultra B300, where SSM layers are the largest part of the decode step up to 100K tokens [8].
The fix works on the unrolled form of the recurrence. Over a window of h tokens, the state is a checkpoint plus a small buffer of recent input records [17]. Skipping the store on non-flush steps removes a 4 MiB write per layer per sequence in the post's example [22]. The post's text, as we have it, ends mid-sentence where the buffer of L records fills, and it includes no speedup figures.
We'd expect the other two benefits to follow from the same layout. Attention rolls back by moving the KV cache pointer, while an SSM update has no easy inverse [12]. A buffer of input records can be cut back to the accepted prefix. A state written on fewer steps also gives rounding error fewer chances to re-enter the recurrence.
The paper presents the full evaluation and is on arXiv, with code on GitHub [5]. NVIDIA says ReplaySSM also runs in production in GLM-5.3-Flash and Kimi-K3 [24].
What to watch
- The paper's measured speedups on hybrid models, and whether independent runs on Qwen3.5 through Qwen3.8 reproduce them.
- Whether speculative decoding on hybrid models at serving batch sizes stops being the marginal gain the post describes today.
- The flush step: how large the buffer length L is and what a flush costs per layer.
Clarity's read
What the record supports and how the coverage leans. The claims behind it follow.
Reality
- Evidence45
- Adoption
- Insufficient
- Hype gap
- Insufficient
- Incentives55
- Confidence50
Claim ledger
Ranked by verification strength, evidence, and original report placement.
- [1]
Hybrid models interleave SSM layers with a few attention layers. NVIDIA Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3 and GLM-5.3-Flash are built this way, with Mamba-2, Gated DeltaNet (GDN) or Kimi Delta Attention (KDA) as the SSM layer.
- [2]
SSM layers are designed to make decoding fast: they use a fixed-size state in place of a key-value cache that grows with context length.
- [3]
At serving batch sizes, SSM layers, not attention, can be the latency bottleneck up to 100K tokens of context.
- [4]
ReplaySSM represents the SSM state as a checkpoint plus a buffer of recent inputs, separating output computation from state updates, and addresses state traffic, rollback for speculative decoding, and error accumulation in a quantized state.
- [5]
ReplaySSM is joint work between NVIDIA and the Dao AI Lab at Princeton. The paper, ReplaySSM: Cache SSM Inputs, Not State, details all three benefits and presents the full evaluation; it is on arXiv, with code on GitHub.
- [6]
Conventional decoding does both of the state's jobs at once: each step updates the state, reads the output from it, and stores it for the next step. All three challenges come from this fusion.
- [7]
An SSM decoding step loads and stores the full state while performing little arithmetic on it, so the kernel is bound by HBM bandwidth, not compute.
- [8]
On Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU, SSM layers are the largest component of the decode step up to 100K tokens.
- [9]
The model has five times as many SSM layers as attention layers.
- [10]
One layer's FP32 SSM state is 4 MiB per sequence.
- [11]
An FP8 KV entry is 0.5 KiB per token.
- [12]
Attention supports speculative-decoding rollback by moving the KV cache pointer, but once an SSM has summarized the drafts into its state, the update has no easy inverse, so rollback becomes expensive.
- [13]
Serving engines such as vLLM store a state snapshot after every draft token and pick the one matching the accepted prefix. That is one read and T writes per step on a memory-bound path, and the snapshots cap how many requests fit in GPU memory.
- [14]
Speculative decoding, a reliable speedup for Transformers, barely helps SSMs at serving batch sizes.
- [15]
A quantized SSM state produces the output and is the input to the next update, so its rounding error re-enters the recurrence at every step. A KV cache has no such loop, which is why KV caches tolerate lower precision than SSM states.
- [16]
Qwen3.5-122B-A10B-NVFP4 uses the NVFP4 4-bit floating-point format for its weights, with an FP8 KV cache and an FP32 SSM state.
- [17]
The SSM recurrence can be unrolled over a window of h tokens. ReplaySSM represents the state as a checkpoint plus a small buffer of recent input records.
- [18]
On most steps, called non-flush steps, ReplaySSM appends the current inputs to the buffer and computes the output from the checkpoint and the buffer without storing a state.
- [19]
One SSM layer's FP32 state equals the KV cache of 8,192 tokens.
- [20]
With five SSM layers per attention layer, the SSM state attached to each attention layer equals the KV cache of 40,960 tokens.
- [21]
Counting both the load and the store of that state each step, SSM traffic per attention layer equals reading a KV cache of about 81,920 tokens, assuming the 0.5 KiB figure applies per attention layer.
- [22]
Because non-flush steps compute the output without storing a state, they skip the 4 MiB per-layer, per-sequence state write in the post's example.
- [23]
ReplaySSM is integrated into NVIDIA TensorRT-LLM, vLLM and FlashInfer.
- [24]
ReplaySSM also runs in production in GLM-5.3-Flash and Kimi-K3.
Sources
1 independent publisher whose own reporting we read for this story.
- developer.nvidia.comReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding
1 article · October 12, 2026
Topics and entities
Follow any of these and your For You feed starts watching them — no settings page required.