[Kernel] ReplaySSM: cache SSM inputs for faster Gated DeltaNet speculative decode

[Johnny-Liou]

Review order. This PR is stacked on #49847 (Mamba2 speculative decode). Because it targets main, the diff also includes that PR's commits. Please review this PR after #49847 is reviewed and merged.

Background

This is the fourth sub-PR split from the larger ReplaySSM draft PR #47576, which contains the complete Mamba2 and Gated DeltaNet implementation for both standard and speculative decode. Mamba2 standard decode landed as 866fea2b9 (#48018); GDN standard decode is #48792, and Mamba2 speculative decode is #49847. This PR adds GDN speculative decode. It is a collaboration between Dao AI Lab and NVIDIA, and the full derivation and figures are in the blog post. AI assistance was used for parts of this work, and the authors have reviewed every changed line.

Core concept

To roll back rejected drafts, vLLM's native GDN spec path keeps a separate recurrent state snapshot for every draft token, so a step with T = 1 + num_speculative_tokens tokens writes T full head_v x head_k states per value head. Cost grows linearly in T.

ReplaySSM replaces the snapshots with one checkpoint plus a small circular buffer of the per-step inputs. For GDN the buffer holds (d, k, g), where d = beta * (v - S k) is the per-token delta-rule update and S is the decayed state before the rank-one update. The spec kernel computes these updates for the whole window as D_spec and writes them to d_cache. The name d is deliberate: GDN prefill already uses U for its chunk-transformed value tensor in the WY representation. Caching d rather than v removes the serial state dependency during replay. This changes three things:

• Rollback is a cursor move.
• Every draft output is computed from the same checkpoint.
• The full state is materialized only on a periodic flush.

What this PR adds

Everything is behind the flag added in #49847. With ReplaySSM disabled, the GDN decode path is unchanged.

• The GDN cached spec kernel gdn_replayssm_spec_decode, its device-side cursor helpers commit_gdn_replayssm_spec and reset_gdn_replayssm_spec_cursors, and the swept Blackwell launch configs.
• Integration. The Qwen3.5 GDN mixer dispatches to the kernel under --use-replayssm-spec, the GDN metadata builder owns the block-keyed ring cursors, and the state-shape and state-dtype calculators size the d/k/g ring on the paged state. Qwen3.5 declares SupportsReplaySSM. The conv path reuses vLLM's causal_conv1d_update unchanged.
• One benchmark change: the spec throughput benchmark routes GDN prefill through Triton instead of FlashInfer.

Design notes

The flush rule, cache sizing, on-device commit, and fixed verify-plus-flush launch pair are inherited from #49847. Three details are GDN-specific.

Paged cache layout

The ring and checkpoint live inside the paged KV block, appended as d_cache, k_cache, and g_cache on the baseline (conv, ssm) state for five tensors total. d_cache stores the delta-rule update defined above. The d and k caches use fp16 when activations are bf16, which reconstructs the state more accurately at the same width. g is always fp32 because it drives the cumulative decay replay.

Only the small per-block cursors live outside the page: write_pos, cache_base, and is_flush, all shaped (num_gpu_blocks,) and indexed by physical block id.

Routing every post-prefill row

Under the flag, the builder marks every post-prefill row as a spec row, so a draft-less row runs as a T=1 window. num_decode_draft_tokens_cpu cannot drive that mask: it is stale on steps where the drafter proposed nothing and is -1 for decode rows whose drafts were dropped. Sending either row to the baseline decode path would read the checkpoint page, which lags the committed ring.

Prefill chunks and captured graphs

A final single-token prompt chunk stays on the prefill path. This differs from GDN standard decode, where leaving it as a prefill lets a shape-uniform batch replay a decode graph against a stale cursor.

That batch is not uniform on the spec path. _is_uniform_decode requires both max_num_scheduled_tokens == 1 + num_speculative_tokens and num_tokens == max * num_reqs. A draft-less step fails the first condition because every row is one token wide, and a mixed batch fails the second because the prompt chunk is shorter than the window. Standard decode has uniform_decode_query_len == 1, so the same batch is uniform there. A builder test covers the classification.

Results

Qwen3.5-122B-A10B-NVFP4 on one B300, TP1, fp32 SSM state, B=16, in-checkpoint MTP drafter.

Kernel-level

Isolated per-layer-step CUDA-graph latency on B300. The baseline is fused_sigmoid_gating_delta_rule_update; ReplaySSM is the full step (verify plus flush) amortized over the exact steady-state flush cadence f = a/(B-T+a), at a mean acceptance of 2 tokens for T=2 and 3 for T=4,6,8. At B=16:

bs	T=2	T=4	T=6	T=8
128	1.79x	2.86x	3.30x	3.86x
256	2.02x	3.05x	3.49x	4.22x
512	2.16x	3.24x	3.56x	4.46x

