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

[sjug]

Summary

W4A16 TC-decode can auto-select the ultra-wide FC2 tile
(tile_k=32, tile_n=512) on 48-SM GB10 geometry. Before this fix,
re-pinning that selected configuration across the custom-op boundary rejected
it because the generic tile validator enforces tile_k >= 64.

The result was a deterministic graph-capture boot failure for
DeepSeek-v4-Flash TP2 on DGX Spark even though sparkinfer itself had selected
the tile.

Root cause

The ultra-wide FC2 override intentionally validates (32, 512) outside
_candidate_tile_fits, because it is the one supported TC-decode tile below
the generic tile_k >= 64 floor. The force_tile_config path did not share
that exception and therefore rejected the selector's output.

Fix

• Forward-port the fix onto the current namespaced W4A16 implementation.
• Extract one _tc_decode_ultra_fc2_tile_fits predicate and use it in both
  automatic selection and forced re-pin validation.
• Keep exact N/K coverage, scale-group alignment, and shared-memory capacity
  checks.
• Preserve generic validation for every other tile configuration.

Regression coverage

test_w4a16_tc_decode_ultra_fc2_tile_round_trips_forced_repin:

• simulates the 48-SM GB10 wave-balance geometry;
• asserts auto-selection produces
  (fc1_k=64, fc1_n=256, fc2_k=32, fc2_n=512);
• feeds that complete configuration back through force_tile_config and
  asserts it is preserved;
• verifies the exception remains rejected when the FC2 problem N is not
  512-aligned.

Validation

• Before fix: exact reproduction selected (64, 256, 32, 512) and raised
  ValueError during forced re-pin.
• After fix: exact reproduction accepts the re-pinned configuration.
• pytest -q tests/moe/test_w4a16_e2e.py::test_w4a16_tc_decode_ultra_fc2_tile_round_trips_forced_repin
  — 1 passed.
• pytest -q tests/moe/test_w4a16_e2e.py -k tc_decode
  — 14 passed, 79 deselected.
• ruff check sparkinfer/moe/_shared/kernels/w4a16/kernel.py tests/moe/test_w4a16_e2e.py
  — passed.
• git diff --check — passed.

The original deployment evidence remains: on 2× GB10, DeepSeek-v4-Flash TP2
failed graph capture on every boot before the fix and booted and served
successfully with the ultra tile admitted.

[coderabbitai]

Warning

Review limit reached

@lukealonso, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 4 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c7d1a9c8-189b-4156-af36-b04a56a573be

📥 Commits

Reviewing files that changed from the base of the PR and between cc9b476 and adcc6aa.

📒 Files selected for processing (2)

• sparkinfer/moe/_shared/kernels/w4a16/kernel.py
• tests/moe/test_w4a16_e2e.py

📝 Walkthrough

Walkthrough

The forced tile configuration validation now recognizes a TC-decode FC2 ultra-wide shape when its geometry, divisibility constraints, and shared-memory footprint match; other configurations continue through the existing candidate-fit validation.

Changes

FC2 tile validation

Layer / File(s)	Summary
Ultra-wide FC2 validation bypass b12x/moe/fused/w4a16/kernel.py	Adds a guarded bypass for the FC2 tile_k=32, tile_n=512 configuration and preserves _candidate_tile_fits(...) validation for other shapes.

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 summarizes the main fix: accepting the ultra-wide FC2 tile during 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.}