[BI][DSv4] Pin sparse-MLA decode and prefill scheduling under batch invariance by aoshen02 · Pull Request #21 · aoshen02/vllm

aoshen02 · 2026-08-15T07:22:48Z

What

Pins sparse-MLA decode and prefill scheduling under
VLLM_BATCH_INVARIANT=1.

Decode: the FlashMLA sparse decode kernel builds its internal KV-split
plan from batch-wide state. Under the flag the tile-scheduler plan is
pinned and each request is split by its own length only; the plan is built
once per layer type per step (host cost 37 ms → 0.2 ms/step).

Prefill: flash_mla_sparse_fwd takes no schedule metadata — it plans
internal splits from the call geometry (chunk query rows, gathered
widths). The greedy prefill chunk packing makes a request's geometry depend
on its neighbors, so identical requests produced different bits depending
on co-scheduled prefills. Under the flag the chunk plan becomes one
request per chunk: the geometry is a function of the request alone,
constructively batch-invariant.

The prefill bug was invisible to the previous e2e because prefix caching
recomputed the victim's prefill only once; a long-sequence sweep with
prefix caching off reproduced it 100% at a fixed 2-request composition, and
per-layer forward hooks localized the first divergence to L0 attention.

Perf

Decode: CUDA-graph cost 0.99–1.04× vs the shipped plan (b=1..17), with
correctness 0/7 mismatching and negative control 6/6 MISMATCH. Prefill
per-request chunks are covered by the e2e throughput A/B in the RFC PR.

Tests

