fix(gemm): restore capture-safe K6 small-M dispatch

[voipmonitor]

Summary

• Restore the established cooperative K6/MCG dense path for M <= 128.
• Keep the gate narrow: K6, MCG, unpaired payload, FP16 compute, and no external H128 callback. SQG/QSRT, paired payloads, BF16, and externally rotated calls remain on the generic scheduler.
• Reuse caller-owned output and rotation buffers, so vLLM decode graph capture does not require the generic W4A16 GEMM scratch arena.

Root cause

master at 9bbae67 removed the K6/MCG small-M dispatch while retaining its implementation and CUDA-graph regression test. vLLM intentionally provisions only a one-element generic scratch placeholder for this qualified small-M path. Falling through to the generic scheduler therefore fails during graph capture with W4A16 GEMM scratch is not initialized for CUDA graph capture.

The fix restores the original ownership contract rather than increasing persistent scratch or disabling CUDA graphs.

Validation

• Current master: test_k6_small_m_cuda_graph_replay_is_stable fails during capture.
• This branch: all K6 small-M tests -> 17 passed, 48 deselected on SM120, including numerical comparison against the generic CuTe path and CUDA-graph replay.
• Ruff and git diff --check: pass.
• Live GLM-5.2 EXL3 3.36 bpw, TP4/DCP1/MTP3: FULL and PIECEWISE graph capture complete and coherent generation passes when combined with the mixed-ABI repair.
• Warm synthetic C1 regression gate: 139.5 and 144.4 tok/s.

Summary by CodeRabbit

• Performance Improvements
  • Improved quantized matrix multiplication tile selection for more efficient low-bit layouts.
  • Added a specialized path for small FP16 inputs to accelerate K6 MCG operations.
  • Reduced unnecessary metadata, transformations, and scheduling for eligible workloads.

[coderabbitai]

Warning

Review limit reached

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

Next review available in: 52 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

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: 8553e1f7-55f6-4e21-9029-d432fb899c70

📥 Commits

Reviewing files that changed from the base of the PR and between 6d2109b and b3f572d.

📒 Files selected for processing (2)

• b12x/moe/_shared/kernels/w4a16/kernel.py
• tests/gemm/test_trellis_linear.py

📝 Walkthrough

Walkthrough

The W4A16 kernel combines Trellis low-bit tile selection conditions and adds a specialized dense K6 MCG path for small FP16-compatible workloads.

Changes

W4A16 kernel execution

Layer / File(s)	Summary
Trellis tile selection b12x/moe/_shared/kernels/w4a16/kernel.py	Equal-bitrate and early-K cases now use one low-bit Trellis selection condition. Later K tiles retain high-bit selection.
K6 MCG dense execution b12x/moe/_shared/kernels/w4a16/kernel.py	Small FP16-compatible workloads prepare FP16 buffers, call run_k6_mcg, convert non-FP16 outputs, and return before generic GEMM processing.

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

Sequence Diagram(s)

sequenceDiagram
  participant DenseExecutionBranch
  participant FP16Buffers
  participant run_k6_mcg
  DenseExecutionBranch->>FP16Buffers: Prepare or reuse FP16 buffers
  DenseExecutionBranch->>run_k6_mcg: Pass reshaped Trellis payload
  run_k6_mcg-->>DenseExecutionBranch: Return K6 MCG results
  DenseExecutionBranch->>FP16Buffers: Convert non-FP16 outputs

Loading

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 and concisely describes the restoration of capture-safe K6 small-M dispatch, which is the main change in the pull request.
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.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@b12x/moe/_shared/kernels/w4a16/kernel.py`:
- Around line 11329-11336: Update the use_k6_mcg_small predicate near run_k6_mcg
to require the current device’s supported sm_120 capability before selecting
this path. When that capability is absent, the predicate must be false so
execution falls back to the generic scheduler instead of invoking run_k6_mcg.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e0cc5b93-83f9-4443-b963-1f6c56336dba

📥 Commits

Reviewing files that changed from the base of the PR and between 9bbae67 and 6d2109b.

📒 Files selected for processing (1)

• b12x/moe/_shared/kernels/w4a16/kernel.py