[Bugfix][V1] Keep still-prefilling rows out of full decode graphs - #60512

Overview

Full CUDA-graph dispatch in Model Runner V1 is decided by batch shape alone, so a one-token prompt looks like a decode step and replays a decode graph. On Mamba-hybrid models the result is garbled output that differs on every run. This gates full graphs on the batch actually being a decode, the way Model Runner V2 already does. Fixes #57720.

Claims

• One-token prompts on Mamba-hybrid models are deterministic again on MRV1 under the default FULL_AND_PIECEWISE: 5 of 5 one-token prompts go from 6 distinct outputs in 6 runs to 1.
• Prompts of two or more tokens are byte-identical before and after the change, so nothing outside the affected batch shape moves.
• MRV2 is untouched: it already rejects these batches ([Bugfix][MRV2] Require all requests to be decoding for uniform-decode dispatch #51865, refined by [Perf][MRV2] Allow FULL decode graphs for one-token prompt tails #58400). This is the MRV1 counterpart of that guard.

Validation

End-to-end reproducer from the issue, tiiuae/Falcon-H1-Tiny-90M-Instruct, VLLM_USE_V2_MODEL_RUNNER=0, cudagraph_mode=FULL_AND_PIECEWISE, temperature=0, 6 runs per prompt, RTX 4070 Laptop (8 GB):

Totals: unstable=5 of 7 before, unstable=0 of 7 after.

import collections
from vllm import LLM, SamplingParams

if __name__ == "__main__":
    llm = LLM(model="tiiuae/Falcon-H1-Tiny-90M-Instruct", max_model_len=4096)
    sp = SamplingParams(temperature=0.0, max_tokens=32)
    for p in ["Hello", "Hi", "Yes", "No", "The", "Hello there"]:
        outs = [llm.generate([p], sp, use_tqdm=False)[0].outputs[0].text for _ in range(6)]
        print(len(collections.Counter(outs)), "distinct in 6:", repr(p), "->", repr(outs[0][:40]))

Test suites, all passing on the same machine:

Also ran facebook/opt-125m on MRV1 end-to-end to confirm the attention-only path is unaffected (stable, coherent output). pre-commit run --files <changed files> passes, including ruff check, ruff format, typos, mypy, and the test-tethering hook. tools/pre_commit/mypy.py 3.12 reports the same 48 pre-existing errors in gpu_model_runner.py with and without the change, so no new type errors.

Validation constraints: single consumer GPU, so no TP/PP, no disaggregated-prefill and no speculative-decoding end-to-end run. The spec-decode shapes are covered by unit tests instead. The local flashinfer JIT cannot build in this environment (CUDA toolkit/compiler header mismatch), so the reproducer ran with VLLM_USE_FLASHINFER_SAMPLER=0; that is unrelated to the dispatch path being changed.

Details

Root cause:

1. GPUModelRunner._is_uniform_decode classifies a batch purely by shape: max_num_scheduled_tokens == uniform_decode_query_len and num_tokens == max * num_reqs. A fresh one-token prompt satisfies both.
2. The batch therefore dispatches to a full graph. For Mamba, AttentionCGSupport.UNIFORM_BATCH resolves the default mode to FULL_AND_PIECEWISE, so full graphs cover exactly these "uniform decode" batches.
3. BaseMambaAttentionMetadataBuilder._compute_common_metadata classifies the same row correctly as a prefill: it only reclassifies one-token rows that have prior state (has_prior_state = seq_lens > query_lens), and explicitly keeps first prompt chunks as prefills.
4. It does not matter, because has_prefill / has_decode in MambaMixer2.forward are Python branches. Graph replay re-runs neither; the captured graph holds only the decode kernels.
5. Those kernels (causal_conv1d_update, selective_state_update) update conv_state / ssm_state in place. The prefill path is what seeds them. A first prompt chunk therefore reads whatever the freshly allocated Mamba block contains, which is why the answer is fluent, unrelated to the prompt, and different on every run.

This also explains the workarounds reported on the issue: PIECEWISE and --enforce-eager re-run the Python branch each step, prompts of two or more tokens are not uniform-shaped, and --mamba-cache-mode and prefix caching are orthogonal.

Design: rather than a Mamba-specific gate, _is_decode_graph_eligible mirrors decode_graph_eligible in MRV2's gather_batch_req_state — a batch may take a full graph only if every row is done prefilling or is a single new token over existing context, with placeholder drafts discounted so padded prompt tails keep their graph. Keeping one semantics across the two runners seemed better than inventing a narrower MRV1-only rule.

MRV1 deprecation: MRV2 already has this guard, so nothing there changes. MRV1 still needs it because it remains selectable via VLLM_USE_V2_MODEL_RUNNER=0 and is still the default for some configurations (ROCm architectures in ROCM_DEFAULT_MRV1_ARCHITECTURES, environments without Triton, and models using features MRV2 does not support).

Tradeoff: as in MRV2, the guard is backend-agnostic, so an attention-only model with a lone fresh one-token prompt in an otherwise decode-shaped batch now runs piecewise instead of full. That shape is rare, the fast path for real decode batches is unchanged, and the early return keeps the check to two array comparisons when no row is prefilling.

Testing approach: the three new cases sit beside test_is_uniform_decode and exercise the predicate directly with a lightweight stub, matching how tests/v1/worker/test_gpu_batch_ordering.py covers the MRV2 equivalent. They mirror the MRV2 assertions: first prompt chunk rejected, one-token tail with prior context kept, draft-padded tail kept while a bare K+1 first chunk is rejected.