fix(w4a16): admit ultra-wide FC2 tile in forced re-pin validation

[sjug]

Summary

The W4A16 TC-decode auto selection can pick the ultra-wide FC2 tile (tile_k=32, tile_n=512) (the wave-balance override in compile_w4a16_fused_moe), which is admitted via a direct shared-memory footprint check because tile_k=32 sits below the generic tile_k>=64 floor in _candidate_tile_fits — the override's comment notes this explicitly.

However, when the selected tiles cross the torch.ops.b12x.w4a16_fused_moe_launch custom-op boundary, the launch side re-pins them via force_tile_config, and that validation path re-checks with the generic _candidate_tile_fits — which rejects the exact shape the auto selection just produced:

ValueError: force_tile_config fc2 tile (tile_k=32, tile_n=512) does not fit problem N/K=4096/1024 at moe_block_size=8

Why upstream never sees it

The ultra override only fires when its wave-balance gates hold, and those compare FC1/FC2 mn-tile counts against the device SM count (fc1_mn_after <= sms, etc.). On 188-SM SM120 parts (RTX PRO 6000) the gates don't fire with typical model dims. On 48-SM GB10 (DGX Spark) they do: DeepSeek-v4-Flash TP2 decode selects the ultra tile and crashes during CUDA graph capture on every boot.

Fix

Teach the forced re-pin path the same exception the selection path makes: accept fc2 (tile_k=32, tile_n=512) under the identical direct footprint check (exact N/K coverage, scale-group alignment, smem fit). No behavior change for any config the generic floor already admits; the tiles still land in the GEMM cache keys as before.

Testing

• Repro/verify on 2x GB10 (SM121, aarch64): DS4-Flash TP2 with the b12x-a16 backend crashed at graph capture on every boot before the patch; with it, the server boots, serves, and the ultra-tile kernel compiles and runs (currently under a multi-hour benchmark campaign).
• Not reproducible on SM120 hardware (gates never fire there), so no regression path for existing deployments.

Summary by CodeRabbit

• Bug Fixes
  • Improved support for ultra-wide FC2 workloads using specific forced tile configurations.
  • Prevented valid configurations from being incorrectly rejected during fused MoE processing.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: bf146bc6-7ea8-42e4-9489-ffbea11909b7

📥 Commits

Reviewing files that changed from the base of the PR and between 2433500 and 9aa1f07.

📒 Files selected for processing (1)

• b12x/moe/fused/w4a16/kernel.py

---

Walkthrough

The forced tile validation in compile_w4a16_fused_moe now permits a specific FC2 TC-decode ultra-wide tile configuration when its alignment, divisibility, and shared-memory constraints are satisfied.

Changes

W4A16 tile validation

Layer / File(s)	Summary
Admit constrained FC2 ultra-wide tiles b12x/moe/fused/w4a16/kernel.py	The force_tile_config validation loop bypasses _candidate_tile_fits for FC2 tiles with tile_k=32 and tile_n=512 when the required constraints pass.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 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 and concisely describes the main change: allowing the ultra-wide FC2 tile through forced re-pin validation.
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.}

[sjug]

Superseded by #39 — identical commit, renamed branch.