Skip to content

[Bugfix][DSv4] Bound the indexer short-context shortcut on max_model_len - #27

Closed
aoshen02 wants to merge 1 commit into
bi/basefrom
bi/indexer-shortcut
Closed

aoshen02 wants to merge 1 commit into
bi/basefrom
bi/indexer-shortcut

Conversation

@aoshen02

@aoshen02 aoshen02 commented Aug 16, 2026 •

Copy link
Copy Markdown
Owner

What

DeepseekV4Indexer.forward short-circuits the Lightning indexer when the whole
candidate set fits in top-k. The predicate reads this step's
indexer_metadata.max_seq_len. Bound it on the configured maximum instead:

-        if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
+        # self.max_model_len is already divided by compress_ratio
+        if self.max_model_len <= self.topk_tokens:

One file, +10/-1 (vllm/models/deepseek_v4/attention.py), 9 of those lines a
comment explaining the hazard.

Why

That if is a host-side branch that runs inside a captured cudagraph
segment
. DeepSeek-V4 auto-enables VLLM_USE_BREAKABLE_CUDAGRAPH=1, which
forces CompilationMode.NONE — no torch.compile, so no Dynamo guards — and the
model's only eager break is DeepseekV4Attention._sparse_indexer_and_attn
(43 breaks for 43 layers, measured). The indexer runs before that break, so
the predicate is evaluated exactly once, at capture, against the dummy batch,
and the chosen side is baked into the graph.

Logged predicate, capture vs. real step (SM100 / GB200, TP4/EP4,
0.26.1rc1.dev668+g3ee2df303):

max_seq_len branch
capture, dummy batch (mode=PIECEWISE) 1 – 2 short=True
real step, eager fallback (mode=NONE) 3268 / 3410 short=False

So every graph-replayed step runs _fill_short_context_topk_indices, which
writes

tl.where(offsets < num_compressed, offsets, -1)

— the earliest compressed positions in order, not the top-k by score. At
3268 tokens that is 817 candidates against topk=512: the step sees roughly the
first 2k tokens of the prompt; the rest of the context is invisible to the
sparse-attention path
(the model still sees it through everything the indexer
does not gate). Steps
that exceed max_cudagraph_capture_size fall back to eager and run the real
indexer, so the same request alternates between two different sparse patterns
depending on how large its step happened to be.

This is a correctness bug on its own, independent of batch invariance.
Ordinary long-context serving under piecewise silently uses the wrong sparse
pattern; without a batch-invariance checker it just looks like unexplained
quality loss.

cudagraph_mode=FULL is already safe under the capture configuration we
tested
(the defaults; we did not probe an explicitly shrunk capture set), for a
reason worth recording: its
capture path passes for_capture=True
(v1/worker/gpu/cudagraph_utils.py:643), and DefaultModelState then uses
max_model_len as the host bound (v1/worker/gpu/model_states/default.py:149)
— measured msl=8192 at FULL capture, so the long branch is what gets baked.
PIECEWISE passes for_capture=False, under a comment that already flags the
risk: "We assume that attention-like operations intended for capture will still
produce capturable metadata, even when for_capture=False. While this assumption
is brittle, it currently works in practice."
Breakable cudagraphs capture every
non-breakpoint op, so it no longer holds. The one-line model fix is the safe
change; whether PIECEWISE capture should also use worst-case host bounds is a
broader call.

Tests

Batch composition pinned by driving LLMEngine with add_request/step by
hand (HTTP-level scheduling changes composition and confounds the measurement).
Victim ~3.3k-token prompt (the probe logs show max_seq_len 3268), N short or
long fillers, compare the victim's first four logprobs against the solo
reference.

DMODES=piecewise BI_TAG=-fixed bash agent_run/scripts/run_probe_gve3_dense.sh
PATCH=0 DMODES=full BI_TAG=-fixed bash agent_run/scripts/run_probe_gve3_dense.sh
scenario before after
PIECEWISE, long fillers (the shared step exceeds the capture size) 8/8 filler counts differ 0/8
PIECEWISE, short fillers (step ≤ capture size → graphed) 2/9 differ 0/9
FULL (no-regression) 19/19 identical 19/19 identical, same values

After the fix, PIECEWISE values are bit-identical to FULL and eager
(-0.6408482789993286). PIECEWISE was confirmed live in the fixed run (zero
mode-pinning messages, two PIECEWISE capture bars) — it is not passing because
the graph was skipped.

