[BI] Close the cudagraph paths that pick a graph per step, and warn about the one that stays

[aoshen02]

What

Batch invariance is an absolute claim: a request's output must not depend on
which other requests share its engine step. Several components schedule
cudagraphs of their own, keyed on what the step happens to carry, and
cudagraph_mode does not reach them. This closes the ones whose condition is
exact, and warns where the answer is a per-model property config validation
cannot establish.

• Microbatching is refused. It decides per step whether to split and which
  graph to replay, and never consults cudagraph_mode — so unlike everything
  else here it is not gated on the mode, and --enforce-eager does not switch
  it off either. Opt-in (enable_dbo / ubatch_size), and the condition is
  exact, so refuse rather than silently disable.
• Speculative-decoding proposers drop to NONE. They pick PIECEWISE
  whenever the target's mixed mode is PIECEWISE or FULL, which puts a
  per-step graph/eager choice back on the draft path.
• A warning when an explicit capture envelope is too small. If a deployment
  sets max_cudagraph_capture_size below its own largest decode step, batches
  above it fall back to eager, so which numeric path a request takes depends on
  how many requests share its step. Only an explicit choice is checked — the
  defaults come from the platform, the model (some override the cap) and, under
  speculative decoding, from query-length alignment, so re-deriving "what the
  default would have been" only produces false alarms on stock configs. The
  answer is snapshotted before try_verify_and_update_config() can let a model
  write its own capture defaults, and stored so it survives replace().
• A warning for piecewise modes. A piecewise mode picks the per-step path
  from batch properties, so a request's decode step runs under a path its
  neighbours chose. That is harmless exactly while every selectable path is
  bit-identical for this model — verified on DeepSeek-V4, unknown elsewhere.
  cudagraph_mode=FULL removes the choice.

Why the mode is no longer pinned

An earlier version of this PR pinned cudagraph_mode to FULL under
VLLM_BATCH_INVARIANT. The leak that motivated it turned out to be a
DeepSeek-V4 model bug — a host-side indexer predicate evaluated once against
the capture-time dummy batch and baked into the graph — fixed upstream in
vllm-project#52492. With that fixed, a 50-round soak of the full 43-layer model under
FULL_AND_PIECEWISE is bitwise clean, where the same soak leaked in 38.7% of
rounds before. Pinning would now cost the mixed-step graph for nothing, so the
mode is left alone and the burden moves to the warning: a deployment that has
not checked its own model should not read silence as safety.

Tests

tests/test_config.py -k batch_invariant → 13 passed, covering: the
requested piecewise mode survives; the warning fires; resolve_cudagraph_mode_ and_sizes is no longer special-cased (BI on and off return the same mode);
microbatching is refused under both FULL and NONE; the capture-envelope
warning fires on an explicit undersized choice and stays quiet on defaults; and
a model-written capture default is not mistaken for a deployment choice.

tests/v1/determinism/test_batch_invariant_dispatchers.py → 1 passed: the
microbatching decision point as a unit test over the branch itself, no weights,
no GPU.

Negative control: with the old pinned build in place, 6 of those 13 fail — the
assertions are testing the removal, not restating it.

Non-BI behaviour unchanged, checked explicitly: with the flag off, both
piecewise-bearing modes survive and the mixed-mode downgrade still picks
FULL_AND_PIECEWISE.

CI: the determinism suite now also runs when vllm/config/,
vllm/v1/spec_decode/ or gpu_model_runner.py change — previously it only
triggered on kernel changes, so config-level regressions here were untested.

Model evaluation

Not applicable to the current contents: nothing here changes numerics on any
configuration we run (microbatching off, speculative decoding off, capture
sizes default). The piecewise measurements that justify not pinning are in
the soak referenced above — full 43-layer DeepSeek-V4-Flash-Base TP4/EP4,
50/50 rounds with 0 differing cases under FULL_AND_PIECEWISE, and GSM8K
1319/1319 per-question agreement across two independent re-runs at concurrency
64.

Scope, and what was deliberately dropped

Two more dispatchers of the same shape were found and then removed from this
PR (2026-08-18): the multimodal encoder's cudagraph manager, and Gemma4's
self-managed centroid graphs. Both were motivated by the pin — it promised no
numeric path is selected from the batch, so every escape had to be closed. With
the pin gone that obligation goes too, and both would change numerics on models
the submitter has never loaded (the Gemma4 one is not a one-liner: skipping the
capture needs an eager centroid path of its own, or the base class silently
substitutes a full-vocab argmax). They are recorded for whoever has that
hardware.

