[Bugfix] Preserve Mamba state for padded prompt tails

[natsala13]

Purpose

Fix recurrent-state corruption in hybrid Mamba models when speculative decoding pads the last prompt token to a uniform K + 1 query length.

This is most visible in prefill/decode-disaggregated serving:

1. The prefill node transfers the Mamba recurrent state after prompt token N - 1.
2. The decode node must process the final prompt token N to produce h(N).
3. If another speculative request is already running, the scheduler pads this one-token prompt tail with K placeholder positions so every row has query length K + 1 and can use the same FULL CUDA graph.
4. The Mamba backend currently classifies that row as a prefill. Prefill persists only the state after the final query position, so it saves h(N + K placeholders) instead of h(N).
5. The scheduler discards the placeholders from token accounting, but the prefill state path cannot roll the recurrent state back. The next real target token therefore reads a corrupted convolution/SSM state.

The requests do not need to share a prefix. They only need to enter the same scheduler step. Both may have MTP enabled: one request is already in steady-state speculative decode, while the newly resumed request still has one real prompt token left.

The trigger is:

• a hybrid Mamba model;
• speculative decoding with K > 0;
• a request with prior recurrent state and exactly one prompt token remaining; and
• another running speculative request, which activates uniform K + 1 padding.

Disaggregated serving makes this common because the NIXL Mamba handoff deliberately transfers h(N - 1) and leaves token N for the decode node. The same state can also arise after a local prefix-cache hit, preemption/replay, or another short prompt extension, so the fix is not NIXL-specific.

Fix

Pass the scheduler's draft-count tag into Mamba metadata construction. When a row is:

• tagged as speculative padding,
• shaped as one real token plus K placeholders, and
• extending an existing recurrent state,

route it through Mamba's existing speculative-decode state layout even though the real token is still part of the prompt.

That layout is already transactional: it keeps h(N) in the running state slot and writes later speculative states to scratch slots. Normal acceptance/rollback therefore discards the placeholder states and retains exactly h(N). Genuine multi-token prefills and first prompt chunks remain on the prefill path.

This preserves the original performance optimization: the scheduler still emits a uniform K + 1 row, so the FULL CUDA-graph batch shape is unchanged. The change only selects the state-management path that can safely roll back speculative positions.

Reproduction

Deploy a hybrid Mamba model with prefill/decode disaggregation and NIXL KV/state transfer. Use separate prefill and decode workers, and enable MTP with five speculative tokens on both roles so their cache layouts remain compatible. Our reproducing deployment used two TP4 prefill workers and two TP4 decode workers, but the essential condition is that a decode worker handles more than one request concurrently.

Run SWE-bench Verified against the P/D router with enough concurrency for a newly transferred request to join a scheduler step that already contains a speculative decode request. The transferred request arrives with recurrent state through prompt token N - 1; the decode worker must process token N. Stock vLLM pads that one-token prompt tail with five placeholders to match the existing six-position speculative row.

On an affected vLLM 0.28 deployment, the padded row is sent through the Mamba prefill state path. The next decoding iteration then reads state after the placeholders instead of state after token N. In the full 500-instance, three-repeat SWE run, this appears as invalid byte-piece output, Unicode replacement characters, leaked literal tool-call XML, missing structured tool calls, and a large resolved-accuracy drop with MTP5. A matched MTP-off run using the same checkpoint, serving topology, sampling settings, and benchmark provides the control.

Apply this PR and repeat the same MTP5 workload without changing the checkpoint or serving parameters. The padded batch shape remains unchanged, but the prompt-tail row uses transactional speculative state handling; the corrupted-output phenotype disappears and SWE accuracy returns to the matched MTP-off level.

Why this is not duplicate work

No open PR found in the required duplicate-work search routes scheduler-tagged, padded prompt-tail rows through Mamba's transactional speculative state path.

• [Bugfix][PD] Pad resumed speculative decode requests #55126 changes the scheduler so a resumed speculative request can retain the uniform verifier shape even when its data-parallel rank has no already-running request. It does not change Mamba state handling. The changes are complementary: [Bugfix][PD] Pad resumed speculative decode requests #55126 can broaden where prompt-tail padding occurs, while this PR makes that padding state-safe for hybrid Mamba models.
• [Bugfix][Kimi-K3] Do not classify a stateless first chunk as a decode #51483 distinguishes stateless first-token KDA chunks from stateful one-token updates. This PR instead handles a row that already owns recurrent state and contains one real prompt token plus K scheduler placeholders, in the shared Mamba metadata path.
• [Bugfix][Spec Decode] Cache the Mamba state at the block-grid position of EAGLE resume #53945 addresses EAGLE/MTP prefix-cache checkpoint placement in Mamba align mode. It does not address placeholder positions being persisted through the prefill state path.

Validation commands and results

The rebased source-only branch passes the following checks:

UV_CACHE_DIR=/tmp/vllm-pr-uv-cache uv run --no-project python -m py_compile \
  vllm/v1/attention/backends/mamba_attn.py \
  vllm/v1/attention/backends/mamba2_attn.py
git diff --check upstream/main...HEAD

A focused metadata regression in the deployment image passed all three cases: a stateful padded prompt tail becomes a speculative-decode row; an ordinary six-token prefill remains a prefill; and a stateless first prompt chunk remains a prefill.

An exact-image CUDA state probe compared three lanes from identical initial recurrent state. The padded speculative-decode lane retained the same accepted state and next-token result as the single-token decode lane. The unpatched padded-prefill lane persisted the placeholder-advanced convolution and SSM states and diverged on the next real token.

