feat(moe): rp-native tiny-decode kernel for w4a8_mx (M=1) — A8 decode reaches A16 parity

[voipmonitor]

Per your feedback on the vLLM-side attempt (local-inference-lab/vllm#70): here's the tiny-decode band as a b12x CuTe DSL kernel integrated into the normal kernel selection — no vLLM changes needed at all.

What

b12x/moe/fused/tiny_rp.py — MoETinyRpKernelBackend, micro-family style, two plain launches:

• FC1: gate+up dots into an fp32 intermediate. Thread (n8c, r8, cgrp) issues one 16 B ld.global.nc.v4 per k32 covering the four 8-apart logical rows of an rp word quad → warps read fully coalesced 512 B runs of the N256/K128 in-place-repacked storage directly (inverse mappings from tests/test_w4a8_rp_inverse_mapping.py); a 4-lane butterfly folds k.
• FC2: inline SiLU from the intermediate, fp4_dot4 hardware decode (f16-pair math, same numerics class as w4a16), router-weighted scatter_add_bf16x2 epilogue straight into the bf16 output.

Activations stay BF16 (no input quantization — at M=1 they're 8 KB; skipping the MXFP8 quant removes work and beats the dynamic path's numerics: cos vs fp32 oracle 0.999980 vs 0.9990).

Integration: _resolve_workspace_layout routes the w4a8_mx tiny band to the compact family when B12X_W4A8_TINY_RP=1 (default off); _launch_micro dispatches to the new kernel; _get_weight_views flat-views rp storage (and never re-flips halves in place). Two new fp4 primitives: red_add_global_f32, cvt_bf16x2_to_f16x2.

Measured (DS4-Flash TP2 shapes, RTX PRO 6000)

	µs/layer @ M=1 (full b12x_moe_fp4 graph)
w4a8_mx dynamic (today)	37.8
w4a16 fused reference	22.5
tiny_rp (this PR)	20.6

E2E serve (hybrid DG-linear base): decode cc1 133.7 → 140.2 tok/s, i.e. A8 now matches the A16 reference (140.5, ITL 7.131 vs 7.136 ms) while keeping the A8 prefill advantage (128k prefill unchanged, 12.1k tok/s; 30k coherence clean, CJK 0).

Notes

• A first single-launch cooperative variant (grid barrier, micro-style) deadlocked under TP2 serving — rank-skewed first-use compile broke CTA co-residency while the resident CTAs spun. The landed two-launch form has no co-residency assumptions and no barrier state; it also measured faster (20.6 vs 30.8 µs).
• M gate is 1 for now (weights re-read per routed row; the dynamic kernel stays better at M≥2). configure() accepts m≤4 for future batched-decode experiments.
• Follow-ups if you want them: fold the two wrapper zeros into the kernels, per-expert FC2 early start, MTP-friendly M=2–4 batching.
• Full working log incl. the Triton prototype journey and falsified variants: https://github.com/voipmonitor/rtx6kpro/blob/master/optimization/b12x-w4a8mx-tiny-decode-kernel.md

Builds on #21 (the verified inverse-mapping test this kernel consumes).

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added support for a new tiny MoE execution path for specific small-shape workloads, improving backend flexibility and performance.
  • Expanded low-level numeric support with new packed BF16/FP16 conversion and FP32 reduction handling.
• Bug Fixes
  • Improved weight layout handling for certain repacked models to better match expected runtime shapes.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 9451e279-2e1e-4235-9aad-2c4077131501

📥 Commits

Reviewing files that changed from the base of the PR and between 8de17f1 and 5f80a9a.

📒 Files selected for processing (3)

• b12x/cute/fp4.py
• b12x/integration/tp_moe.py
• b12x/moe/fused/tiny_rp.py

---

Walkthrough

Adds two PTX helper intrinsics (global FP32 reduction, BF16x2→F16x2 conversion), a new CUTLASS/CUTE-backed MoETinyRpKernelBackend two-phase kernel for tiny W4A8-MX MoE decode, and tp_moe integration wiring (workspace sizing, weight-view reshaping, backend routing, custom-op launch) to route qualifying tiny-decode shapes through this new "tiny_rp" path.

Changes

Tiny-RP W4A8-MX decode path

Layer / File(s)	Summary
New PTX intrinsics b12x/cute/fp4.py	Adds red_add_global_f32 (relaxed-scope global FP32 reduction) and cvt_bf16x2_to_f16x2 (packed BF16x2→F16x2 conversion) PTX helpers.
MoETinyRpKernelBackend kernel b12x/moe/fused/tiny_rp.py	Implements a two-phase (FC1/FC2) CUTLASS/CUTE kernel backend for M<=4 SM120 decode: constructor/configure() validation, _row_block_dot() packed-weight decode helper, FC1 dot-partial accumulation into an FP32 intermediate, FC2 SiLU/gating and bf16 scatter-add output, a __call__ tensor-building wrapper, and a static launch() entry point.
tp_moe integration and routing b12x/integration/tp_moe.py	Adds B12X_W4A8_TINY_RP-gated capability checks, a cached kernel compiler, a flat launch wrapper, and a tp_moe_tiny_rp_launch custom op (real + fake); wires workspace sizing, weight-view reshaping for W4A8-MX, backend layout resolution, and _launch_micro to route supported shapes to the new tiny-rp launch path.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant LaunchMicro as _launch_micro
  participant Op as tp_moe_tiny_rp_launch (custom op)
  participant FlatWrapper as _launch_tiny_rp_flat
  participant Backend as MoETinyRpKernelBackend
  participant Inter as inter_fp32 buffer
  participant Out as output tensor

  LaunchMicro->>LaunchMicro: check w4a8_mx + tiny_rp enabled/supported
  LaunchMicro->>Op: invoke with weight pointers, ids, barriers
  Op->>FlatWrapper: dispatch launch
  FlatWrapper->>Backend: MoETinyRpKernelBackend.launch(...)
  Backend->>Inter: zero + write FC1 router-weighted partials
  Backend->>Out: zero + FC2 scatter_add bf16 results
  Backend-->>LaunchMicro: kernel execution complete

Loading

Compact metadata

• Related issues: None provided
• Related PRs: None provided
• Suggested labels: None provided
• Suggested reviewers: None provided

Poem
A rabbit taps four rows so tight,
FP4 decoded, bf16 in flight,
FC1 whispers into inter's floor,
FC2 sings SiLU, scatters more,
Tiny-RP hops through the MoE night. 🐇

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

[voipmonitor]

Done in de15849: renamed everything tiny_rp → tiny_decode ("rp" was shorthand for the in-place-repacked weight storage from _logical_weight_to_w4a8_rp_inplace; now spelled out in the docstring and gone from public names) and flipped the default to on — w4a8_mx M=1 decode uses the new kernel out of the box, B12X_W4A8_TINY_DECODE=0 is the kill switch. Re-verified both ways on the binding path (20.6 µs engaged / 34.8 µs dynamic with the kill switch) and E2E with no env at all: 140.16 tok/s decode cc1.