Songlin Yang releases ReplaySSM, cutting hybrid SSM memory traffic in half to enable 2x faster speculative decoding

[1/N] Blog post: https://dao-lab.ai/blog/2026/replayssm/ https://tridao.me/blog/2026/replayssm/ Code: https://github.com/Johnny-Liou/ReplaySSM

@zwljohnny ? Snakes and Ladders: Accelerating State Space Model Inference with Speculative Decoding

https://assets.amazon.science/45/72/4848937d41b0b152ab24f1ca7d41/snakes-and-ladders-accelerating-state-space-model-inference-with-speculative-decoding.pdf

[2/N] In hybrid models, the SSM layers dominate latency up to 100K tokens (not attention).  The SSM layer is deeply memory-bound. Most of its latency comes from memory traffic, loading and storing the state every step, not actual computation.

[3/N] Speculative decoding doesn’t work well on SSMs, and the reason is rollback.

When a draft token gets rejected you have to undo it. In a Transformer, you just trivially move pointer back in the KV cache. In SSM, the state is irreversible (raw inputs are unrecoverable), so vLLM has to keep a separate state for every single draft token just to be able to roll back.

[6/N] ReplaySSM works for both standard and speculative decoding, and it generalizes to GDN. We built it on vLLM and tested at serving batch sizes, 4B to 550B. Standard decoding gets up to 1.48x, and speculative decoding hits 1.96x. The speculative decoding result is particularly important because vLLM's current SSM speculative decode is actually slower than just decoding token-by-token.

[4/N] ReplaySSM addresses both with a simple idea: instead of storing the state every step, it caches recent inputs in a small buffer and only updates the state when the buffer fills.

Same output, but state traffic roughly halves. Each step still loads the state once, but instead of writing the whole state back it just appends the new inputs to the buffer. Rollback also comes for free, because we explicitly cache the recent inputs and rollback becomes a pointer move.

[5/N] ReplaySSM also changes what must be produced. Before, states and outputs were both needed. Now, most steps only need the outputs.

So we compute the output directly from the buffer and never materialize the state at all, which is what the figure below shows for a single decode step. Speculative decoding is where this really pays off, since the same output-only form removes the serial state dependence that made SSMs hard to parallelize and turns the whole draft verification into GEMMs.

@zwljohnny Cool! Is this conceptually similar to the chunkwise parallel form but applied to decoding instead of training?

@zwljohnny for the flush route, I think you can directly compute the new state via a single GEMM by assembling the v_t and k_t vectors into two matrices and multiplying them. This computes the sum of rank-1 outer products in a single GEMM rather than recurrently. Are you already doing this?

Hi, previous rollback approaches still materialize and write at least one recurrent state back to HBM at every step. Committing that state each step makes them still suffer from sequential state dependency, since the state for the last committed tokens has to be rebuilt before the current drafts are verified.

@adityastomar_ Yes, in flush route, both state and output are needed, so we multiply VK^T for the state, and read the output by multiplying with q_t.

In most of the step, we directly calculate the output (we don't need the state)

The key idea is that storing the updated state is the natural design choice, but SSM actually has the flexibility to store recent input (k and v) or store the states.

The output-only route is a benefit from the concept. This part is more similar to chunkwise parallel training (outputs within a chunk are computed from a shared initial state)

@zwljohnny really cool work!

@zwljohnny Nice! 👏🏻👏🏻

@tri_dao Flashrecompute

@tri_dao Latency-critical systems like trading bots live or die on this tradeoff. Do you have numbers on the actual end-to-end inference latency you're hitting at scale?

@VipulS_1 The lack of parallelism and the state materialization cost become more pronounced in GDN and at larger batch sizes, where the benefit of amortizing weight loads becomes smaller.

@zwljohnny @VipulS_1 Thanks @VipulS_1 for mentioning our paper! @zwljohnny I am also wondering whether the replay scheme handles prefix caching in vLLM since we might be storing the stale state for prefix caching if our state is not most up to date.

@tri_dao so it is trading a bit of recompute for skipping the state writeback to HBM, on the bet that for SSMs the memory traffic was the real cost. is that the right read