feat(moe): rp-native Triton tiny-decode path for b12x w4a8_mx (M=1) — DS4 A8 decode 133.7 → 138.7 tok/s

[voipmonitor]

What

New b12x_tiny_decode.py: two Triton kernels + maybe_run_tiny_w4a8mx_moe, hooked at the top of _run_b12x_moe_fp4. At M=1 with quant_mode=w4a8_mx (silu, w31 layout, shape multiples of 256) it replaces the b12x dynamic grouped kernel; everything else falls through unchanged. Env-gated, default off: VLLM_B12X_W4A8_MX_TINY_DECODE=1.

The kernels read the N256/K128 in-place-repacked FP4 weights and e8m0 sfb grids directly via the verified inverse bit mappings (b12x PR local-inference-lab/b12x#21, tests/test_w4a8_rp_inverse_mapping.py) — zero extra weight memory. Activations stay BF16 (no input quant), accumulation fp32: FC1+SiLU staged in a 48 KB fp32 scratch, FC2 scatters router-weighted fp32 atomics.

Stacked on #69 (retarget to dev/eldritch-enlightenment after it merges).

Why

The whole DS4 A8-vs-A16 decode gap (~130 vs ~140 tok/s) is the MoE at M=1: w4a8_mx tiny decode is hard-wired to the dynamic kernel (34.8 µs/layer) because the in-place repack blocks the fused path (22.5 µs). All config/env knobs around it are falsified — full analysis in b12x-dense-fp8-gemm-vs-deepgemm.md §6.8.

Measured

• Isolated (graph replay, DS4 TP2 shapes, M=1): 28.8 µs/layer vs dynamic 34.8 (+~3 µs wrapper it also bypasses).
• Numerics: cos vs fp32 oracle 0.999999 (dynamic kernel: 0.9990 — it quantizes activations); cos vs dynamic 0.999.
• E2E DS4-Flash TP2 A8 (hybrid DG-linear base): decode cc1 138.7–138.8 tok/s (baseline 133.7; 4 runs across 3 restarts), prefill unchanged (12,202 @128k vs 12,182 ref), 30k coherence clean.
• Ready-to-run image: voipmonitor/vllm:eldritch-enlightenment-v8e6e417c-b12x77bd50e-tinymoe-overlay-cu132-20260702.

Notes for review

• Second commit reverts sem='relaxed' on the atomic scatters: isolated-neutral but −7 % E2E (131.3). PDL overlap (gdc_wait/launch_dependents) measured the same −7 % and is deliberately absent. Both documented in the kernel wiki page with the full optimization journey (flat-load vectorization is the load-pattern that matters; num_stages/num_warps traps listed).
• Buffers pre-allocate on the first eager call (b12x prewarm covers this before capture); if somehow first hit during capture, it logs once and falls back to the dynamic path.
• M=2–4 intentionally not routed (kernel re-reads weights per routed row; dynamic wins there). cc1 serving decode is M=1.
• Longer term this band belongs in the b12x dynamic-kernel family (CuTe DSL, true FC1+FC2 fusion, target ~22–24 µs) — this PR is the working floor/reference until then.

🤖 Generated with Claude Code

[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: ce7f1509-107b-455e-8289-0aa0c17f81c4

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fable/b12x-w4a8mx-tiny-decode-20260702

---

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

[voipmonitor]

Superseded per Luke's feedback by the b12x-native CuTe DSL version: local-inference-lab/b12x#22 — same design, integrated into the normal b12x kernel selection (B12X_W4A8_TINY_RP=1), zero vLLM changes needed, and faster: 20.6 µs/layer vs this PR's 28.8 (full-path), E2E decode 140.2 tok/s (A16 parity) vs 138.7 here. Closing; keeping the branch for reference — the Triton prototype journey that led to the layout tricks is documented in the wiki.