cuda: EXL3 linear and grouped routed experts, every codebook and width, a width per expert (#593)

[justincasey]

What this changes

EXL3 on the native CUDA engine (#593): hand-written kernels for the linear and the grouped routed experts, at every codebook (mul1, mcg, 3inst) and width. Each routed expert keeps its own width.

• exl3_decode.cuh, exl3_linear.cu, cuda/exl3.zig:
  • the trellis decode straight into mma.m16n8k16 B fragments, and the input rotation;
  • a row-invariant linear for 1 to 128 rows (the mul1 half steps included), with the K split planned from the shape alone;
  • prompt rows in 128-row blocks;
  • an .exl3 arm in qlinear.
• exl3_experts.cu, cuda/exl3_experts.zig: grouped experts (member lists, each expert's rotation, gate and up, GLM's SwiGLU roundings, down, the weighted combine).
  • Widths can mix within a layer: vcruz305/RED-SNOW-5.3-FLASH-EXL3-SAGE-4.91bpw mixes 4- to 8-bit experts in 41 of its 43 MoE layers.
  • An expert's member count never changes a row's bits.
• core/exl3_load.zig, cuda/exl3_upload.zig: a checkpoint's trellis groups from its safetensors headers, and their upload to the device.
• tf-cuda-test exl3 [DIR] and exl3-tp2 RANK HOST:PORT [DIR], plus a section in docs/recipes/exl3.md.

No family loads these yet. One commit on main (c126fe16), with no conflicts.

Receipt

• Environment: this head on c126fe16; DGX Spark (GB10, sm_121), CUDA 13.0; checkpoint vcruz305/RED-SNOW-5.3-FLASH-EXL3-SAGE-4.91bpw.
• Exactness:
  • Float64 references built on core/exl3_format.zig: the linear at every width and codebook, and the experts at 4/5/6 bits and mixed widths, each against a worst-element bound.
  • A row's bytes are equal across row counts 1 to 128, K splits, expert member counts and prompt chunking.
  • With DIR, real RED-SNOW layers byte for byte against the Python engine's EXL3 kernels: five linears, a dense prefill, an expert layer, and whole KDA and DSA layers with all 288 experts under real routing.
  • exl3-tp2: the GLM split rules on two GB10s over NCCL, the same way.
• Decode / prompt speed: n/a. No family serves these yet.
• Tests run:
  • tf-cuda-test exl3: 28 pass, 0 fail, 0 skip.
  • zig build test-cuda-host: pass.
  • zig fmt --check zig: clean.
  • tools/zig/lean_check.py: 538, the same as c126fe16.
• Not run:
  • Metal: no Metal code changes.
  • The lane checks: no family uses this yet.
  • Without DIR, the Python-engine comparisons SKIP and the float64 and invariance checks still run.

Checklist

• Every token still goes through the lane rounds. (Operators only; no serving path is added.)
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (n/a until a family uses it)
• No precision traded for speed.
• Prompt processing is no slower than the last release. (No existing path changes.)
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency. (They skip cleanly without fixtures; the code is new, so there is no "before" failure.)
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thank you, @justincasey. This is careful work: hand-written EXL3 kernels for the linear and the grouped routed experts at every codebook and width, a width per expert, and a receipt that holds a row's bytes equal across row counts 1 to 128, K splits, member counts and prompt chunking, with real RED-SNOW layers byte for byte against the Python engine's kernels. That is exactly the bar we want.

Two things before it lands. Your headers say these port the Python engine's decode.cuh, linear.cu and experts.cu, which @vcruz305 wrote in #42, so the landing commit will carry his Co-authored-by line beside your authorship. We add it ourselves; there is nothing for you to rebase. And @BHCC2025's #740, #741 and #806 build the same projection behind qlinear from core/exl3_format.zig. We want one EXL3 module, so we will run both on the same tensors, bits against the float64 oracle and bandwidth against ExLlamaV3's LinearEXL3, keep the faster at equal bits, and credit both authors on it.

We are gathering the EXL3 work in #832, where Cruz is invited to lead. It would be great to have you there.

[justincasey]

Ok sounds good. Let's make sure that the mixed widths don't get lost. Happy to bench or compare anything. I have 4 GB10s I help with if needed.

[ashhart]

Thank you, @justincasey, the per-expert descriptors and broader format coverage make #837 the common consolidation base.

We checked 224406d. Before composition, we will repair the packed-scale scan/load mismatch, allocate scratch replacements before releasing their owner, and refund upload buffers when list insertion fails. We can make those bounded fixes with failure regressions at landing.

Then we will carry #806's FP32 contract, #741's Grouped integration and #729's Metal path, preserving your work, @BHCC2025, @mrv777 and @vcruz305's credit. The equal-bit comparison remains required to choose the faster correct implementation per operation; selecting a coverage base does not supply that timing result. Mixed expert widths and full-token cache/draft gates stay together.

Posted by Ash's AI Assistant Synthena