[BI][DSv4] Batch-invariant GateLinear router with fp32 persistent matmul

[aoshen02]

What

Batch-invariant GateLinear router under VLLM_BATCH_INVARIANT=1, plus a
small-N config for the shared batch-invariant persistent matmul.

The default cuBLAS route for the V4 router GEMM (4096, 256) runs
cublasLt::splitKreduce_kernel and switches algorithm across M
(ncu-verified), so row-0 routing scores change bitwise with batch size — a
discrete expert-routing flip. Under the flag GateLinear skips the tiered
dispatch and calls linear_batch_invariant directly in fp32: the
persistent matmul computes in the input dtype, and V4's router explicitly
wants fp32 scores — routing through bf16 first would round them (caught in
review, covered by a tolerance that a bf16 round-trip fails by ~40×).

batch_invariant.py gains a small-N persistent-matmul config (fp32,
N≤256 → BLOCK 64/64/128, w4s3) keyed only on (dtype, N) — M-independent
by construction, so it cannot break invariance; every small GEMM on the BI
path shares the gain.

Perf (CUDA-graph + events, GB200)

	baseline	BI before	BI after
n=1..128	8.2–20.5 us	45.1–49.2 us	28.7 us flat

≈ +0.65 ms/step over ~58 layers (was +1.7 ms). ncu: grid 2→4, kernel
89.5→43.8 us; the residual is persistent-tile residency, config choice
justified by a two-round sweep (production weight.t() layout — contiguous
layout ranks configs completely differently).

Tests

tests/v1/determinism/test_gate_linear_batch_invariant.py — 6 passed:
bitwise stability of fp32 scores across batch boundaries [1,8,15,16,17,32,
33,64], negative control on the default path that actually fails
(catches the cuBLAS split-k algorithm flip on this exact shape), and
correctness vs fp32 F.linear at reordering-level tolerance
(atol 2e-3 / rtol 1e-4; measured 8e-5 worst element).

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) found nothing to change here: batch
invariance bypasses every M-dispatched tier, bias and skip_bias_add are
preserved, requested fp32 router logits do not round through bf16, and the
fixed (dtype, N, K) persistent-matmul config does not depend on M.

Notes

• Not a duplicate: upstream BI covers the model's main GEMMs via aten
  overrides; the V4 tiered router bypasses them, and no upstream work
  touches it.
• 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_gate_linear_batch_invariant.py 6 passed,
test_matmul_batch_invariant.py 30 passed.

Update: audit round 16 (2026-08-16)

The two tests that cover the bias fold and the fp32 branch decision had never
actually been executed in the container -- the harness died on a circular
import between vllm.config and vllm.transformers_utils .model_arch_config_convertor, and separately dropped the determinism
directory's autouse fixture, so the suite would have run with batch invariance
off. Both were harness bugs, not code bugs. With them fixed the suite is
6 passed on GB200, test_matmul_batch_invariant.py 30 passed alongside it.

That also caught a silent gap in the container sync list: it did not include
vllm/model_executor/layers/batch_invariant.py, which this PR changes, so
earlier runs would have tested a half-updated tree. The sync step now warns
when a file the branch changes is missing from the list.

Update: is the M=4096 cost avoidable? (2026-08-16)

The tile config is keyed on (dtype, N, K) and never on M, so one config has
to serve every batch size, and the one picked here is faster at decode and
slower at M=4096 than the default. Three more sweep rounds, ~30 configs, on
GB200, profiler device time for the kernel alone (wall time sits on a ~300us
Triton launch floor and shows nothing), production weight.t() layout:

config	M <= 1024	M = 4096
16/32/256 w4	12.2us	150.1us
64/32/128 w4	20.6us	77.5us
32/64/256 w4	22.1us	83.5us
64/64/128 w4 (this PR)	24.4us	47.4us
64/128/64 w4	28.0us	29.2us
default 128/128/32 w8	40.8us	41.9us

No config wins at both ends, and the reason is structural: A is 4096x4096 fp32
and is re-read once per N-tile, so M=4096 wants few wide tiles, while decode is
latency-bound and wants many. 64/128/64 is the one config that beats the
default at every M -- it would turn the M=4096 number from 13% slower than
default into 30% faster -- and it costs 15% at decode. Over ~58 layers that is
+209us per decode step against -1.06ms per 4096-token prefill step.

Keeping 64/64/128: this work exists for RL rollout alignment, where decode
steps vastly outnumber prefill steps. 64/128/64 is the config to switch to for
a prefill-heavy deployment, and the numbers above are what that costs.

A correction to why the key excludes M

The comment above the branch used to say that keying on M would make BLOCK_K a
function of the batch, i.e. the reduction-order dependence this module removes.
That is not true on this path. acc = tl.dot(a, b, acc) accumulates over k
tiles in increasing k and the sum does not depend on the chunking: three configs
with two distinct BLOCK_K values give bitwise equal output at every M tested,
and so does deliberately switching the tile at M=1024 while watching one row.

The rule is still right, for a different reason: that equality is a property of
how Triton currently lowers an fp32 dot, not of the contract. bf16 and fp16 go
through tensor cores where it has no reason to hold, this helper serves every
dtype, and a regression would be silent. The comment now says that instead.

[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.