[Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv

[jschmied]

Purpose

Speculative decoding on Gated DeltaNet today snapshots the full recurrent state after every draft position, and
reserves one Mamba block per draft position. This PR brings the RecoverSSM protocol that Kimi-K3 KDA already uses
to the Qwen GDN layers of Qwen4Exp (Qwen3.8-Flash-Next), plus its PLE dilated short conv:

• Verify runs from a read-only checkpoint and stores a small per-token replay record
  [correction (V), normalized key (K), gate (1)] in fp32.
• Commit (after sampling) replays only the accepted tokens and writes the recurrent state once, plus the
  block-boundary state in align mode. The PLE short conv compacts its accepted window with the same plan.
• num_speculative_blocks = 0: no per-draft Mamba blocks.

Enabled like KDA's: --use-replayssm together with speculative decoding. Qwen4Exp (CausalLM, ConditionalGeneration,
MTP) declares SupportsReplaySSM; without spec decoding the flag is rejected on Qwen4Exp (there is no plain ReplaySSM
path for its GDN/PLE layers). Requires the V2 model runner for align mode, the Triton mamba backend and PP=1 (the same
checks as KDA). The GDN and PLE RecoverSSM builders declare AttentionCGSupport.UNIFORM_BATCH, so the default
FULL_AND_PIECEWISE captures FULL graphs for uniform spec-decode batches (and PIECEWISE for mixed ones). Graph padding
rows get a null state slot and a zero-length window, so the verify writes nothing for them. With breakable CUDA graphs
(auto-enabled on this setup; they split at attention with an empty splitting_ops) a PIECEWISE fallback used to drop
to NONE; it now resolves to PIECEWISE there.

Note: under FULL decode graphs resolve_cudagraph_mode_and_sizes requires max_num_seqs <= num_blocks. The replay
record makes the unified page about 6 % larger, so a pool that just fits the default max_num_seqs under native MTP
can fall short with --use-replayssm (seen on one H200 at TP1, 983 Mamba blocks); lower --max-num-seqs there.