Review follow-up

An external review (codex, gpt-5.6) raised five points on this branch; each
was checked against the source before acting.

• True, fixed. The microbatching refusal fired on use_ubatching, the
  configured flag, not on whether a step can actually be split. should_ubatch
  stays False unless data_parallel_size > 1, and check_ubatch_thresholds
  needs a step at least as large as one of the two DBO thresholds — unreachable
  once both exceed the token budget. Either case was a startup failure for a
  configuration that could never microbatch. Now refused only when it can run.
• True, fixed. The capture-envelope warning multiplied max_num_seqs by
  the global num_speculative_tokens. Dynamic speculative decoding pairs the
  large draft counts with small batches, so that product is a step the schedule
  forbids, and the warning fired on an envelope that already covered every real
  step. The estimate now follows the schedule, extracted into
  _largest_uniform_decode_step so the arithmetic is testable without a draft
  model.
• True, warned. cudagraph_mm_encoder greedy-packs the images scheduled
  this step and picks a captured budget from their combined token count,
  consulting neither cudagraph_mode nor --enforce-eager. Warn rather than
  refuse: it is opt-in and no model exercised here has an encoder, so a refusal
  would be a promise about a deployment that was never tested.
• True, warned. The two draft-side dispatchers switched off here are the
  legacy ones; the V2 speculators (AutoRegressiveSpeculator,
  DFlashSpeculator) build cudagraph managers of their own that this check
  never reaches, and upstream [Feature]: Batch Invariant Feature and Performance Optimization vllm-project/vllm#27433 records that batch invariance under
  speculative decoding is unfinished. Said so, rather than let partial coverage
  read as coverage.
• Out of scope, deliberately. Gemma4's centroid graphs are the same class
  of gap, and skipping their capture needs an eager centroid path or the base
  class silently substitutes a full-vocab argmax. Changing numerics on a model
  nobody here has run is not something the submitter can defend; the analysis is
  kept as a draft.

The review also confirmed clean: _capture_sizes_user_specified is invisible to
dataclass equality and compute_hash(), survives with_hf_config() and the
engine's deepcopy/cloudpickle round-trips, and the snapshot ordering does
separate a CLI choice from GPT-OSS's model hook.

Notes

• Not a duplicate: [Bugfix][DSv4] Keep indexer scoring in breakable graphs vllm-project/vllm#52492 fixes the DeepSeek-V4 indexer predicate, a model bug;
  this is about dispatchers that schedule graphs outside cudagraph_mode's
  reach.
