WIP - feat(cuda): CUTLASS grouped block-scaled NVFP4 MoE GEMM

[TheTom]

Summary

• Add a CUDA-backend CUTLASS grouped block-scaled NVFP4 (e2m1) MoE GEMM — cutlass_moe_fp4.cu — with f16 output and per-group device alpha pointers, exposed over the metaltile-runtime FFI as one-shot and persistent-handle (prepare / run / release) forms.
• Bake every M-independent section (expert pointers, strides, worst-case scale-factor layouts, alpha pointer array) into a persistent device blob + workspace in prepare; run is ONE small fill kernel that derives the per-group problem/operand/scale/output pointers from the router offsets on device, then gemm.run. Zero host work per call, no per-call allocations, graph-safe.
• Wire the CUTLASS_DIR build hook in build.rs to compile the fp4 unit with the block-scaled mma availability define for the target arch (sm_120a / sm_121a); the non-fp4 unit is unchanged.
• Ship the deeper TILE_K=256 config (ThreadBlockShape<128,128,256>) as the default.
• Allow the block-scaled FP format names (e2m1, ue4m3, ue8m0) in the typos config via a format-family identifier regex (the leading ue otherwise reads as a typo for use).

Motivation

The grouped NVFP4 MoE GEMM is the foundational prefill kernel for block-scaled MoE inference on the GB10 (sm_121a) CUDA path. The deeper K-tile (TILE_K=256) halves TMA load iterations for the K=2688 contraction on this weight-streaming-bound GEMM.

Originated from / validated against thewafflehaus/butter#20 (the Nemotron-Nano-30B-A3B NVFP4 GB10 prefill/decode work) — these kernels were produced and validated while testing that prefill path.

Validation

Validated on GB10 per the originating branch (CI to confirm the CUDA build):

• Bit-exact at TILE_K=256 vs TILE_K=128: argmax 1104 EXACT, cosine 0.971 unchanged.
• −8.8% kernel time: nsys cuda_gpu_kern_sum (decode-free, 20 prefill replays) grouped-problem kernel avg 2191.9us → 1999.7us/call; MoE-GEMM prefill bucket 26.6% → 24.2%.

Locally (M5 Max, no CUDA toolkit) the Rust side is green:

Test plan

• cargo fmt --check -p metaltile-runtime
• cargo check -p metaltile-runtime (default / no-cuda)
• cargo check -p metaltile-runtime --features cuda (Rust FFI typechecks; build.rs skips nvcc when CUTLASS_DIR is unset)
• cargo clippy -p metaltile-runtime --features cuda -- -D warnings
• typos --config .github/configs/typos-cli.toml
• CUDA build + grouped-GEMM bit-accuracy + nsys prefill bench on GB10 (CUTLASS_DIR set) — CI / hardware runner

[0xClandestine]

CUTLASS kernels not wired into the model executor

After checking out the branch: the CUDA codegen already exists (metaltile-codegen/src/cuda/) and all 3 backend PRs are merged — my earlier context was wrong, sorry.

The actual concern is different. moe_grouped_cutlass from #274 and this new FP4 variant both follow the same pattern: raw CudaDevice methods with no callers anywhere outside tests — not in std, not in the model executor, not in the fusion planner. Neither kernel is reachable from inference today.

The grouped NVFP4 MoE GEMM isn't a leaf utility — it sits in the middle of a full MoE forward pass: router → topk → sort → grouped_gemm(gate) → silu → grouped_gemm(up) → scatter. For it to actually run during inference the model executor needs to know it exists, which means it either needs an Op::InlineCuda escape hatch (parallel to the existing Op::InlineMsl) or a KernelEntry-registered stub the planner can dispatch by name.

If there's a wiring PR already planned, happy to merge this as-is — just want to make sure these don't get stranded as FFI stubs. If not, worth at minimum opening a tracking issue before this lands so the pattern from #274 doesn't keep compounding.

Smaller notes (still apply)

• 11 positional args mixing CUdeviceptr / host-slices / usize — a typed descriptor struct would prevent silent pointer transposition at call sites.
• release is a manual unsafe free — a Drop impl on a handle newtype makes it non-opt-in.
• Three CI checks still failing; GB10 hardware validation still [ ].

[0xClandestine]

sorry my agent tried closing this lol

[TheTom]

Addressed the grouped-GEMM dispatch concern. Opened #285 to track full executor integration for the grouped-cutlass family, and pushed a follow-up (branch tom/cuda/grouped-moe-wiring, commit 427a647) that wires both grouped GEMMs so they are reachable by registered name from the planner/executor rather than only as raw FFI methods — built and tested green on Blackwell (sm_120a):

• a small CUDA by-name dispatch seam (dispatch_grouped_moe) registers both the f16 grouped GEMM and its block-scaled NVFP4 sibling, so neither stays a test-only FFI stub;
• repr(C) argument descriptors replace the positional mixed-type arg lists, so a transposed pointer is a compile error rather than a silent wrong-buffer dispatch;
• an owning prepared-handle newtype whose Drop releases the persistent NVFP4 handle — the _release entry was previously never bound on the Rust side, so this also fixes a leak;
• the existing NVFP4 correctness test now drives the by-name seam (cos 1.000000 vs the f32 oracle), plus a new name-resolution + descriptor-validation test.

A first-class Op::InlineCuda IR node (the CUDA analog of Op::InlineMsl, for full fused-graph integration) is the remaining work, tracked in #285 — intentionally not half-built here.

GB10 validation: the NVFP4 grouped GEMM runs as the production MoE path on GB10, K256 config bit-exact vs the f32 oracle at the 12288-row x 128-expert grouped shape.

With the family-level integration tracked in #285 and the by-name wiring + ergonomics on the follow-up branch, this should be safe to merge as-is, consistent with #274.

[TheTom]

Migrating active metaltile development to thewafflehaus/metaltile per thewafflehaus/iron#12 — dev there is now the canonical branch. This draft moves to thewafflehaus/iron#13 (same three commits on that fork's dev at 14d1b30). Closing here; continuing on the fork.