[BI][DSv4] Make mHC tilelang ops batch invariant

[aoshen02]

What

Makes the mHC (hyper-connections) tilelang ops batch-invariant under
VLLM_BATCH_INVARIANT=1.

compute_num_split derives the K-split count from
n_sms // cdiv(num_tokens, 64), so the GEMM reduction tree changes with
the batch — the same defect class as the RMSNorm block-size fix in
upstream vllm-project#48391. use_small_fma additionally switches to a second
implementation at num_tokens <= 16. Under the flag the split count is
pinned to a batch-independent constant (min(cap, n_sms//4) rounded down
to a power of two → 32 on GB200: divides the 256 K-blocks, single wave up
to M=256) and the small-fma fork is disabled. The cross-split merge is
already a T.serial ordered loop, so pinning the count pins the tree.

Perf (CUDA-graph + events, GB200)

n	baseline	BI (s=32)
≤128	10.27 us	10.26 us (0%)
129–192	12.30 us	10.26–11.05 us (faster)
193–256	12.30 us	12.30 us (0%)

Only n=1 keeps +2 us (the fused kernel splits into post+GEMM — mechanism-
inherent). ncu: grid 192→96 single wave; the earlier +17% was attributed
by probe to the second wave turning real past grid ~150, and eliminated by
the s=32 choice.

Tests

tests/v1/determinism/test_mhc_batch_invariance.py — 10 passed: bitwise
stability across batch boundaries [1,7,8,15,16,17,63,64,65,128,129,192,
193,256] for mhc_pre / mhc_fused_post_pre / first-layer broadcast /
fused-RMSNorm variants, mhc_post and hc_head regressions, and negative
controls that actually fail — with checkpoint-realistic magnitudes
(hc_attn_scale≈[2.08,0.019,0.245]): synthetic *1e-4 weights wash the
10-ULP reordering signal below fp32 resolution and made the first control
pass vacuously.

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 the fail-closed discipline had a
hole, and it was verified against the source before acting:
tf32_hc_prenorm_gemm is DeepGEMM's and this path has no other
implementation of it, so a deep_gemm without set_batch_invariant leaves
it free to select its config from the batch. Disabling DeepGEMM MoE, the
loader's answer to such a build, does nothing here. Every entry point that
reaches the kernel now refuses — mhc_pre_tilelang,
mhc_pre_broadcast_tilelang and mhc_fused_post_pre_tilelang, one call each —
behind a functools.cache so the per-layer path pays the probe once. This adds a dependency on
deep_gemm_batch_invariant_enabled() from #24.

Notes

• Not a duplicate: SGLang's mHC implementation has the same
  compute_num_split defect and no deterministic branch; nothing upstream
  covers mHC BI.
• 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_mhc_batch_invariance.py 15 passed.

Scope (added 2026-08-16 after audit)

This validates CUDA TileLang mHC batch invariance on SM100 (GB200). The change
itself is correctness-first and platform-independent: it stops the two fast
paths from switching kernels on the token count, which is what made a row's
result move with the batch. That applies wherever TileLang runs, including
ROCm -- limiting it to CUDA would leave ROCm batch-variant under the flag,
which is the worse outcome.

What is not claimed:

• No numerical or performance evidence on ROCm TileLang/AITER. The new tests
  are all skip_if_not_cuda.
• No performance claim for SM120 or any other Blackwell part. The generic
  schedule is used there for correctness; the GB200 numbers are not
  extrapolated to unmeasured hardware.
• ROCm AITER is not covered at all. Production ROCm prefers AITER when
  HAS_AITER_MHC and hidden_size % 256 == 0, so this PR must not be read as
  "ROCm mHC is batch invariant".
• No real-model end-to-end evaluation. This is a kernel-level correctness fix
  and the regressions are kernel-level.

Added after the audit:

• The production fused-norm entry (mhc_fused_post_pre_tilelang with
  norm_weight) now has an invariance sweep, a negative control, and a
  reference composed from the kernel suite's own post/pre references plus an
  fp32 RMSNorm. Its epilogue is not shared with the norm-free entry.
• The non-DeepGEMM prenorm GEMM's own dispatch is now exercised directly, since
  B200 always takes DeepGEMM and nothing else in the file reached it. Of its two
  branches only the < 128 one moves a row's bits; the >= 1024 block-M
  variant keeps the same n_thr and tile_n, reassociates K in the same order,
  and comes out bitwise equal at DSv4's shape. It stays pinned anyway, and the
  negative control says which of the two it can actually prove.

tests/v1/determinism/test_mhc_batch_invariance.py is now listed in the B200
Batch Invariance job; it was previously not executed by any job, since that job
names its files individually.

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