• Path-disjoint from the other five BI PRs (vllm/config/vllm.py,
  vllm/v1/spec_decode/* and the two test files only).
• AI assistance was used (Claude); every line human-reviewed before merge.

🤖 Generated with Claude Code

[aoshen02]

Added a second commit: [BI] Warn when a decode step can escape the cudagraph capture set.

Pinning cudagraph_mode=FULL under VLLM_BATCH_INVARIANT closes the leak for
the default capture set, but it leaves one hole. FULL resolves to
FULL_DECODE_ONLY here, so path selection depends on the kind of step rather
than its size — and that only holds because max_cudagraph_capture_size defaults
to min(max_num_seqs * decode_query_len * 2, 512), which a decode step can never
reach. If someone sets cudagraph_capture_sizes / max_cudagraph_capture_size
explicitly below max_num_seqs * decode_query_len, large decode batches fall
back to eager and the same neighbour-dependent leak returns under FULL.

The new commit warns at config time when that combination is requested. It only
warns — it does not raise, and it does not raise the ceiling on its own, since
either would change behaviour for configs that are otherwise fine.

Placement matters: the check sits after _set_cudagraph_sizes() so
max_cudagraph_capture_size is already final. Two regressions added in
tests/test_config.py (one asserting the warning fires, one asserting it stays
quiet when the capture set does cover the largest decode step). Both need a real
ModelConfig — with model_config=None, max_cudagraph_capture_size resolves
to 0 and the guard correctly skips.

Verified: tests/test_config.py -k batch_invariant → 5 passed.

---

Separately, the root cause behind this whole issue has since been found, and
it is a model bug rather than a cudagraph-policy one: DeepseekV4Indexer.forward
evaluates a host-side short-context predicate that gets frozen at capture time,
so every graph-replayed step silently runs the "select the earliest candidates"
fallback instead of the real top-k. That fix is a separate one-line change on
bi/indexer-shortcut — it is independent of batch invariance, so it is not
bundled here. With it applied, PIECEWISE itself becomes bit-identical to FULL
and eager (8/8 and 2/9 differing filler counts → 0/8 and 0/9).

This PR is still worth having on its own: it is the defence-in-depth layer, and
it is what keeps a batch-invariant run from silently depending on the capture-set
configuration.

[aoshen02]

Pushed 8aa2f9f0: the pin stays, but the warning no longer claims piecewise
cudagraphs break batch invariance, because on DSv4 they measurably no longer do.

Once bi/indexer-shortcut bounds the indexer's short-context predicate on a
static length, a 50-round soak of the full 43-layer model (TP4/EP4, GB200) under
the production-default FULL_AND_PIECEWISE is bitwise clean. The same soak
leaked at a 38.7% per-round rate before that fix, so 50 clean rounds is ~1.5e-11
by chance.

Getting that measurement required stepping over this very pin, so the run carries
its own vacuity guard — it is only counted if all three hold: the pin warning
never fires, the resolved mode really is FULL_AND_PIECEWISE, and both PIECEWISE
and FULL captures are non-zero. Any one missing and the cell aborts rather than
reporting a pass, because "the path under test was never taken" is exactly how
this investigation produced a false negative once before.

So what changed is the justification, not the behaviour:

• before: piecewise breaks invariance, therefore pin it.
• now: piecewise makes a request's numeric path follow the token count of
  the batch it lands in. That is harmless only while every selectable path is
  bit-identical. It is bit-identical here because no host-side value read inside
  a captured region is frozen against the capture-time dummy batch — the indexer
  predicate was the one instance, found by a tree-wide search that turned up no
  second one. Nothing in __post_init__ can check that property for the next
  model or the next patch, so the pin is kept as a policy.

Reviewers who want the graph back on a model they have audited themselves now
have an accurate warning to reason about, and the code comment says what the pin
is actually buying.

Measured cost of keeping it: none detectable on our workload — BI=1 warm median
0.980s with the pin against 0.99s before it. Decode steps keep their FULL graph;
only mixed prefill/decode steps fall back to eager.

[aoshen02]

Pushed 58111df8 and retitled. The FULL pin is gone; the PR is now about the
cudagraph paths that pick a graph per step without consulting cudagraph_mode,
plus a warning for the one path that does.

Owner decision, and the measurement supports it: with the indexer predicate fixed
(upstream vllm-project#52492, merged 2026-08-17), a 50-round soak of the full
DeepSeek-V4-Flash-Base TP4/EP4 under the production-default FULL_AND_PIECEWISE
is bitwise clean — the same soak leaked in 38.7% of rounds before. Pinning to FULL
was costing the mixed-step graph for a guarantee it was not actually providing.

What went with the pin, and why each had to:

• the "can use neither piecewise nor full" refusal — it existed only because
  the pin ran late and would silently undo vLLM's own downgrades. Nothing to undo
  now.
• _full_cudagraph_unsupported_reason — that refusal was its only caller.
  Leaving a helper nothing calls would be worse than deleting it.
• the batch-invariant branches in resolve_cudagraph_mode_and_sizes — two
  tests now assert those fallbacks resolve identically with and without
  invariance, which is the inverse of what they asserted before.

What stayed, because none of it depends on the pin: the capture-set escape
warning, and the refusals for microbatching, prefill context parallel, the
multimodal encoder dispatcher and the speculative-decode dispatchers. Those close
paths where a graph is chosen per step and cudagraph_mode is never consulted, so
no mode setting can fix them.

Tests: tests/test_config.py -k the cudagraph/batch-invariant selection, 13
passed. Negative control — re-inserting the pin fails exactly the two
mode-survives cases and nothing else, so the test is watching the behaviour and
not the wording.

One gap a reviewer should not inherit. The soak that justifies this ran with
our own version of the indexer fix, which bounded the predicate on a static length
so both paths did full scoring — agreement by construction. vllm-project#52492 instead keeps
the eager shortcut, so a graphed step does full scoring while a non-capturing step
may still take the shortcut. Those agree only while "candidates <= topk" makes the
shortcut's set identical to the real top-k, and an identical set is not an
identical order — the order sensitivity of FlashMLA sparse decode measured
321-361 bf16 ULP under permutation in this project and was never explained.
Enabling piecewise on top of vllm-project#52492 warrants re-running that soak rather than
inheriting this result. The warning text is deliberately worded so a deployment
that has not done so does not read silence as safety.

[aoshen02]

Retracting the "Not claimed" paragraph in 58111df8's commit message. It says a
deployment enabling piecewise on top of upstream vllm-project#52492 should re-run the
50-round soak, because vllm-project#52492 keeps the eager shortcut while a captured step
does full scoring, and an identical top-k set is not an identical order.

That premise is wrong, and reading the kernel settles it without a GPU. Both
top-k kernels this model can dispatch to write ascending indices when the
candidates fit — csrc/libtorch_stable/persistent_topk.cuh:915 and
cooperative_topk.cuh:491 carry the same trivial branch:

if (seq_len <= TopK) { row_output[i] = (i < seq_len) ? i : -1; }

which is exactly what _fill_short_context_topk_indices writes. So wherever the
shortcut applies, vllm-project#52492's two branches emit a bitwise identical index array —
same set, same order — and a graphed step cannot disagree with an eager one on
that account.

It also means the two fixes are equivalent on this deployment
(index_topk=512, compress_ratios ∈ {4, 128}, max_model_len=8192): the
ratio-128 layers never exceed 64 candidates so both paths take the trivial
branch, and the ratio-4 layers sit at 2048 > 512 so both do real scoring.

The probe that motivated the caveat measured 321-361 bf16 ULP under a random
permutation of the index list. No code path produces a non-monotonic list, so
that input does not occur in production and the measurement says nothing about
vllm-project#52492. That line is closed; its record is marked accordingly. Still unchecked,
and out of scope here: whether the kernels require monotonic indices at all, and
the third dispatch path top_k_per_row_decode, which this model never takes
(index_topk is 512).

Nothing in the code or tests changes — the retraction is of a caveat, not of a
claim the PR makes.

[aoshen02]

Dropped the multimodal encoder and Gemma4 dispatcher changes (6494fe8e6c), and
rewrote the body — it still described pinning cudagraph_mode to FULL, which
was removed in 58111df82e.

Both dropped changes existed because of the pin. The pin promised that no
numeric path is selected from the batch, so every dispatcher that schedules
graphs of its own had to be closed too — otherwise the promise was false in a
way a deployment could not see. What replaced the pin is a warning, which makes
no promise on the deployment's behalf, so that obligation is gone.

What is left is two changes to models nobody on this side has ever loaded. The
Gemma4 one is not a one-liner either: skipping the centroid capture requires an
eager centroid path of its own, or the base class silently substitutes a
full-vocab argmax. Changing numerics on unrun models is not something the
submitter can defend line by line, which is the bar this repo asks for.

Kept: microbatching (the only one that never consults cudagraph_mode and that
--enforce-eager does not reach either), the speculative-decoding proposer, the
capture-envelope warning, the piecewise warning, and the CI wiring.

Verification after the change: tests/test_config.py -k batch_invariant
13 passed, test_batch_invariant_dispatchers.py 1 passed, and the non-BI
control unchanged on all three assertions. Negative control: against the old
pinned build, 6 of the 13 fail — so they are testing the removal, not restating
it.

[aoshen02]

Rebased onto fork/bi/base, which now points at upstream aa9903490 — the exact
commit the nightly container image is built from, so the files copied into
site-packages can no longer be half-old. The previous base was 3ee2df303,
218 commits behind, and upstream had touched every file this stack changes.

Two things upstream did for us in that window:

• [Bugfix][DSv4] Keep indexer scoring in breakable graphs vllm-project/vllm#52492 fixes the DeepSeek-V4 indexer short-context predicate being frozen at
  capture. That was bi/indexer-shortcut here; the branch is deleted.
• [Bugfix][DSv4] Revert adaptive C128A metadata packing vllm-project/vllm#51318 reverts the adaptive C128A metadata packing, so active_topk_width
  no longer exists and is no longer computed from max_seq_len of the step. The
  commit that pinned it (Size the C128A top-k width from the config) is dropped
  as redundant — that was the only rebase conflict in the stack.

Post-rebase verification: tests/v1/determinism (the six files this stack adds
or touches) → 120 passed, 0 failed on GB200 against the new image.