moe: bound W4A16 scratch and repair tail correctness

[malaiwah]

Summary

Bound unpinned W4A16 scratch planning to route block sizes reachable within the declared token capacity, mirror the correction in replicated-input EP, and expose a cheap public arena-sizing query.

This is stacked on #92 because the fixed-capacity Trellis planning API under review there is the source base. It addresses the safe planner work from #93; it does not clamp Trellis capacity or implement caller-side prefill chunking.

Related: #93

Problem

fused_moe and ep_moe sized unpinned W4A16 route buffers and both fp32 c_tmp planes across every allowed block size (8, 16, 32, 48, 64), even when runtime selection cannot reach most of them within max_tokens.

At max_tokens=64, topk=8, E=160, SM count 188, hidden 6144, intermediate/rank 512, the current full sweep prices the non-full W4A16 arena at 94.94 MiB. Runtime can select only block 8, whose exact arena is 30.03 MiB: 64.91 MiB/rank of unreachable reserve.

The review also exposed two sizing hazards that a naive narrowing would make visible:

1. packed direct-topk/TC-decode uses m * topk * block_m scratch geometry instead of packed-route geometry;
2. non-Trellis binding cannot carry an expert map, so core route packing always uses weight_E, even when the router workspace uses a smaller route_num_experts.

Both are covered explicitly rather than relying on the old oversized sweep.

Change

• Add route_block_sizes_for_capacity(max_tokens, topk, num_experts) using the same monotone selector as runtime.
• Size unpinned fused and EP route buffers only across that conservative reachable interval.
• Size fused c_tmp independently for both packed-route and every reachable direct-topk/TC-decode (m, block) geometry.
• Use weight_E for every non-Trellis W4A16 core route allocation; retain requested route_E for mapped full-rotation Trellis.
• Preserve explicit block pins exactly.
• Export fused_moe.required_nbytes(Caps), which returns the complete caller-owned arena requirement without launch compilation or retained storage.

Safety invariants

• Kernel selection, launch geometry, route packing, and numerical math are unchanged.
• Every runtime-selected block for 1 <= m <= max_tokens is included.
• Block-8 doubled fp32 scratch accounting remains intact.
• EP still reserves power-of-two route-pack capacity before applying the reachable block set.
• Direct-topk affects only c_tmp; packed route-index and block-expert buffers are not inflated by direct geometry.
• Non-Trellis core expert_offsets and route buffers match the expert count the runtime actually packs.
• Full-rotation Trellis mapped-route behavior remains unchanged.

Additional correctness value found during GPU qualification

The 64.91 MiB/rank planner reduction is one value axis. GPU qualification also found four pre-existing W4A16 correctness/memory-safety defects that are worth landing independently of how reviewers value that memory saving. They were not introduced by the planner patch: initial commit fa7a6ad changed host-side planning/allocation, while the affected micro.py and prepare.py implementation was unchanged from base 06032ce.

The smaller, correctly bounded arena made the first latent lifetime violation reproducible; poisoned replay, boundary, and Compute Sanitizer tests then exposed the remaining inherited defects:

inherited defect	observed consequence
stale direct-FC2 scratch lanes	deterministic non-finite output on same-binding M=1/M=3 ReLU2 and SiLU replay
truncated 64-row ModelOpt scale tails	wrong numerical output for supported non-64 ReLU2 intermediates such as I=32/224
coarse non-64 W/scale bounds and E8M0 scale stride	out-of-bounds reads; the uncorrected source produced 278 invalid global reads in the exact-source sanitizer control
wide activation-tail validity derived from padded rather than logical rows	stale/NaN consumption for supported I=464/480/496 residue classes

The current GLM I=1856 production shape is aligned and does not exercise those tail cases; matched A/B is therefore byte-identical and throughput-neutral there. The generic fixes still remove real wrong-output/OOB behavior for supported W4A16 shapes. In short: the planner change recovers memory, while the later commits make the inherited kernel implementation safe and correct.

Observed sizing

CPU-isolated query against the r12 image dependencies and this source:

auto=31491908 (30.03 MiB)
block8=31491908 (30.03 MiB)
block64=99553972 (94.94 MiB)
recovered_vs_worst=68062064 (64.91 MiB)
trellis_3072_required=1106622184 (1055.36 MiB)

The shipped vLLM EXL3 integration pins both decode and prefill block sizes, so the 64.91 MiB correction does not change that deployment directly. The new sizing query is the prerequisite for pricing the separate caller-side prefill-capacity tradeoff.

Verification

All final validation ran in a separate CPU-only container (--runtime crun, NVIDIA_VISIBLE_DEVICES=void, empty CUDA_VISIBLE_DEVICES) using the exact pulled r12 image dependencies:

32 passed, 80 skipped

Covered suites:

• test_fused_moe_planning.py
• test_fused_moe_trellis.py
• test_fused_moe.py
• test_ep_moe_api.py
• test_w4a16_route_pack.py

Changed-file Ruff lint and git diff --check pass.

Independent review initially found the direct-topk and non-Trellis expert-count hazards above. Regression tests were added for both exact cases; after both corrections the final verdict was APPROVE, with no remaining actionable defect. Residual risk is live-GPU execution, which is intentionally not claimed for this host-side planner-only change while the production deployment remains active.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: dd9e4cb7-8fed-4234-9171-103239bfff90

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[malaiwah]

Rental-machine validation — repaired composition now PASS