The ratio grows with T because the baseline writes T states per step while ReplaySSM's traffic is flat in T.

End-to-end throughput

GSM8K, greedy, ignore_eos, 256 output tokens, max_model_len 2048, CUDA graphs on. Produced with benchmarks/replayssm/e2e_spec_decode_throughput.py.

ReplaySSM spec over native spec:

bs	T=2	T=4	T=6	T=8
1	1.00x	1.12x	0.87x	0.96x
32	1.05x	1.22x	1.17x	1.12x
128	1.35x	1.43x	1.42x	1.84x
256	1.36x	1.79x	1.99x	2.09x
512	1.73x	2.10x	OOM	OOM

ReplaySSM spec over standard (no-spec) decode:

bs	T=2	T=4	T=6	T=8
1	1.37x	1.97x	1.76x	1.77x
32	1.44x	2.08x	2.07x	1.94x
128	1.56x	1.87x	1.86x	1.83x
256	1.45x	1.84x	1.87x	1.76x
512	1.53x	1.90x	OOM	OOM

Compared with native spec, the gain grows with batch size, where its T state writes per step cost the most. Compared with standard decode, ReplaySSM is 1.4x to 2.1x faster at every batch size. At bs512, native spec reaches 9,223 tok/s at T=2 and 9,395 tok/s at T=4, below standard decode at 10,398 tok/s. ReplaySSM spec reaches 19,707 tok/s.

The two bs512 runs at T=6 and T=8 fail to start. The failure happens in _init_minimal_kv_cache_for_profiling, which allocates a temporary KV cache before vLLM determines the size of the real one. This isn't a ReplaySSM memory regression. The temporary cache is block-based, so it charges the larger ReplaySSM page size without accounting for the reduced per-request footprint when num_speculative_blocks = 0.

Acceptance length

At bs256:

T	native spec	ReplaySSM
2	1.95	1.95
4	3.56	3.57
6	4.73	4.76
8	5.52	5.53

Serving capacity

Native spec holds 1 + num_speculative_tokens recurrent states per request for rollback. ReplaySSM holds one checkpoint plus the input ring. Maximum decode concurrency at max_model_len 2048, (B, T) = (16, 4), one B300:

native spec	ReplaySSM	ratio
233.2	656.5	2.81x

These are worst-case caps at the full request length.

Accuracy

Official model-card presets with paired seeds. Qwen3.5-122B-A10B-NVFP4 on one B300 with CUDA graphs (FP32 state, d/k cache in FP16 instead of BF16 for precision purpose):

dataset	standard	native spec	ReplaySSM spec
gsm8k	0.9606	0.9621	0.9575
mmlu_pro	0.8590	0.8590	0.8550
gpqa	0.8510	0.8460	0.8422
hmmt	0.8729	0.8708	0.8771

Every ReplaySSM result is within 0.9 percentage points of the baseline, and degenerate-loop counts stay within the baseline range. Qwen3.5-35B-A3B on 2xH100 TP2 shows the same parity. Native spec scores 85.8% on MMLU-Pro against a model-card value of 85.3%, and 87.9% on HMMT against 89.0%.

Acceptance length is only comparable between runs with identical settings. With the model-card presets, native spec and ReplaySSM both average 2.6 to 2.8 accepted tokens; the greedy throughput runs average about 3.4 on the same model.

Tests

Per-step kernel correctness is covered by tests/kernels/test_replayssm_spec_decode_gdn.py. A T-window verify step is checked against decoding the same window one token at a time, across three state and activation precision combinations, the production geometries (small, Qwen3.5-4B, Qwen3.5-122B), B in {16, 32}, and T in {2, 4, 6, 8}. It also covers rollback across a flush and continuous batching with per-row cursors and null-padded slots. The reference is the upstream fused_recurrent_gated_delta_rule_packed_decode, run on a clone from the checkpoint through the history and then the window. The committed history is built with draft-less T=1 spec steps. One case corrupts a cached history key and asserts that the output changes.

Metadata builder unit tests are in tests/v1/attention/test_gdn_spec_metadata_builder.py. They cover commit and flush cursor updates, the early-flush boundary, rejected drafts, recycled blocks, resume after preemption, mixed batches, draft-less rows, final prefill chunks, persistent block-keyed cursor buffers under graph capture, and matching page sizes from the layer and config paths. The last check guards against silent truncation when those paths are zipped.

Engine-level parity lives in tests/v1/e2e/test_replayssm_decode.py, which compares native spec against ReplaySSM spec with greedy decode and check_logprobs_close on Qwen3.5-4B.

Results: the GDN kernel file is 67 passed (131 seconds on one A100), and the GDN builder file is 13 passed. Together with the Mamba2 suites inherited from #49847, the combined run is 273 passed. Engine parity on Qwen3.5-4B passes on one H100.