tests/v1/determinism/test_sparse_mla_batch_invariance.py: decode plan
pinning + invariance, production extra-KV shape, per-request prefill chunk
plan (each request's (chunk_N, chunk_M) equals its solo plan).

E2e: online BI checker (long sequences, prefix caching off) 0 diff on
1 GPU and TP4/EP4, including the full 43-layer model.

GSM8K (full 43-layer Flash-Base, 1319 questions, 5-shot, greedy): BI=1
accuracy 0.911 vs BI=0 0.904 (0 invalid both; within noise).

Review follow-up

An external review (codex, gpt-5.6) read the pinned scheduling end to end and
found no surviving batch dependency: _split_budget derives from the
configured max_num_seqs, _blocks_per_split from the configured top-k caps,
_split_plan cuts only at fixed multiples of each row's own block count, and
_plan_skeleton never splits a request — including for batches above the
partition count. It also confirmed the flashmla.py cold start: the shipped-plan
result is overwritten by a second call after record_num_sm_parts, both pinned
metadata tensors are installed, and a failure to rebuild raises rather than
returning the first result. No changes were needed.

Notes

Not a duplicate: upstream FLASHINFER_MLA BI ([Feature]: Batch Invariant FLASHINFER_MLA support vllm-project/vllm#47359) is dense-MLA and
stalled; no upstream work touches the DSA sparse prefill chunk plan.
AI assistance was used (Claude); every line human-reviewed before merge.

🤖 Generated with Claude Code

Update: consolidation round (2026-08-15)

Branch rebuilt on a clean bi/base parent with the consolidation pass
folded in (shared test helpers, pure Triton key fn, tl.constexpr
constants — plain global ints fail to compile under Triton 3.7 — and
repo-pinned ruff format), adversarially reviewed (codex r10/r10b).
Container suite on GB200: test_sparse_mla_batch_invariance.py 21 passed.

Stacked on `bi/indexer-topk` (2026-08-16)

This branch is based on bi/indexer-topk, not on bi/base, and the PR targets
it. The dependency is real rather than a note in this description: the indices
this scheduling plan operates on come from the indexer's top-k, whose dispatch
switches implementation on num_rows <= 32 (sparse_attn_indexer.py,
cooperative vs persistent). Both resolve exact score ties through shared-memory
atomics, so on a tie plateau the selected set depends on how many rows share the
step. bi/indexer-topk collapses both into one deterministic implementation
under VLLM_BATCH_INVARIANT.

Pinned on its own, this branch would fix the schedule and leave the indices
batch-dependent. Basing it on the fix means that cannot be deployed by accident.

Combined HEAD verified on GB200: test_sparse_mla_batch_invariance.py 21
passed, test_indexer_topk_batch_invariant.py 39 passed.

Scope

Fixed here:

The decode plan no longer falls back to the shipped (batch-dependent) planner
above the SM partition count; whole requests are grouped per partition.
The cold-start step's shipped-planner output is never the one observed.
The C128A top-k width comes from the configured maximum, not from this step's
longest sequence.
Above the SM partition count the grouped plan is compared value-by-value
against the shipped planner on GPU, at the same 8 bf16 ULP bar the
below-the-count comparison uses. That is what says the hand-written metadata
columns mean what the kernel reads them as: getting one wrong moves the
output. Leaving one request unowned in the last partition makes it fail.

Cost of the constant C128A width

Pinning the width means a short request at a large max_model_len writes the
full row instead of a bucket sized to this step. The width is
align128(max_model_len / 128), so it is 128 -- i.e. unchanged -- for every
max_model_len up to 16k, and reaches 1024 at 128k.

Measured on GB200 by calling build_c128a_topk_metadata directly, device time
from the profiler (wall time sits on a ~50us launch floor below width 2048 and
shows nothing there):

width	max_model_len	decode 8192 tok	prefill 8192 tok
128	<= 16k	13.32us (1.00x)	5.74us (1.00x)
512	64k	13.42us (1.01x)	5.79us (1.01x)
1024	128k	13.54us (1.02x)	5.92us (1.03x)
2048	256k	21.58us (1.62x)	9.75us (1.70x)
4096	512k	38.23us (2.87x)	19.53us (3.40x)
8192	1M (config default)	70.63us (5.30x)	37.89us (6.60x)

The kernel writes every column of the row whatever the valid length is, but its
BLOCK_SIZE is 1024, so widths at or below that are the same vector store with
more of it masked -- the traffic grows and the trip count does not. Past 1024
each doubling adds a loop iteration, and the cost is real: at the width implied
by DeepseekV4Config's default max_position_embeddings of 1048576 the
metadata kernel is 5-7x its non-batch-invariant cost.

So: free up to max_model_len 128k, which covers the deployments this was
built for, and a genuine cost above it. Serving a 1M context under
VLLM_BATCH_INVARIANT should expect this, and a width-aware plan (bucket by
the request's own length rather than by the step's) is the obvious follow-up if
that matters -- it is deliberately not attempted here, because a per-request
bucket is exactly the batch-dependent quantity this PR removes and getting it
right needs the bucket to come from the request, not from the step.
Not claimed:

Runtime cost of the constant width in the prefill combined-index workspace and
the decode split cap. The table above covers the metadata kernel only.
Kernel-side coverage on SM90. The backend accepts compute capability 9 and
10; every FlashMLA test here was run on SM100 only, and the timings quoted
above are GB200's. The structural and control-flow tests are no longer gated
on SM100, so they do run on the SM90 half, but the kernel ones are skipped
there and nothing in this PR should be read as validated on H100.

All three CUDA top-k kernels (cooperative/persistent/per-row, prefill included) resolve exact-score ties through shared-memory atomics, so on tie plateaus both the selected set and its order change run to run — and relu(q·k) DSA scores make exact-0.0 plateaus routine. Selecting different KV is divergence, not rounding noise. Under VLLM_BATCH_INVARIANT=1, decode and prefill dispatch to a deterministic per-row top-k: score descending, exact ties by ascending column index, emitted ascending with a -1 tail, row-local (rowStart- relative) coordinates matching the CUDA prefill kernel's write contract. Fast path is one Triton program per row: the row is register-resident, a 32-step binary search on a monotone int32 key space finds the exact k-th key, and one ascending scan emits key>T plus the first budget ties by column index via exclusive cumsum — no sort anywhere. -0.0 canonicalizes to +0.0 before keying (torch.sort tie semantics); fixed num_warps, no autotune; the path is gated to the CUDA kernels' envelope (topk in 512/1024/2048, <=8192 columns), row bounds clamp to [0, num_cols] before int32 narrowing, and everything else falls back to the torch stable-sort reference (kept for other platforms and as the equivalence oracle). CUDA-graph events: 18.5us flat for n=1..64 (was 90-133us; nondeterministic CUDA baseline 8.2-14.4us); captures cleanly under PIECEWISE+FULL graphs. Tests (12): bitwise-equivalence matrix vs the reference (tie plateaus, quantized ties, mixed signed zeros, denormals, empty rows, narrow chunks, shifted cu_seqlen_ks), NCOLS bucket edges, all-equal and all--inf rows, strided/padded layouts, fallback dispatch past 8192 columns, int64 bounds past int32 range. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Consolidation pass folded in (codex r10/r10b reviewed): shared test helpers, pure Triton key fn, tl.constexpr-instantiated constants (plain global ints fail to compile under Triton 3.7), repo-pinned ruff format. Container suite 39/39 green; topk equivalence probe 240/240 bitwise.

The fused top-k masked lanes outside [row_start, row_end) to -inf and only gave true out-of-bounds lanes the INT32_MIN sentinel. An in-window score of -inf therefore shared a key with them, and since the tie budget is spent in column order, the masked lanes sitting before row_start consumed it first and were then dropped at emission. The number of indices a row emitted fell by exactly row_start -- which is where the packer happened to put that request, i.e. a function of its neighbours. The torch reference made the same choice, so comparing the two agreed on the wrong answer. Both now key every out-of-window lane at INT32_MIN, which sits strictly below _topk_bi_key(-inf). The reference switches to the kernel's integer key rather than sorting floats, so the two are equivalent by construction instead of by argument. tests: test_masked_lanes_do_not_spend_the_tie_budget asserts the invariant directly -- the same logical row at six different row_start offsets must select the same local indices -- because a reference comparison cannot catch a defect the reference shares. Negative control on the pre-fix kernel: 2 failed ("row_start=1 selected different indices than row_start=0"); after: 14 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Round 2 of the audit found the sentinel is not unique: a negative NaN (0xffffffff) maps to exactly INT32_MIN, the same key out-of-window lanes get. The kernel is unaffected -- its tie scan is already restricted to in-window lanes -- but the reference sorts first and drops after, so a masked lane could take a slot from an in-window NaN. Concretely, 600 in-window -NaN lanes at row_start=1 gave 512 indices from the kernel and 511 + a -1 from the reference. Lift in-window keys off the sentinel in the reference. Scores are not expected to be NaN; this is about the two implementations agreeing on every bit pattern rather than about a production input. tests: negative-NaN plateau against the sentinel, plus a window narrower than top-k inside a wider row (valid_count 30 < k_eff 512, num_cols 630 > TOPK) -- the shape where the cutoff lands on the sentinel. Negative control without the lift: 1 failed. With: 15 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The previous commit lifted in-window keys off INT32_MIN so they could not collide with the out-of-window sentinel. That is not sound: INT32_MIN and INT32_MIN + 1 are both reachable keys -- 0xffffffff and 0xfffffffe, two distinct negative NaNs -- so the lift merged them, and a row holding both ordered differently from the kernel. No sentinel value can be reserved here. Rank on validity first and key second instead. The kernel selects only from in-window lanes, so validity is already its implicit primary key; two stable sorts compose lexicographically and reproduce that for every bit pattern. tests: a window holding both sentinel-adjacent NaN patterns. Negative control with the lift: 1 failed. 15 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

… tests Round 4 accepted the validity-first ranking as equivalent to the kernel on all fp32 bit patterns but noted the tests only exercised two of them, both negative quiet NaNs. Parametrize over the classes that actually sit near the integer key's edges: signed zeros (canonicalized), infinities (its bounds), quiet and signaling NaNs, min and max denormals, and the two sentinel-adjacent patterns -- across three window shapes and both index_topk sizes. Written as raw int32 views: float32 -> Python float quiets signaling NaNs, so they cannot be constructed through .item(). Two comments contradicted the implementation and are corrected: the kernel no longer masks invalid columns to -inf before keying, and the reference no longer lifts in-window keys off a sentinel. Negative control against the pre-validity-ranking reference: 2 of the classes fail. 31 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

2048 is a supported index_topk and the selector specializes on it, but the bit-pattern sweep only ran 512 and 1024. Add it, along with a window wide enough (2600 in-window lanes) that the selection actually truncates at every topk under test rather than emitting the whole window. 39 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

_INT32_MIN was instantiated unconditionally at module scope. Without Triton `tl.constexpr` is None, so importing this module on a CPU-only path raised TypeError -- including during test collection for test_mla_short_prefill_indexer.py. Guard it on HAS_TRITON, the pattern already used by gumbel.py and lamport.py. 39 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The B200 batch-invariance job lists its test files one by one; source_file_dependencies only decides whether the job runs, it does not discover new files. This one was never executed by any job. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…nvariance Decode: the FlashMLA sparse decode kernel plans its internal KV splits from batch-wide state; under VLLM_BATCH_INVARIANT=1 the tile-scheduler plan is pinned and each request is split by its own length only, built once per layer type per step (host cost 37ms -> 0.2ms/step). CUDA-graph cost 0.99-1.04x vs the shipped plan. Prefill: flash_mla_sparse_fwd takes no schedule metadata and plans its internal splits from the call geometry (chunk query rows / gathered widths); the greedy prefill chunk packing makes a request's geometry depend on its neighbors. Under BI the chunk plan becomes one request per chunk, so the geometry is a function of the request alone — constructively batch invariant. Found by a long-sequence e2e sweep with prefix caching off (first divergence localized to L0 attention via per-layer forward hooks). Tests: decode plan pinning, per-request prefill chunk plan (each request's chunk equals its solo plan), production extra-KV shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Consolidation pass folded in (codex r10/r10b reviewed): shared test helpers, pure Triton key fn, tl.constexpr-instantiated constants (plain global ints fail to compile under Triton 3.7), repo-pinned ruff format. Container suite 39/39 green; topk equivalence probe 240/240 bitwise.

…e SM count Round 1 flagged this branch BLOCK; this is the planner half of it. The pinned plan refused to build once the batch exceeded the partition count and handed the step back to the shipped planner -- whose layout is a function of the batch, which is the thing this pin exists to remove. The refusal was only there because the multi-request case was untested, not because it is unsound. Partition p owns requests [floor(p*b/P), floor((p+1)*b/P)). Those spans are contiguous, disjoint and cover every request exactly once for any b: below the partition count each is 0 or 1 wide, above it a partition simply holds several whole requests. No request is ever split across partitions either way, so a row's reduction order still depends only on its own length. end_block_idx belongs to the partition's last request, which is the one index the caller fills in per call. tests: the covering property over b in (1, 3, 4, 5, 11) against P=4 -- once, contiguously, in order; idle rows left begin > end; nothing marked split; one split per request. Negative control restoring the b <= P assumption: the two above-count cases fail. 14 passed. Added to the B200 job, which lists its files individually. Still open on this branch and not claimed fixed: the cold-start step still lets the shipped planner's output through before record_num_sm_parts has seen a plan, and the C128A active_topk_width is derived from a batch-wide max_seq_len in sparse_mla.py. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

num_sm_parts is only observable from a real plan, so the very first step for a given (h_q, s_q) had to run the shipped planner -- and its result went straight into the output buffer. That layout is a function of the batch, so the first request through a fresh shape got an answer no later request would reproduce. Warmup hides it in practice, which is why it read as harmless. Record the partition count from that call as before, then rebuild the pinned plan and redo the kernel call with it. The shipped planner still runs once; its output is no longer what anyone sees. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The B200 Batch Invariance job lists vllm/v1/attention and vllm/model_executor/layers, so the planner change in sparse_swa.py is covered but the two DSv4 model files this PR also changes are not. A later edit to flashmla.py or sparse_mla.py alone would skip the suite that guards them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…ove the SM count Three gaps, all in the same direction -- code that looked verified but was not. If build_pinned_sched_meta still returns None after the shipped planner has run and the partition count has been recorded, the decode path kept whatever the shipped planner produced. That is a batch-dependent layout returned under a flag whose whole contract is that it is not, and it happened silently. Raise instead. Every GPU test stopped at batch 65 while a GB200 plans 152 partitions, so the branch that groups whole requests per partition -- the one that replaced the fallback -- had only a CPU assertion on the skeleton's shape. Add invariance and coverage tests at the partition count and above, reading the count from the cache so they follow the hardware. The cold-start rerun has no test of its control flow, which needs the model. Test the fact it rests on instead: an unobserved (h_q, s_q) yields no plan, and one observed plan makes the same shape pinnable. Also fix the annotations on _plan_skeleton and _plan_cache, which declared the second element a Tensor where it is a (busy_idx, last_owned) pair. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…overage The cold-start rerun -- run the shipped planner once because the partition count is not observable otherwise, then redo the call with the pinned plan so its batch-dependent output is never the one returned -- had no test of the control flow itself, only of the fact it rests on. It does not need a model: stub the kernel and the planner and watch the sequence. Removing the rerun makes both new tests fail, which is the point of them. Also stop gating the whole module on SM100. The backend accepts compute capability 9 and 10, so declaring the file "SM100-only" was wrong; what is true is that the kernel tests were only ever run there and the numbers in the docstring are GB200's. Gate those individually and let the structural and control-flow coverage run across the supported range. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…ct on SM90 The stub proved less than it looked. Its fake kernel never left a plan behind the way the in-kernel planner does, its fake record was a self-assignment, and its fake build returned the pinned plan on call count alone -- so deleting record_num_sm_parts, or calling it before the kernel, left the test green while the real cold start would fail closed forever. It also never checked that num_splits was installed, so dropping that line stayed green too. Model each participant instead: the kernel leaves a shipped plan on its first call, record only counts as an observation if it was handed one, and build returns pinned only after an observation. Assert both installed fields by identity. All three regressions now turn the tests red, including the mis-ordering, which reproduces exactly as the permanent fail-closed it would be in production. Separately, test_prefill_chunk_plan_batch_invariant is pure CPU metadata work and never launches FlashMLA, so it should not have carried the SM100 gate; and splitting the markers only helps if some lane collects the file below SM100. Add it to the H100 job, where the kernel tests skip and the rest run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…just against itself test_pinned_plan_does_not_drop_work checks the pinned plan value-by-value against FlashMLA's own planner, but only below the SM partition count, where each request still owns a partition. Above it the layout groups whole requests per partition -- a shape the shipped planner never emits and whose metadata columns this PR writes by hand -- and the only coverage there was invariance plus a span check on the metadata. A plan that quietly skipped KV in that regime would satisfy both. Run the same comparison at P+1 and 2P+1. Agreement is also what says each hand-written column means what the kernel reads it as, since getting one wrong moves the output. Leaving one request unowned in the last partition makes the test fail, which is the check it is there for. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

aoshen02 · 2026-08-18T15:51:57Z

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.

aoshen02 force-pushed the bi/sparse-mla branch 2 times, most recently from aca66aa to 412c3d3 Compare August 16, 2026 11:53

aoshen02 changed the base branch from bi/base to bi/indexer-topk August 16, 2026 11:53

aoshen02 and others added 16 commits August 18, 2026 15:08

aoshen02 force-pushed the bi/indexer-topk branch from fe5d557 to 14b96ea Compare August 18, 2026 15:11

aoshen02 force-pushed the bi/sparse-mla branch from 9a075db to 755c4ee Compare August 18, 2026 15:11

Navigation Menu

Repository navigation

[BI][DSv4] Pin sparse-MLA decode and prefill scheduling under batch invariance - #21

[BI][DSv4] Pin sparse-MLA decode and prefill scheduling under batch invariance#21
aoshen02 wants to merge 16 commits into
bi/indexer-topkfrom
bi/sparse-mla

aoshen02 commented Aug 15, 2026 •

edited

Loading

Uh oh!

aoshen02 commented Aug 18, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

aoshen02 commented Aug 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Perf

Tests

Review follow-up

Notes

Update: consolidation round (2026-08-15)

Stacked on bi/indexer-topk (2026-08-16)

Scope

Cost of the constant C128A width

Uh oh!

aoshen02 commented Aug 18, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

aoshen02 commented Aug 15, 2026 •

edited

Loading

Stacked on `bi/indexer-topk` (2026-08-16)