Initial tested head: fa7a6ad4e8bbcf241661009697a879630cddb554
Current tested head: c2f135c66032ac5e9d0778067dd19ae1910cff47
Tested base: 06032ce74afba49a683d01339ed2c971568cec3f (PR #92 branch)

Host: Vast.ai 46335896; 4x RTX PRO 6000 Blackwell, all NODE/NUMA 0 with real P2P copies passing.
Driver/CUDA/PyTorch: 610.43.02 / CUDA UMD 13.3 / PyTorch 2.12.0+cu132.

Initial result — isolated PASS, composition FAIL

• Required CPU MoE union: 32 passed, 80 skipped; Ruff clean.
• Same affected suites with one GPU: 112 passed cold (68.87 s), 112 passed warm (7.61 s).
• Exact planner delta: 68,062,064 bytes = 64.90904 MiB/rank of unreachable reserve removed.
• Clean moe: bound W4A16 scratch and repair tail correctness #100 + fix(moe): reject unresolved capture launches #102 initially failed two ordered W4A16 GPU cases with NaN oracle metrics. The smaller arena exposed stale intermediate_cache2 reuse.

That failure was real and promotion was blocked. The four defects found during the repair campaign were inherited W4A16 correctness/memory-safety bugs, not bugs introduced by the planner PR. Initial commit fa7a6ad changed host-side planning/allocation; the affected micro.py and prepare.py code was unchanged from base 06032ce. The smaller arena made stale direct-FC2 lane consumption deterministic, and the expanded poisoned/boundary/sanitizer matrix then exposed three more latent base defects: truncated 64-row ModelOpt scale tails, exact non-64 W/scale/E8M0 bounds, and wide activation-tail residue at I=464/480/496. Their consequences were non-finite or wrong output and, in the uncorrected exact-source sanitizer control, 278 invalid global reads. These four repairs have correctness value even if 64.91 MiB/rank seems modest; two independent reviews rejected intermediate candidates before the final source/test stack was approved.

Current exact-head result — PASS

• Ruff/format/diff: clean.
• Contract + poisoned eager/16-replay graph + bounds matrix: 88/88.
• Bounds matrix under Compute Sanitizer: 24/24, 0 errors.
• Broader pure-moe: bound W4A16 scratch and repair tail correctness #100 file: 196 passed, 16 skipped, with one known inactive-relu2 helper failure fixed by fix(moe): reject unresolved capture launches #102.
• Exact moe: bound W4A16 scratch and repair tail correctness #100 + fix(moe): reject unresolved capture launches #102 composition cae3aad137bd739881833c064e1912177698156f: 220 passed, 16 skipped cold in 171.98 s and warm in 5.91 s.

The planner saving remains 64.90904 MiB/rank. A matched three-run 600 W A/B on aligned GLM-like I=1856 shapes found byte-identical output/allocation/peak-memory and latency deltas from -0.084% to +0.068%; no throughput claim is made for the safety repair.

Full chronology, exact commands, raw XML/logs, sanitizer evidence, and A/B JSON: repair ledger at d4b0a90.

Exact r14 deployable-source replay

The immutable r14 SparkInfer release tree was independently reconstructed from blackwell-llm-docker@2464cc03: base 9b852b281250123fe323f63ccb1df3cac0f3bbca plus the release lock produced the published tree 8110e3ea417794bfb08aff1fba20135102e5536b exactly. Applying the accepted #100/#102 repairs produced candidate e2205cba8d78db03427a3945a75a1f647b51ef5a (tree 1d378d599538daecfc1e8c11a44010b27d3ddfe5).

On physical GPU 2 with a fresh candidate-specific extension cache, the complete seven-file union passed 332 passed / 16 intentional skips cold in 237.03 s and identically warm in 10.03 s. This closes the image-source composition boundary for the W4 repair; no new numerical, compile-cache, capture, or mixed-K regression appeared.

Exact source provenance and immutable raw logs: r14 reconstruction ledger, cold log, warm log.

[malaiwah]

Final cross-stack evidence for exact head c2f135c: the byte-exact planner query confirms 68,062,064 bytes, or 64.90904 MiB/rank, of unreachable W4A16 reserve removed. The repaired #100 plus #102 composition passed 88/88 boundary/poison tests, 24/24 Compute Sanitizer cases with zero errors, and 220 passed/16 skipped cold and warm; matched production-shape latency was neutral within noise. This 64.909 MiB is an independent sizing claim, not yet a measured combined full-model delta with vLLM #210. Final evidence: https://github.com/malaiwah/glm52-exl3-vast/blob/5c76a2536e7fc9a5f1cb6bf182531889f5385e65/docs/field-review-results/2026-07-30-vast-46335896/UPSTREAM-REPAIR-CAMPAIGN.md Reproduction: https://github.com/malaiwah/glm52-exl3-vast/blob/5c76a2536e7fc9a5f1cb6bf182531889f5385e65/docs/field-review-results/2026-07-30-vast-46335896/COUNTER-VALIDATION.md

[malaiwah]

Review-base note

This PR still targets fix/trellis-stable-profile-20260730, whose #92 content reached master through a non-ancestral merge/squash. I checked the graph before changing the base: a naive retarget is not patch-neutral. It changes the visible series from 16 to 19 commits and adds tests/moe/test_tp_moe_scratch_bindings.py to the hypothetical three-dot diff.

I am therefore leaving the tested head/base untouched instead of manufacturing a misleading clean check. Reviewers should use exact head c2f135c66032ac5e9d0778067dd19ae1910cff47 and the immutable manifest/patch digest in the counter-validation guide. A clean master replay should be a separately range-diffed successor, not an in-place base flip.

Guide: https://github.com/malaiwah/glm52-exl3-vast/blob/d821001f97dd15f0ae6ace626cc1f334a2069553/docs/field-review-results/2026-07-30-vast-46335896/COUNTER-VALIDATION.md

[voipmonitor]

Superseded by #109. The original PR was stacked on an obsolete pre-#92 base. #109 replays only the unmerged planner and tail-correctness commits directly onto current master, with the same focused test result and no already-merged history.