pytest tests/kernels/test_replayssm_spec_decode_gdn.py \
       tests/v1/attention/test_gdn_spec_metadata_builder.py \
       tests/v1/e2e/test_replayssm_decode.py -q

Duplicate check

No open PR addresses ReplaySSM GDN speculative decode. This PR is split from the large draft PR #47576, which stays open as the full design. The related prior ReplaySSM PRs we are aware of cover standard decode only:

• #43102 and #43439, earlier PRs from our collaborator on the same ReplaySSM effort. This is the newer, more complete version.
• #46239, opened by another contributor based on our original fork. We have since made substantial changes, intend to land the feature ourselves, and have notified the author.

Out of scope

• Prefix caching. Rejected by the validator added in [Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 speculative decode #49847.
• P/D disaggregation and KV connectors. Rejected by the validator.
• Stochastic rounding. Rejected by the validator.
• Low-precision state. The code handles fp16 and bf16 state, but only fp32 state is evaluated here.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Johnny-Liou.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[F21HGG]

Hi @Johnny-Liou, thanks for working on this. We independently explored the same GDN speculative-state traffic bottleneck, and I would like to coordinate rather than open a duplicate PR.

We have a complementary exact-order design:

• Keep the baseline sequential FP16 GDN recurrence unchanged during verification.
• Store only the per-token a/b/k/v inputs in stable scratch.
• After target acceptance is known, replay and commit only the accepted-length state to the canonical GDN/Mamba cache.
• No periodic checkpoint flush, hidden dtype promotion, D2H synchronization, or change to token acceptance/sampling semantics.

The main difference from #49887 is the intended contract. #49887 currently requires mamba_cache_mode=none and lists prefix caching and low-precision state evaluation as out of scope. Our path is specifically for mamba_cache_mode=align, prefix caching, static MTP K1-K3, TP1, and FP16 recurrent state.

We first ported the current ReplaySSM GDN design and tested it with the public Qwen/Qwen3.5-4B model under deterministic FP16-state gates. K1 and K3 both first diverged from the native sequential baseline at output token 8. The direct kernel tests indicated that the chunked replay transform changes arithmetic order, and FP16 checkpoint rounding makes that difference observable. This motivated the exact-order alternative instead of accepting tolerance or promoting the state dtype.

Preliminary evidence for the exact-order path:

• Same-process shadow checks covered K0/K1/K3, zero/partial/full acceptance, and draft-less rows; per-step GDN outputs and accepted FP16 states were bit-exact against the sequential baseline.
• A controlled two-replica RTX PRO 5000 A-B-B-A serving test on Qwen3.5-4B showed per-GPU throughput +4.29%, total token throughput +4.38%, average latency -3.25%, and P95 latency -2.98%, with zero failed requests and zero preemptions.
• The measured recurrent-state traffic reduction was 38.1%.
• The trade-off is 1.137 GiB of stable scratch per GPU, reducing logical KV capacity by 2.23% in that configuration.

Would you and the maintainers prefer this as:

1. a small stacked/follow-up PR adding an FP16-exact align-mode path after [Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 speculative decode #49847/[Kernel] ReplaySSM: cache SSM inputs for faster Gated DeltaNet speculative decode #49887, or
2. tests and an exact-order backend contributed directly to the ReplaySSM series?

I am happy to rebase the minimal implementation onto current main, keep it default-off/fail-closed initially, and add upstream tests for prefix hits, zero accept, draft-less rows, final single-token prompt chunks, preemption/abort, row compaction, and CUDA Graph capture/replay.

[hnover-nv]

@F21HGG regarding #49887 (comment)

First, I'll say we are working on adding cache mode align to replay, just not in that PR. I don't think adding align support is any easier or harder in a replay implementation vs what I think you're proposing (see next); both require a kernel that converts an older state and an input buffer to some intermediate state when old state + # accepted tokens passes an alignment boundary.

If I'm following you correctly, I think you're proposing what we've been calling "recover ssm", which was implemented for Kimi K3 in #51855, and which there's some in-progress work to speed up, although our long-term goal is to move to a replay model. There are 3 paradigms for MTP when you don't write every state:

• "recover ssm", where after acceptance you write the "correct" output state. Any buffered inputs are ephemeral. This has worse performance because you have to do that state computation and write at the end, although the semantics are perhaps cleaner.
• "every step replay" where the state stored is the state two steps back, along with the input buffers. As part of the per-layer forward, you use # of accepted tokens to advance the old state to the last-step accepted state, write that out, and then generate the new output tokens. here the buffered inputs and other metadata like # of accepted tokens are durable between steps. Note the vLLM MTP default where we write every possible state is sort of like this, as vLLM doesn't copy the winning state from all the possible states over at step end, but just remembers which of the (draft_len + 1) output states was the winner and uses that as the input to the next step
• "checkpointing replay" aka "replay ssm", where the checkpoint is arbitrarily far back and we only flush when the buffer is near full

Then there is an orthogonal algorithmic question to the kernel internals for any of these, whether the math is done recurrently or via a prefill-ish UT / WY transform, and if it's the prefill-ish one when do we use a precision other than fp32.

In our experiments, even for GDN we have not seen an issue with accuracy using replay, although we were using the Flashinfer GDN replay kernels that will be dropped in on a follow-up PR.

It's true you can get some divergence, but that's true just by switching between recurrent MTP and recurrent STP, because STP will quantize to fp16 every token whereas recurrent MTP carries state w/in step in fp32, and replay-ish methods accumulate in fp32 state.

ReplaySSM should actually be more accurate with fp16 state than what you propose, because state writing and quantization happens only when the buffer is full, vs every step. It's true that some math happens in 16b for tensor cores, but prefill does that as well, and prefill accuracy is fine.

[F21HGG]

Thanks — yes, this is a GDN "recover SSM" path.

I suggested contributing it to this series rather than opening a competing PR because the integration surface is the same: Qwen3.5 dispatch, accepted-token metadata, cache lifecycle, preemption/compaction, CUDA Graph handling, and cache-mode support. The independent part is the commit policy/kernel. It can live beside ReplaySSM as an exact policy for native-MTP FP16 semantics and align/prefix-cache deployments.

On 5b96cdbc, K1 and K3 first diverged from native at token 8 with FP16 state; this path remained bit-exact for GDN output, selected token, chosen logprob, and accepted state. In our A-B-B-A serving test it reduced recurrent-state traffic by 38.1% and improved QPH/GPU by 4.29%, at the cost of 1.137 GiB/GPU scratch and 2.23% KV capacity.

The implementation is rebased onto current main as one isolated commit:

• commit 02f53f240c
• compare against main

RTX PRO 5000 validation passed 11 kernel tests, 13 metadata tests, and 5 runner lifecycle tests. Would you be open to taking this as an isolated commit set in the ReplaySSM series, sharing the existing metadata and lifecycle infrastructure?

[askliar]

@F21HGG thanks for starting the discussion!
Few questions:

1. Any particular reason you want implementation to be bit-exact?
2. We are working on having two paths enabled in vLLM: Johnny's path (with this PR and [Kernel] ReplaySSM: cache SSM inputs for faster Gated DeltaNet standard decode #48792) + FlashInfer path (PRs are to follow in the next weeks). FlashInfer PR will include prefix caching. Those should address most of the concerns you have and have not "RecoverSSM" but "ReplaySSM" directly. This way we can avoid the overhead of agreeing on everything and spending extra time and effort on that. Would it be acceptable as a solution?

[F21HGG]

Bit-exactness is not a hard requirement for me; it was a conservative correctness gate. If the planned ReplaySSM + FlashInfer path supports Qwen3.5 speculative decode with mamba_cache_mode=align and prefix caching, that addresses my use case. I’m fine using the unified path rather than adding a separate implementation.

[benchislett]

+1, let's focus our effort on trying to get the implementation right the first time. To me, that's ReplaySSM not RecoverSSM. Prefix caching is a hard requirement and we will need to have a concrete plan for this if we do not intend to have support after landing the first PR in the sequence. It sounds like we do.

[erdholion]

Production datapoint in support of "prefix caching is a hard requirement" — from a config that runs hybrid GDN + MTP + prefix caching together today, since numbers for that combination seem scarce in this thread.

Config: Qwen3.8-27B (hybrid GDN, NVFP4/ModelOpt), vLLM v0.27.1, RTX 5090 (SM_120, 32 GB), FLASHINFER, fp8 e4m3 KV, MTP num_speculative_tokens=3, prefix caching enabled, max_model_len 187,872, serving a multi-agent coding-harness gateway (long shared system prompts, heavy multi-turn reuse).

Measured over a 5h serving window today:

• prefix-cache hit ratio 49.8% (5,668,800 / 11,389,975 queries)
• MTP acceptance 81.8% overall = 2.45 of 3 drafted tokens accepted per proposal (per-position 93.0% / 84.8% / 67.5%)

So both features carry their weight simultaneously on this workload — losing prefix caching to land spec-decode improvements would cost us roughly half our prefill on cached turns, which is why the hard-requirement stance matters to downstream users like us.

The known interaction tax we eat today is #54360 (spec decode enlarging the effective block-hash unit, so cache hits quantize coarsely under MTP). If the ReplaySSM sequence changes that behavior in either direction, we are happy to A/B candidate PRs on this box — 5090/SM_120 with fp8 KV is a config the CI matrix does not cover, and we can run identical multi-hour gateway load against both arms.