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

docs/recipes/exl3.md

@@ -172,4 +172,30 @@ reproduces the numpy oracle on every codebook and width and on real checkpoint b
 (the oracle vectors ship beside the test). Its real checkpoint bytes are one 128x128 block from
 `UnstableLlama/Qwen3.5-2B-exl3-8.00bpw` (mcg codebook, 8.00 bpw), in `zig/tests/fixtures/exl3_oracle.bin`.
 
-The scalar decoder and its oracle tests do not register an EXL3 model or provide a native GPU projection.
+The scalar decoder and its oracle tests do not register an EXL3 model on their own; the native CUDA kernels below
+are checked against them.
+
+## Native CUDA (Zig)
+
+The Zig engine carries the same arithmetic as hand-written CUDA, ported from the Python engine's kernels so a
+row's bits match it:
+
+- `zig/kernels/cuda/exl3_decode.cuh`, `exl3_linear.cu`, `zig/src/cuda/exl3.zig`: the input rotation and the linear
+  for 1 to 128 rows, every codebook and width (the `mul1` half steps included), with the K split planned from the
+  layer's shape alone. `exl3.Dense` runs prompt rows in 128-row blocks of the same kernel, so a prompt row's bits
+  equal its decode-row bits wherever its chunk starts; `qlinear.Weight` has an `.exl3` arm for both.
+- `zig/kernels/cuda/exl3_experts.cu`, `zig/src/cuda/exl3_experts.zig`: grouped routed experts (member lists,
+  each expert's rotation, gate and up, GLM's SwiGLU roundings, down and the weighted combine), widths mixed within
+  a layer; an expert's member count never changes a row's bits.
+- `zig/src/core/exl3_load.zig`: every trellis group of a checkpoint from its safetensors headers;
+  `zig/src/cuda/exl3_upload.zig` puts a group, or a layer's experts, on the device.
+
+`tf-cuda-test exl3 DIR` checks them: every width and codebook against float64 references built on
+`exl3_format.zig`, row and chunk invariance, and, with `DIR` from
+`zig/tests/cuda/oracle/exl3_oracle.py` (run with the Python engine on a CUDA box), real checkpoint layers bit for
+bit against the Python engine, two whole layers of RED-SNOW-5.3-Flash (a KDA and a DSA layer, all 288 experts under
+real routing) included. `tf-cuda-test exl3-tp2 RANK HOST:PORT DIR` runs the GLM split rules on two ranks over NCCL
+(output-column cuts, input-row cuts with fp32 partials summed rank 0 first, and the experts' intermediate cut) and
+checks every rank's bytes and the sums the same way.
+
+No family serves EXL3 natively yet: a family's engine takes these pieces through `qlinear` and `exl3_experts`.

zig/build/cuda.zig

@@ -41,6 +41,8 @@ const kernels = [_]Kernel{
     .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3
     .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts
+    .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5
+    .{ .name = "exl3_experts", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/experts.cu's device code, tensorfold_exl3_experts
     .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops },
     .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops },
     .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops },

zig/kernels/cuda/exl3_decode.cuh

