fix(moe): realign mixed Trellis with QSRT ABI

[voipmonitor]

Summary

• Realign the existing mixed K3/K4 MCG launcher with the shared W4A16 ABI added by the fixed-payload QSRT work.
• Supply the new LUT and paired-M8 helper arguments at every direct call site, while keeping the established mixed checkpoint path explicitly on the MCG codebook.
• Restore qualification of the existing 512x32 ultra-wide FC2 tile used by the GLM-5.2 mixed prefill geometry.
• Add AST contract tests so future shared-kernel signature changes fail in unit tests instead of at CuTe compilation or model startup.

Root cause

master at 9bbae67 changed several shared W4A16 signatures without updating the existing mixed K3/K4 direct calls. Arguments shifted position, so active_m received a callable and compilation failed with functools.partial * Int32. The forced GLM FC2 tile was also rejected by the generic tile floor despite having a valid e4m3-k32 geometry.

This PR is deliberately narrower than #129: it restores the already-supported MCG mixed path on current master; it does not adopt the broader Fruit/SQG serving feature. If #129 is rebased, it should preserve or supersede these ABI regression contracts explicitly.

Validation

• Current master: tests/moe/test_w4a16_mixed_trellis.py -> 10 failed, 5 passed.
• This branch: mixed suite plus new ABI contracts -> 18 passed on SM120.
• Ruff and git diff --check: pass.
• Live GLM-5.2 EXL3 3.36 bpw, TP4/DCP1/MTP3: mixed tiers load, FULL and PIECEWISE graph capture complete, coherent generation passes.
• Warm synthetic C1 regression gate after the paired K6 fix: 139.5 and 144.4 tok/s; this is an MTP-friendly stability gate, not a general workload headline.

Summary by CodeRabbit

• Bug Fixes

  • Improved mixed-trellis processing for low-bit layouts and equal-bitrate tile pairs.
  • Enabled an additional FC2 specialization for supported K/32 scaling configurations.
  • Corrected codebook and dynamic pair handling for more reliable kernel execution.

• Tests

  • Added coverage to validate mixed-trellis dispatch, argument contracts, and internal kernel calls.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 0552d6e4-8a16-4e9d-974b-062bf3f6bca8

📥 Commits

Reviewing files that changed from the base of the PR and between 9bbae67 and 850bc0d.

📒 Files selected for processing (4)

• b12x/moe/_shared/kernels/w4a16/kernel.py
• b12x/moe/_shared/kernels/w4a16/mixed_trellis.py
• tests/moe/test_mixed_trellis_source_contract.py
• tests/moe/test_w4a16_mixed_trellis.py

---

📝 Walkthrough

Walkthrough

The PR updates W4A16 pair execution, mixed Trellis LUT ABI handling, MCG codebook selection, and AST-based ABI contract tests.

Changes

Mixed Trellis and W4A16 pipeline

Layer / File(s)	Summary
W4A16 pair pipeline updates b12x/moe/_shared/kernels/w4a16/kernel.py	The kernel accepts the 512×32×256-thread K/32 tile. Pair overrides now propagate through dispatch, scale loading, prefetch, MMA execution, and fragment lookahead. K-axis loading selects the low-bit record for equal bitrates or the low-bit K region.
Mixed Trellis ABI wiring b12x/moe/_shared/kernels/w4a16/mixed_trellis.py	Mixed Trellis dispatch passes zero LUT ABI values. Generated kernels use the mcg codebook.
ABI contract validation tests/moe/test_mixed_trellis_source_contract.py, tests/moe/test_w4a16_mixed_trellis.py	AST checks validate method bindings, dispatch arity, internal kernel calls, and the expanded shared MoE body arguments.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant MixedTrellisDispatch
  participant SharedMoEBody
  participant W4A16GemmKernel
  participant MCGDecoder
  MixedTrellisDispatch->>SharedMoEBody: pass zero LUT ABI values
  SharedMoEBody->>W4A16GemmKernel: dispatch paired or non-paired tiles
  W4A16GemmKernel->>MCGDecoder: decode Trellis data through MCG hooks

Loading

Possibly related PRs

• local-inference-lab/b12x#117: Related mixed Trellis ABI and dispatch argument changes.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the primary change: aligning the mixed Trellis launcher with the QSRT ABI.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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