Related configuration suite: tests/test_config.py -k batch_invariant → 7
passed
(run in the pinned container; the host venv's huggingface_hub is too
old for that file's imports). Those cover the cudagraph-mode policy change, not
this predicate — this PR still has no unit test that exercises the shortcut
directly
; the evidence for it is the pinned-composition probe above.

Model evaluation

The fix is an identity transform for cudagraph_mode=FULL on this config —
before: 8192 // 4 = 2048 <= 512 is False; after: self.max_model_len (=2048) <= 512 is also False — same branch. Confirmed empirically by the 19/19 row
above. The accuracy and throughput numbers below were measured in the
FULL-pinned configuration and carry over by that identity argument -- they are
not a fresh full-eval rerun under PIECEWISE:

  • gsm8k full set, two independent runs: acc 0.9113 / 0.9113, per-question
    agreement 1319/1319
  • online batch-invariance soak: 47/47 rounds with zero differing logprobs
    (baseline single-round hit rate 38.7%)
  • throughput: BI=1 median 0.980s (0.99s archived pre-fix) — no measurable cost

Notes

AI assistance was used to produce this change. ⚠ Before this goes upstream the
AGENTS.md duplicate checks (gh issue view, gh pr list --search by issue and
by area) still have to be run and their results pasted here. Related but, on the
evidence we have, not duplicated:
vllm-project#52109 is a different defect (gfx942 indexer K-cache
FLAT-write / SHUFFLE-read layout mismatch; that reporter rules out
capture/replay with an --enforce-eager control).

DeepseekV4Indexer.forward took the "every candidate is selected" shortcut
when this step's max_seq_len // compress_ratio fit in topk_tokens. That
predicate is a host value, and the indexer runs inside a captured cudagraph
segment under breakable PIECEWISE, so it is evaluated once against the
capture-time dummy batch and baked in. The dummy batch has max_seq_len 1-2
(piecewise capture builds attention metadata with for_capture=False, so the
host bound is the dummy's own length rather than the worst case), so every
replayed step takes the short-context path. With a 3268-token context that
writes indices 0..topk-1 instead of the score-based top-k, while any step
that exceeds max_cudagraph_capture_size falls back to eager and runs the
real indexer -- two different sparse-attention patterns for the same request.

self.max_model_len is already divided by compress_ratio, so bounding on it
takes the shortcut only when it holds for every request the engine can
serve, which is capture-safe in every cudagraph mode. FULL was already safe
because its capture path passes for_capture=True and DefaultModelState then
uses max_model_len as the host bound.

Verified on DeepSeek-V4-Flash-Base, TP4/EP4, 4xGB200, VLLM_BATCH_INVARIANT=1,
cudagraph_mode=PIECEWISE, with batch composition pinned by a hand-driven
LLMEngine: the victim's logprobs went from 8/8 and 2/9 filler counts
differing to 0/8 and 0/9, and now match the FULL/eager values bit-for-bit.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@aoshen02

Copy link
Copy Markdown
Owner Author

Closing: superseded upstream. vllm-project#52492 ("[Bugfix][DSv4] Keep
indexer scoring in breakable graphs") was merged on 2026-08-17, and it fixes the
same defect — the short-context predicate is a host value read inside the region
that PIECEWISE captures, so the capture-time dummy batch (max_seq_len=1..2)
freezes it as "short context" and every graphed step then writes the earliest
topk candidates instead of the score top-k.

Per AGENTS.md we do not carry a competing PR, and the upstream fix is the better
shape: and not torch.cuda.is_current_stream_capturing() keeps the eager
shortcut, where this branch gave it up everywhere by bounding the predicate on
the static self.max_model_len.

One difference worth recording for whoever picks this up, because it is not
covered by the evidence gathered here. Under this branch both paths do full
scoring, so graph and eager agree by construction — that is what the 50-round
FULL_AND_PIECEWISE soak measured. Under vllm-project#52492 the graphed step does full
scoring while a non-capturing step may still take the shortcut, so the two agree
only while "candidates <= topk" makes the shortcut's set identical to the real
top-k. Identical set is not identical order, and the order sensitivity of
FlashMLA sparse decode is an open item in this project (a probe measured 321-361
bf16 ULP under index permutation and the mechanism was never explained). Anyone
relying on piecewise + BI on top of vllm-project#52492 should re-run that soak against the
upstream fix rather than inherit this branch's result.

Evidence and reasoning: agent_run/reports/batch_invariance/ASSUMPTIONS.md §4
and results/batch_invariance/fp8-flash-base/_e2e/soak-piecewise-postfix/.

@aoshen02 aoshen02 closed this Aug 17, 2026
@aoshen02
aoshen02 deleted the bi/indexer-shortcut branch August 18, 2026 15:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant