[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 vllm-project/vllm#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 vllm-project/vllm#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 vllm-project/vllm#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 vllm-project/vllm#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%