The bounded P/D MTP5 reproducer changed from 7/8 corrupted responses before the fix to 0/8 after it while retaining the uniform K + 1 batch shape. The full SWE-bench Verified qualification is reported below.

AI assistance

OpenAI Codex assisted with root-cause analysis, implementation, validation harnesses, experiment analysis, and PR preparation. I reviewed and understand every changed line and am responsible for the contribution and the results reported here.

Test Result

mean/resolved over all 1,500 SWE trajectories:

Cohort	vLLM state handling	MTP	Resolved	Accuracy
Original regression	Before fix	off	580 / 1,500	38.67%
Original regression	Before fix	5	57 / 1,500	3.80%
Post-fix qualification	Fixed image	off	573 / 1,500	38.20%
Post-fix qualification	Fixed image	5	576 / 1,500	38.40%

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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: c03927e7-bd14-47d7-9efa-a097b8a19d46

📥 Commits

Reviewing files that changed from the base of the PR and between 8f269a9 and d734c8f.

📒 Files selected for processing (2)

• vllm/v1/attention/backends/mamba2_attn.py
• vllm/v1/attention/backends/mamba_attn.py

🚧 Files skipped from review as they are similar to previous changes (2)

• vllm/v1/attention/backends/mamba2_attn.py
• vllm/v1/attention/backends/mamba_attn.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved Mamba attention handling during speculative decoding.
  • Correctly processes remote prompt tails padded with draft tokens through the appropriate decode and state-update flow.
  • Improved handling of mixed prefill and decode workloads when draft tokens are present.
  • Ensured attention state remains accurate across successive decoding steps.
  • Improved metadata consistency across Mamba-based attention backends for more reliable token processing.

Walkthrough

Mamba attention metadata builders now propagate optional draft-token counts. Prefill-to-decode classification uses prior Mamba state and supports padded speculative prompt tails.

Changes

Mamba draft-token routing

Layer / File(s)	Summary
Propagate draft-token metadata vllm/v1/attention/backends/mamba2_attn.py, vllm/v1/attention/backends/mamba_attn.py	Build methods accept optional CPU draft-token counts and pass them to common metadata computation.
Route stateful prompt tails vllm/v1/attention/backends/mamba_attn.py	Prefill-to-decode routing checks prior Mamba state and routes padded speculative prompt tails when their draft count and query length match.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to d734c

The change routes padded stateful Mamba prompt tails through transactional speculative decoding to preserve recurrent state. No concrete unresolved merge-blocking risk is recorded.

Suggested reviewers: njhill

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the main change: preserving Mamba state for padded prompt tails.
Description check	✅ Passed	The description directly explains the recurrent-state corruption, the Mamba metadata fix, affected scenarios, and validation results.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[natsala13]

@njhill @qianlihuang Could you please review this or route it to the appropriate Mamba reviewer? This fixes silent recurrent-state corruption when the K+1 padding introduced in #45237 is applied to a stateful Mamba prompt tail, while preserving the uniform CUDA-graph shape and its performance benefit. It is complementary to #55126: that PR broadens when padding occurs, while this PR makes the padding state-safe for Mamba. The bounded reproduction changed from 7/8 corrupted responses to 0/8, with full SWE MTP5 accuracy recovery. If the approach looks right, could you mark it ready or authorize CI?

[njhill]

Thanks @natsala13.

Did you try this with the v2 model runner (now the default)? I think it does not apply with that.

But I think this is a clean fix and probably good to have the additional protection here.

See also #50591 which was a fix in a different place. cc @sungsooha

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[natsala13]

Thanks, I had not yet run a matched MRV2 control. I see that #50591 found the unpatched MRV2 path clean. I’ll run the exact concurrent P/D reproducer with VLLM_USE_V2_MODEL_RUNNER=1 and verify both unpatched and patched behavior, then update the PR scope accordingly.

[natsala13]

MRV2 SWE-bench validation

Hi @njhill I ran the full 1,500-rollout SWE-bench Verified evaluation on upstream main, using Model Runner V2 and the same 2P/2D disaggregated configuration for every cohort.

MRV2 configuration	MTP	Resolved	Accuracy
Without this fix	5	234/1500	15.60%
With this fix	5	587/1500	39.13%
With this fix	Off	584/1500	38.93%

MRV2 mitigates the older issue described in #50591, but it does not fix this bug. Stock MRV2 still classifies a stateful prompt-tail row padded to K+1 as prefill, causing Mamba to commit the state after the placeholder positions rather than after the real token.

This fix routes that row through Mamba's existing transactional speculative-decode state layout. It preserves the uniform CUDA-graph shape while committing only the accepted state. The fixed MTP5 result matches the MTP-off control and eliminates the observed empty responses, unparsed tool XML, reasoning-only outputs, and random byte corruption.

[njhill]

Thanks @natsala13

[github-actions]

✅ @natsala13, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[njhill]

/ci run

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[github-actions]

✅ Triggered Buildkite CI #87334 for commit d734c8f3f7f8.

[natsala13]

@njhill @khluu Would it be possible to consider #55178 for the v0.29.0 cherry-pick milestone? This is time-sensitive for us: we have a downstream workload scheduled to begin shortly, and we need a released vLLM build containing this fix to proceed safely.

The PR fixes silent Mamba recurrent-state corruption under speculative prompt-tail padding. On latest-main MRV2, it restored SWE-bench MTP5 from 15.60% to 39.13%, matching the 38.93% MTP-off control, with all 4,500 rollouts completing.

I realize rc4 has already been cut, but if there is still room for one more critical-fix cherry-pick/RC, we would greatly appreciate its inclusion.

cc. @TomerBN-Nvidia @Naveassaf