The larger page also widens the attention block that hybrid models align to the Mamba page (Qwen3.8-Flash-Next on
one GB10: 1,600 → 1,664 tokens). Prefix reuse counts whole blocks, so a repeated prompt can reuse up to one block
less (a 9,698-token repeat reused 6,656 instead of 8,000 tokens in @k3dani's run); long repeats are unaffected.
The decode gain comes from dropping the per-draft state write-back, so it grows with the draft length: measured at
MTP K=3 and K=5 here; at K=2 with PIECEWISE graphs @k3dani measured c=1 decode 1.9–4 % slower.

Commits:

1. [Mamba] Share RecoverSSM commit kernels and helpers across models: the commit-plan and conv-compaction kernels
   move from models/kimi_k3 to model_executor/layers/mamba/ops/recoverssm.py (Kimi keeps its private names as
   aliases; kernel source unchanged); mamba/recoverssm_utils.py; generic request-mapping helpers in
   v1/attention/backends/recoverssm_metadata.py; CacheConfig.use_kda_recoverssm → use_recoverssm. No functional
   change for Kimi-K3.
2. [Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv: GDN verify/commit kernels
   (mamba/gdn/recoverssm_gdn.py), metadata builders (gdn_recoverssm.py, ple_recoverssm.py), the layer and model
   integration. The verify reads strided q/k/v views of the conv output and writes straight into the layer output.
   The generic Mamba metadata builder no longer allocates the Mamba2 ReplaySSM ring-buffer workspaces when
   --use-replayssm selects RecoverSSM (the PLE short conv uses that builder).
3. [Mamba] RecoverSSM align: keep a boundary-ending state in its block: a window whose accepted tokens end exactly on
   a block boundary commits into the block holding its last token ((n - 1) // block), not the next, possibly
   unallocated column.
4. Three fixes by @lucifer1004 ([Spec Decode][Qwen4Exp] RecoverSSM fixes: profiling OOM, flag guard, faster align-mode commit jschmied/vllm#2): release the recorded RecoverSSM step in _teardown_profiling_state
   (its commit contexts kept the profiling KV cache alive; ~26 GiB on H200 and a startup OOM); read the RecoverSSM flag
   with getattr in _forward_core; one fold per program in the align-mode commit (290 → 197 µs on H200).
5. [Spec Decode][Qwen4Exp] RecoverSSM: FULL cudagraphs for uniform spec-decode batches (above).
6. [Bugfix][Mamba] RecoverSSM align commit: fold final states before boundary states: the split commit ran both folds
   in one launch, and a boundary state can land in the window's source slot, so a program could read a source tile
   another had overwritten. Final folds now run first, boundary folds second (skipping windows whose boundary slot is
   their final slot).

Related: #49887 and #56466 (ReplaySSM for GDN) take the ring-buffer route; this PR reuses the merged RecoverSSM
protocol instead (discussed in #56466 (comment)).

Test Plan

• tests/kernels/mamba/test_recoverssm_gdn.py: verify bit-identical to the native fused kernel (checkpoint untouched);
  commit bit-identical to the native per-token states (accepted 1–4); align boundaries (8 cases, incl.
  exact-boundary and crossing);
  mixed prefill/decode mapping and non-contiguous request rows; extreme gates × 3 seeds; strided q/k/v views + out=
  (uniform and ragged windows 4/1/3) and the layer's view/fallback decision; BF16 recurrent state (verify exact; commit
  rounds once: within one BF16 ulp of the FP32 reference and no further from it than the native path);
  malformed-input rejection.
• tests/kernels/mamba/test_recoverssm_ple.py: spec conv + commit vs the stock token-by-token decode (accepted 1, 2,
  full window; 4 align cases; non-contiguous request rows), bit-identical states and outputs.
• tests/config/test_recoverssm_config.py: selection through the real validator (target and the MTP drafter's derived
  config), runtime-check rejections, no spec decoding, unsupported architectures, the FULL-graph declaration.
• tests/kernels/mamba/test_recoverssm_gdn.py also pads a batch the way graph capture does and checks that padding rows
  change no state; tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py checks the teardown releases the
  recorded step.
• Server: Qwen3.8-Flash-Next (NVFP4 experts, FP8 GDN), one DGX Spark GB10 (sm_121), MTP K=3, align mode + prefix
  caching, 2 starts per arm: output hashes, cache-hit replay, decode and agent-loop timing.

Test Result

Tests on GB10 (sm_121): GDN + PLE kernel tests 72 passed, config tests 18 passed. ruff, mypy (tools/pre_commit/mypy.py 3.12), typos: clean.

Server, native GDN spec decode (align) vs RecoverSSM (align), 2 starts per arm
(data):

	native (align)	RecoverSSM (align)	Δ
c=1 decode, ms per verify cycle	55.71–56.65	54.37–54.38	−1.3…−2.3 ms
c=4 decode, ms per verify cycle	105.07–105.41	99.59–99.91	−5.2 %
agent loop (8 dependent turns), s/turn	1.31	1.10	−16 %
prefix-cache hit rate	77.4 %	82.1 %	+4.7 pp
KV tokens in the same 4 GiB	75,678	103,953	+37 %

• Node trace (§4u):
  the per-draft state write-back is gone (GDN out_proj 101 → 74 µs, no dirty-L2 aftershock); the commit runs at the
  DRAM floor (224 GB/s).
• The zero-copy verify adds −1.2…−1.6 % at c=1 with identical outputs
  (§5g).
• Output: cache-hit replay reproduces the first pass 8/8 in every run. Since 514102a1db (suggested by @ArtyomITA)
  the verify step records the decay it multiplied by, exp(g), and the commit replays the accepted tokens forward with
  the same update, h = fma(c, k, h * decay), so the committed state equals the native per-token state bit for bit:
  the FP32 kernel tests now assert torch.equal (the earlier closed form fails 28 of those 62 cases on a GB10).
  @ArtyomITA measured greedy tokens and top-20 logprobs identical to a run without --use-replayssm on sm_120
  (8 prompts of 120–245k tokens, 1 and 4 users, matched block size and a deterministic top-k, [Kernel] persistent_topk: deterministic select, faster than the exact-topk workaround #55122); not repeated
  here. Before that change, text differed from the native path at the size of a reduction-order change (first greedy
  divergence at a median of 26 tokens).
• This PR's final shape on the same box (one start each): --use-replayssm with explicit PIECEWISE: c=1 21.22 ms/tok,
  output hashes and cache-hit replay identical to the validated reference; default FULL_AND_PIECEWISE: resolves to
  PIECEWISE, c=1 21.31 ms/tok, c=1 and c=4 hashes identical to the reference, cache-hit replay 8/8.
• FULL graphs (GB10, MTP K=5, 2 starts per arm): FULL_AND_PIECEWISE vs PIECEWISE code c=1 14.72 / 14.70 vs 15.00 / 14.88
  ms/tok (−1.1…−2.1 %), c=4 level, identical greedy text. On 2× RTX PRO 6000 (TP2) and one H200 (TP1) @lucifer1004
  measured this PR plus the FULL graphs and a fused CUDA verify at parity with native MTP3 at C1, −1…−2 % at C2–C16,
  and up to +17 % output tok/s at high concurrency (RTX C32, H200 C64), with H200 preemptions at C64 going from 188 to 0
  ([Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv #58863 (comment)).
• Head 3388ba1a25 on GB10: the five test files above 121 passed; test_recoverssm_gdn.py 20 of 20 repeated runs
  pass (the split commit without the ordering fix failed 9 of 12); greedy hashes identical to the reference.

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Developed with AI assistance (Claude Code); all changes reviewed by the author.

[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.

[TSUMUGI-XE]

Opened the two follow-ups as drafts against main, stacked on this PR: #59982 (XPU routing) and #59983 (PP > 1). Happy to retarget them to your branch instead if you'd rather carry them here.

— TSUMUGI-XE

[TSUMUGI-XE]

Question: would you take a change that makes the RecoverSSM commit cheaper to launch, either by grouping it across KV cache groups or by putting it into the step's graph? Here or as a follow-up? Details below (we can only measure on XPU).

On XPU (2× Arc Pro B70, Qwen3.8-Flash-Next, PP=2 with our follow-up #59983, MTP K=1) the commit shows up in single-stream latency on the last PP rank, and most of its cost is per-group launch overhead:

• commit_step calls commit_recoverssm_state once per KV cache group. Here each pipeline stage has 3 GDN groups, so the commit is 3 × 3 Triton launches (plan, conv, state) per step, about 0.5 ms of host time on the last rank. In a micro-benchmark (synthetic tensors, one card), three 6-layer commit contexts take 378 µs of host time per step versus 125 µs for one 18-layer context.
• To check whether this is on the critical path, we added a busy-wait at the start of commit_step: 1.5 ms on the last rank added 0.9 ms/token, while 1.5 ms on the first rank added nothing (it is hidden behind the last rank's forward).
• The per-group loop runs the same way with PP=1, where there is no other stage to hide it behind, but we haven't measured PP=1.
• Locally we replay the last rank's commit from an XPU graph, which recovered about 0.4 ms/token. That points to launch overhead rather than the commit's GPU work. It relies on XPU graph capture outside the model's graph, so we are not proposing it as is.

So two options we can see: (1) a grouped commit, one set of launches covering all groups of a rank (e.g. with a [group, batch] state-index table), or (2) moving the commit into the step's captured graph. On CUDA the per-launch host cost is probably smaller, so this may matter mainly for XPU. We can measure either option on XPU if that helps.

Written with AI assistance (Claude Code).

— TSUMUGI-XE

[jschmied]

@TSUMUGI-XE thanks for measuring it. Yes, please send it as a follow-up PR; this one is big enough already.

Of the two options, I'd prefer the grouped commit: one set of plan/conv/state launches over a [group, batch] state-index table. It helps every backend in the same way, and it keeps commit_step where it is now, outside graph capture. commit_step runs after sampling, once num_sampled is known, and it returns early during capture.

On the GB10 (PP=1, CUDA graphs) I can measure the host time of the per-group loop before and after on your branch.

(Written with AI assistance.)

[mergify]

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

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

[TSUMUGI-XE]

@jschmied the grouped commit is up as #61017, stacked on this PR (only the last commit is new). On XPU, for 3 KV cache groups × 6 GDN layers, it cut the commit's host time from about 400 to 140 µs per step. With PP=2 it lowered the single-stream decode slope by 0.19 ms/token. Your GB10 numbers (PP=1) would show how much it helps on CUDA.

Written with AI assistance (Claude Code).

— TSUMUGI-XE

[mergify]

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

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

[Yuchi017]

Thanks for this PR. We ported the same verify/commit protocol to the Qwen3.5 GDN layers on top of v0.31.0 locally. The GDN layer code is shared with Qwen4Exp; only SupportsReplaySSM on the Qwen3.5 classes differs. Below are numbers on the public checkpoint, plus one observation about the replay record living in the Mamba page.

Setup:

• 1× L40S, Qwen/Qwen3.5-4B with its native MTP head, BF16, MRV2.
• --mamba-cache-mode align --prefix-match-unit 16, TRITON_ATTN, prefix caching on.
• Spec-Bench (480 prompts), 256 output tokens, temperature 0.

Output tok/s at c=32, stock spec-decode kernels → RecoverSSM:

MTP K	Stock	RecoverSSM	Δ
3 (mean of 3 runs)	2607	2871	+10.1% (runs: +16 / +5 / +9%)
5 (mean of 3 runs)	2529	2723	+7.6% (runs: +10 / +9 / +4%)
7 (1 run)	2436	2693	+11%
1 (1 run)	2062	2010	−3%

At c=1 the two are within ±1% for every K. At c=8 we measured K=3 +8%, K=5 +2%, K=7 +3% (single runs).

Accept length is unchanged in every run (e.g. K=5: 3.941 vs 3.939; per-position .856/.694/.558/.457/.376 vs .856/.693/.558/.456/.375).

The gain grows with batch size and K. At c=1, and at K=1, the commit's fixed cost cancels the saving, which matches the K=2 / c=1 slowdown you mention. It may be worth noting in the docs that RecoverSSM pays off at K ≥ 3 under load.

Correctness. Outputs aren't bit-identical to the stock kernels, as expected when float math is reordered. We checked with greedy decoding + top-5 logprobs:

• Every first divergence is at a near-tie: the stock run's top-1 − top-2 gap there is ≤ 0.875 nats (median 0.25), versus a median of ~11 over all positions.
• This also holds when a follow-up turn's prefix hits blocks whose boundary state was written by the commit during decode.

Observation: record in the page vs. prefix reuse. You note the record makes the Mamba page larger, which widens the attention block. In our setup it grew the block from 576 to 656 tokens. On a workload where requests share a ~1.9k-token system prompt ending mid-block, shared-prefix hits land on block multiples:

• Every request reused 1,312 tokens (2 × 656) instead of 1,728 (3 × 576).
• On a sequential probe of 40 requests, the prefix hit rate dropped from 0.489 to 0.381.

The record only lives from a step's verify to the same step's commit, so we moved it out of the KV pool into a per-layer buffer ([max_num_seqs, HV, K+1, V] fp32 correction + [max_num_seqs, HV, K+1, K+1] key/gate). The Mamba page is then the native size, and the block and hit rate match stock. Our prototype indexes rows by spec-decode position, which isn't safe once PP keeps several steps in flight (#59983); an upstream version would key on the request-state slot.

Proposal: two small follow-ups stacked on this PR:

1. Declare SupportsReplaySSM for Qwen3.5 / Qwen3-Next (+ MTP), with tests on the public checkpoint.
2. Keep the replay record out of the Mamba page (request-slot indexed).

Happy to open these once this lands, or adjust if you'd rather fold either into this PR. Does that direction sound OK?

[Yuchi017]

Opened #61309 as a draft stacked on this PR to enable Qwen3.5 / Qwen3-Next (proposal 1 above). It also opts out the subclasses that would otherwise inherit SupportsReplaySSM from the Qwen3.5 base (Intern-S2, and the AMD Qwen4Exp model).

[mergify]

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

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