Overview

Fix Mamba "align" state copies reading another request's block ids under MRV2 + PP + MTP, which poisons recurrent state (all-NaN logits, constant-token output, acceptance collapse). Consolidates #55506 and #57807 into a smaller change; both authors are co-authors. Refs #54173.

Claims

• The precopy/postprocess state-copy kernels index the persistent per-request-slot block tables by req_idx, so a postprocess that runs steps after its batch (non-last PP rank) no longer reads tables re-gathered for a later batch.
• compute_aligned_state_indices is unchanged: it still reads the gathered batch-order tables, so padded rows and CUDA graph capture behave as on main.
• The align context is built once in initialize_kv_cache instead of lazily from whichever caller runs first, which removes the capture-time / profiling binding issues raised on [Bugfix] Index mamba spec-decode block tables by request slot, not batch row #55506 and [Bugfix][Model Runner V2] Separate Mamba aligned-metadata and state-copy contexts #57807.

Validation

• pytest tests/kernels/mamba/test_precopy_mamba_align.py tests/v1/worker/test_mamba_utils.py tests/v1/worker/test_mamba_hybrid_model_state.py tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py: 146 passed (GB200).
• With the kernels reverted to batch-row indexing, the new/updated kernel tests fail (26 failures), so they cover the bug.
• E2E on GB200, Qwen3.6-27B, MRV2, MTP 4, --async-scheduling --enable-prefix-caching --mamba-cache-mode align (reproduction from @tomasruizt on [Bugfix] Index mamba spec-decode block tables by request slot, not batch row #55506). main = upstream 31e2443.

The stuck requests on main are the constant-token loop from poisoned recurrent state; with a 4096 max-token cap, main PP=2 averaged 235 output tokens per GSM8K question vs 82 for this PR at PP=2 and PP=1. The fix does not change kernel work; the throughput gain comes from the recovered acceptance length.

VLLM_USE_V2_MODEL_RUNNER=1 vllm serve Qwen/Qwen3.6-27B --language-model-only \
  --pipeline-parallel-size 2 --async-scheduling --enable-prefix-caching \
  --mamba-cache-mode align \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}'
# PP=1 reference: drop --pipeline-parallel-size 2

vllm bench serve --backend vllm --model Qwen/Qwen3.6-27B --dataset-name random \
  --random-input-len 2048 --random-output-len 1024 --random-range-ratio 0.75 \
  --random-prefix-len 1024 --ignore-eos --temperature 0 --seed 42 \
  --num-prompts 64 --max-concurrency 8

python tests/evals/gsm8k/gsm8k_eval.py --num-shots 5 --num-questions 1319 \
  --max-concurrency 64   # acceptance length column; accuracy from mqa gsm8k (5-shot, chat, no thinking)

Details

Root cause: MambaSpecDecodeGPUContext captures block-table pointers once. MRV2 bound it to input_block_tables, a single buffer re-gathered in batch order every step, and the copy kernels indexed it by batch row. On a non-last PP rank the postprocess for step N runs after step N+k has re-gathered that buffer, so batch rows point at other requests' blocks. With the unified hybrid cache layout those copies land in live state.

Why not just #55506 or #57807: both keep lazy binding and add machinery around it (temporary bindings / bind flag, or two lazily bound contexts plus profiling teardown), and #55506 remaps the aligned-index kernel through idx_mapping, which changes padded-row behavior. Binding both table sets eagerly avoids all of that. The source tables are safe to read late because a request is not rescheduled (so its row is not rewritten or its slot reused) until its in-flight step's output is processed.

MRV1 is unaffected: it launches both kernels without idx_mapping, where req_idx == batch_idx, and passes a single table set, which stays the default for the aligned-index tables.

Credit: root cause first identified by @zebgop-ops; fix direction and kernel tests from @Karl0007 (#55506); separate metadata/copy tables and the profiling lifetime issue from @yongqinwang-cmd (#57807).