@@ -0,0 +1,141 @@
+// EXL3 tiles: lane L decodes values 8L..8L+7 straight into its two mma B fragments; port of tensorfold's decode.cuh.
+#pragma once
+
+#include <cstdint>
+#include <cuda_fp16.h>
+
+namespace tf_exl3 {
+
+enum Codebook : int { CB_3INST = 0, CB_MCG = 1, CB_MUL1 = 2 };
+
+// 32-bit words of a tile: 256 values of K2 / 2 bits.
+template <int K2>
+__host__ __device__ constexpr int tile_words() {
+    return 4 * K2;
+}
+
+// E(p): one past the last stream bit of value p's 16-bit window (exl3_format.zig's streamEnds).
+template <int K2>
+__host__ __device__ constexpr int stream_end(int p) {
+    return (K2 & 1) ? ((p + 1) * K2 - ((p + 1) & 1)) / 2 : (p + 1) * (K2 / 2);
+}
+
+// Words a lane reads for its 8 values: the most any lane's windows span (2, or 3 for 3.5 bits and 5 to 8 bits).
+template <int K2>
+__host__ __device__ constexpr int lane_words() {
+    int most = 0;
+    for (int l = 0; l < 32; ++l) {
+        const int first = (stream_end<K2>(8 * l) - 16 + 128 * K2) % 32;
+        const int need = first + stream_end<K2>(8 * l + 7) - stream_end<K2>(8 * l) + 16;
+        most = need > most ? need : most;
+    }
+    return (most + 31) / 32;
+}
+
+// The index of the first word lane `lane` reads, and the bit offset of its first window in that word.
+template <int K2>
+__device__ __forceinline__ void lane_start(int lane, int& word, int& offset) {
+    const int first = 4 * lane * K2 + K2 / 2 - 16 + 128 * K2;   // E(8 lane) - 16, made non-negative
+    word = (first >> 5) % tile_words<K2>();
+    offset = first & 31;
+}
+
+// The same from global memory through the read-only cache.
+template <int K2>
+__device__ __forceinline__ void ldg_lane_words(const uint32_t* tile, int lane, uint32_t (&w)[lane_words<K2>()]) {
+    int word, offset;
+    lane_start<K2>(lane, word, offset);
+#pragma unroll
+    for (int i = 0; i < lane_words<K2>(); ++i) w[i] = __ldg(tile + (word + i) % tile_words<K2>());
+}
+
+// Bits [d, d + 16) of the 96-bit big-endian stream hi:mid:lo (d a compile-time constant after unrolling).
+__device__ __forceinline__ uint32_t window16(uint32_t hi, uint32_t mid, uint32_t lo, int d) {
+    if (d <= 16) return (hi >> (16 - d)) & 0xffffu;
+    if (d < 32) return __funnelshift_l(mid, hi, d) >> 16;
+    if (d <= 48) return (mid >> (48 - d)) & 0xffffu;
+    return __funnelshift_l(lo, mid, d - 32) >> 16;
+}
+
+// The 16-bit states of this lane's 8 values (8 lane + j, j = 0..7) from its words.
+template <int K2>
+__device__ __forceinline__ void lane_states(const uint32_t (&w)[lane_words<K2>()], int lane, uint32_t (&s)[8]) {
+    int word, offset;
+    lane_start<K2>(lane, word, offset);
+    const uint32_t c = lane_words<K2>() > 2 ? w[lane_words<K2>() > 2 ? 2 : 0] : 0u;
+    const uint32_t hi = __funnelshift_l(w[1], w[0], offset);   // the stream from the lane's first window on
+    const uint32_t mid = __funnelshift_l(c, w[1], offset);
+    const uint32_t lo = c << offset;
+#pragma unroll
+    for (int j = 0; j < 8; ++j) s[j] = window16(hi, mid, lo, stream_end<K2>(j) - stream_end<K2>(0));
+}
+
+// Two codebook values from two 16-bit states, as a half2 in a uint32 (the first state in the low half).
+template <int CB>
+__device__ __forceinline__ uint32_t decode2(uint32_t s0, uint32_t s1) {
+    if constexpr (CB == CB_MUL1) {
+        const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du;
+        const uint32_t h0 = __dp4a(x0, 0x01010101u, 0x6400u);    // fp16 bits of 1024 + the byte sum
+        const uint32_t h1 = __dp4a(x1, 0x01010101u, 0x6400u);
+        const uint32_t hh = (h0 & 0xffffu) | (h1 << 16);
+        const half2 r = __hfma2(*reinterpret_cast<const half2*>(&hh), __half2half2(__ushort_as_half(0x1eee)),
+                                __half2half2(__ushort_as_half(0xc931)));
+        return *reinterpret_cast<const uint32_t*>(&r);
+    } else {
+        uint32_t x0, x1;
+        if constexpr (CB == CB_MCG) {
+            x0 = s0 * 0xCBAC1FEDu;
+            x1 = s1 * 0xCBAC1FEDu;
+        } else {
+            x0 = s0 * 89226354u + 64248484u;
+            x1 = s1 * 89226354u + 64248484u;
+        }
+        x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u;
+        x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u;
+        const uint32_t lo = __byte_perm(x0, x1, 0x5410);
+        const uint32_t hi = __byte_perm(x0, x1, 0x7632);
+        const half2 r = __hadd2(*reinterpret_cast<const half2*>(&lo), *reinterpret_cast<const half2*>(&hi));
+        return *reinterpret_cast<const uint32_t*>(&r);
+    }
+}
+
+// This lane's 8 values of a tile as the B fragments of the tile's two n8 halves (columns 0-7 in b0, 8-15 in b1).
+template <int K2, int CB>
+__device__ __forceinline__ void decode_lane(const uint32_t (&w)[lane_words<K2>()], int lane, uint32_t (&b0)[2],
+                                            uint32_t (&b1)[2]) {
+    uint32_t s[8];
+    lane_states<K2>(w, lane, s);
+    b0[0] = decode2<CB>(s[0], s[1]);
+    b0[1] = decode2<CB>(s[2], s[3]);
+    b1[0] = decode2<CB>(s[4], s[5]);
+    b1[1] = decode2<CB>(s[6], s[7]);
+}
+
+// mma.m16n8k16, fp16 inputs, fp32 accumulators: d += a @ b.
+__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) {
+    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, "
+                 "{%0,%1,%2,%3};\n"
+                 : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+                 : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
+}
+
+// Unscaled Walsh-Hadamard transform of 128 fp32 values, 4 a lane, in a fixed butterfly order (multiply by 1/sqrt(128)).
+__device__ __forceinline__ void fwht128(float (&v)[4], int lane) {
+    const float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3];
+    v[0] = a + c;
+    v[1] = b + d;
+    v[2] = a - c;
+    v[3] = b - d;
+#pragma unroll
+    for (int m = 1; m < 32; m <<= 1) {
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const float o = __shfl_xor_sync(0xffffffffu, v[j], m);
+            v[j] = (lane & m) ? o - v[j] : v[j] + o;
+        }
+    }
+}
+
+constexpr float HAD_SCALE = 0.08838834764831845f;   // 1 / sqrt(128)
+
+}  // namespace tf_exl3