What ReplaySSM is solving
Hybrid SSM–Transformer models keep a fixed-size recurrent state for the SSM part. In production, that state becomes a systems problem, not just a modeling detail.
The baseline pattern is “eager state management”: on every decode step you:
- Read the entire SSM state from memory.
- Update it with the current token/input.
- Write the entire state back.
ReplaySSM’s authors point out three problems with this pattern for modern hybrid models 1:
- State I/O dominates: you move a big chunk of memory every token with relatively little math in between (“low arithmetic intensity”), which bottlenecks on memory bandwidth.
- Speculative decoding is expensive to roll back: if you optimistically decode multiple tokens and then need to roll back, you’ve already mutated the recurrent state and must restore it somehow.
- Quantization error compounds: if you store the recurrent state in low precision, the rounding error is fed back into the next recurrent updates, step after step.
ReplaySSM is a different execution scheme for SSMs that attacks all three by changing *what* you cache.
Instead of caching the recurrent state at every step, ReplaySSM:
“unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs.” 1
This is an implementation and orchestration change, not an architecture change.
The paper reports that implementing ReplaySSM in vLLM for Mamba‑2 and Gated DeltaNet yields 1:
- 1.20–1.47× speedup over a tuned vLLM baseline on standard decoding at batch size 256 across four hybrid models.
- About 1.8× the throughput of vLLM’s speculative decoding on two 120B‑class models.
- Better output KL divergence and downstream accuracy when using separate state precisions vs a single low-precision checkpoint.
How ReplaySSM works, step by step
The core idea: treat the SSM state as the result of folding over the input sequence. Instead of storing *every* fold result, you store:
- A checkpoint: the recurrent state as of some earlier position.
- A buffer of recent inputs: the raw SSM inputs since that checkpoint.
Then you replay the SSM update over just the buffered inputs when you need the latest state.
The abstract describes the mechanism at a systems level 1:
“This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic.”
Concretely, decoding with ReplaySSM looks like this:
1. Maintain two pieces of SSM “state”
For each sequence in your batch, maintain:
checkpoint_state: a snapshot of the SSM state at some earlier token position.input_buffer: a ring buffer of the most recent SSM inputs since that checkpoint.
You no longer do a read–modify–write of the full state on every token. Instead:
- The checkpoint is updated only periodically.
- The buffer is updated every token, but it’s just appending inputs, not large state vectors.
2. Compute outputs without eagerly updating the checkpoint
Given a new token:
- Encode it into the SSM input for the current step.
- Append that SSM input to
input_buffer. - To get the *effective* SSM state for this step, conceptually:
- Start from
checkpoint_state. - “Replay” the SSM transition over every input in
input_buffer.
Operationally, you don’t have to materialize intermediate states for each buffered input; the key is that the SSM update for this step is computed as a function of:
- The checkpoint.
- The small, recent set of cached inputs.
The output head can use the resulting SSM contribution to compute the token’s logits. Crucially, you can do this *without* writing a new full checkpointed state back to memory every step.
ReplaySSM postpones the write of the new recurrent state by only updating the checkpoint periodically 1:
“By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic.”
So your per-token cost becomes:
- Writes: just the new input into the buffer.
- Reads: the checkpoint (less often) plus a handful of recent inputs.
3. Periodically “commit” a new checkpoint
After some number of tokens (a tunable interval), you:
- Apply the SSM recurrence from
checkpoint_stateacross all inputs ininput_bufferto obtain the state at the buffer tail. - Write this as the new
checkpoint_state. - Clear or reset
input_buffer.
That’s your amortized state update. Instead of N state writes over N tokens, you get ~1 state write plus N tiny input-buffer writes.
The abstract attributes the throughput gains to this reduction in state traffic 1.
4. Cheap speculative decoding rollback
In speculative decoding, you generate tokens with a fast draft model and later verify them with a target model. If verification rejects some suffix, you need to roll back both tokens and model state.
ReplaySSM’s “inputs not state” strategy makes this nearly free 1:
“Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states.”
Mechanically:
- As you speculatively decode, you only extend
input_buffer(no checkpoint updates). - If verification rejects the last k tokens, you just move the buffer tail pointer back by k entries. The checkpoint never changed.
There is no need to restore a previous full recurrent state from a snapshot. Rollback is equivalent to “forget the last k inputs”.
For verification, because you can replay from the checkpoint plus inputs, you can:
- Recompute just the outputs you need (“output-only verification”).
- Avoid storing or reloading intermediate SSM states for each speculative step.
This is where the ~1.8× throughput vs vLLM’s speculative decoding on two 120B‑class models likely comes from: the state rollback and recompute path is significantly cheaper 1.
5. Separate precisions for serving vs updating state
The same separation between “checkpoint” and “replayed updates” unlocks a quantization trick 1:
“The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates.”
You maintain:
- A high-precision checkpoint: used when you periodically update the true recurrent state.
- A low-precision copy: used for most per-token output computation.
Because the high-precision checkpoint is never derived from the low-precision one, rounding error from the low-precision path does not get fed back into subsequent recurrent updates.
The paper reports that 1:
- Using separate state precisions reduces output KL divergence.
- It also reduces downstream accuracy losses compared with a single shared low-precision checkpoint.
For practitioners: this is a signal that you can push state quantization more aggressively for throughput or memory, without the same compounding degradation that comes from feeding low-precision state back into the recurrence every step—provided you keep a higher-precision update path.
How ReplaySSM fits into an inference stack
The paper implements ReplaySSM in vLLM for Mamba‑2 and Gated DeltaNet 1.
That gives a clear picture of where it lives in a typical stack:
- Model architecture: a hybrid SSM–Transformer model (e.g., with SSM blocks like Mamba‑2 or Gated DeltaNet and attention blocks).
- Inference runtime: vLLM, which already handles KV caching, paged attention, batching, and speculative decoding for transformers.
ReplaySSM is:
- A runtime-level change to how SSM state is represented and updated.
- Orthogonal to transformer KV cache strategies (like attention KV compression) and to weight quantization schemes.
You can think of it as “KV cache-style optimization, but for the SSM state”:
- Traditional transformers cache keys/values per layer, per token.
- SSMs instead cache a fixed-size recurrent state per layer, per sequence.
- ReplaySSM introduces a *temporal unrolling cache* for SSMs: checkpoint + input buffer.
Because the authors integrated it into vLLM and benchmarked at batch size 256 1, it’s meant to work in multi-tenant server settings, not just offline generation.
Why this matters for agentic systems now
For people building agents that run long tasks, call tools, and orchestrate multiple models, ReplaySSM hits a few pain points directly:
- Memory-bound inference in large hybrids
For 120B‑class hybrids, the decoding step is often bound by memory traffic rather than compute. Cutting the “dominant state memory traffic” by around half changes the throughput–latency tradeoff 1.
The reported 1.20–1.47× speedup at batch size 256 for standard decoding and ~1.8× vs vLLM speculative decoding on some 120B‑class models is in the range where you might:
- Shrink the GPU count for a given QPS.
- Or increase concurrency for the same hardware budget.
- Speculative decoding that doesn’t fight your SSM
If you’re already using speculative decoding (e.g., vLLM’s implementation) to reduce latency, ReplaySSM makes it cheaper *specifically for SSM-heavy models* by making rollback pointer-based and enabling “output-only verification” 1.
That’s especially relevant as more reasoning and tool-use models use hybrid or SSM-based backbones where transformer-only speculative tricks don’t map directly.
- Quantized serving without catastrophic drift
Many production deployments run in 8-bit or lower to hit cost/throughput targets. For recurrent models, naive low-precision state storage can destroy long-horizon behavior via accumulated error.
ReplaySSM’s split-precision design (low precision for output computation, higher precision for recurrent state updates) explicitly avoids feeding low-precision errors back into the recurrence 1. The reported improvements in output KL divergence and downstream accuracy relative to a single low-precision checkpoint mean:
- You can keep SSM states aggressively quantized in the *serving* path.
- While preserving long-term behavior via a cleaner *update* path.
That’s directly valuable for agent workloads involving long tool-use traces or multi-step reasoning.
- Better behavior under long, branching traces
Agents that explore multiple candidate futures (planning, search, tool forks) need to:
- Branch model state cheaply.
- Abandon branches without expensive state restoration.
ReplaySSM’s “buffer-pointer rollback” for speculative decoding is exactly the kind of primitive you’d want to adapt for such branching: replay from a shared checkpoint along different input buffers, then discard by moving pointers 1.
How to think about using it
While the paper’s implementation is inside vLLM for Mamba‑2 and Gated DeltaNet 1, the underlying pattern is general:
- If your model exposes an SSM-like recurrent state with a clear update rule per step, and you currently read–update–write that state each token,
- Then you can, in principle, refactor your runtime to:
- Store a less-frequently-updated checkpoint.
- Store recent SSM inputs in a compact buffer.
- Replay the recurrence over the buffer as needed.
- Update checkpoints periodically instead of every step.
- Turn rollback into buffer-pointer manipulation.
The trade-offs you’ll need to surface in your own system (not quantified in the abstract):
- Replay cost vs memory traffic: the more steps you defer checkpoint updates, the more SSM computation you replay when you do update.
- Buffer sizing: how many recent inputs you keep, and per-layer vs shared buffers.
- Precision choices: which format for the high-precision checkpoint vs low-precision serving copy.
The paper’s abstract indicates that with a reasonable implementation inside vLLM, the net effect was positive in both standard and speculative decoding regimes 1.
What to watch next
From the abstract alone, a few likely directions for builders:
- Runtime integrations beyond vLLM: today, ReplaySSM is reported only inside vLLM for Mamba‑2 and Gated DeltaNet 1. Expect other stacks to adopt similar “inputs not state” strategies for their SSM components.
- Combination with other memory optimizations: ReplaySSM naturally composes with:
- KV cache compression on the transformer side.
- Weight quantization on both SSM and Transformer blocks.
- Agent harness support: speculative decoding improvements matter more when your harness (OpenAI-compatible backends, local agent frameworks) can exploit them transparently. vLLM already sits under many of these, so there’s a good pathway for ReplaySSM to show up in end-to-end agents without changes at the orchestration layer.
As hybrid models become the default backbone for agentic systems, treating SSM state as a first-class systems object—rather than a modeling afterthought—will be necessary to keep scaling throughput under control.
What is not documented
The abstract for ReplaySSM 1 does not document:
- The exact SSM update equations or architectural details of Mamba‑2 and Gated DeltaNet.
- The specific length of the input buffer, checkpoint interval, or how those are tuned.
- Implementation details of the replay kernel (e.g., vectorization strategies, fusion with other ops).
- Exact hardware setups, GPU types, or memory hierarchies used for the reported benchmarks.
- Detailed benchmark numbers beyond the reported 1.20–1.47× speedups at batch size 256 and ~1.8× speculative throughput on two 120B‑class models.
- Any integration details with higher-level agent frameworks; only the vLLM implementation is mentioned.
Any further quantitative claims or engineering specifics would require the full paper or implementation code beyond what is summarized in the abstract.