glm: GLM-5.3-Flash on CUDA, prompts 267.6 → 447.9 tok/s and decode 19.9 → 25.1 tok/s, every reply unchanged (#593)

docs/recipes/exl3.md

@@ -172,4 +172,32 @@ 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. A block of more than 16 rows on a `mul1`
+  weight takes the prompt kernel: 64-row passes that decode each weight tile once, keeping the decode-row kernel's
+  sums in its order. A prompt row's bits therefore 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`.

docs/recipes/glm-5.3-flash.md

@@ -3,7 +3,8 @@
 This recipe describes the Python 0.6.6 engine, maintained on
 [`python-0.6`](https://github.com/ashhart/TensorFold/tree/python-0.6).
 Its CUDA EXL3 reader, two-rank flags and DFlash2 policies do not apply to the native Zig binary.
-The native binary serves GLM on two M5 Ultra Macs with MLX affine 4-bit/group-64 weights; it has no GLM CUDA backend yet.
+The native binary serves GLM on two M5 Ultra Macs with MLX affine 4-bit/group-64 weights; it does not serve GLM on
+CUDA yet ([its CUDA engine](#native-cuda-engine-zig) runs under `tf-cuda-test` only).
 That is not an engine-only replacement for a CUDA EXL3 deployment.
 For native serving, use the [qualified model table](../../README.md#qualified-models) and
 [paired Metal runbook](../../RUNBOOK.md#flash-next-and-paired-metal-serving).
@@ -61,6 +62,38 @@ The expert decoder and BF16 target matmul keep row arithmetic fixed. A quantized
 propose drafts, but target verification retains the BF16 head. EXL3 speed, capacity and long-context
 qualification are TBD [release-0.3.5].
 
+### Native CUDA engine (Zig)
+
+The Zig side carries a GLM-5.3-Flash engine for two CUDA ranks that reads
+`vcruz305/RED-SNOW-5.3-FLASH-EXL3-SAGE-4.91bpw` (EXL3 `mul1` at mixed widths, unsplit). Each rank takes its
+shard at load. Its kernels are hand-written CUDA ported from the Python engine on this repository's `redsnow-v1`
+branch. They cover the bf16 matmul and glue, KDA's decode chain and chunked prompt kernels, the DSA indexer and
+top-pool selection, latent attention, hyper-connection mixing, the MoE router and EXL3 experts, the dense MLP of
+layers 0-2, and the head with its greedy pick. Each one keeps the Python engine's order of operations, so a
+greedy reply is token for token the Python engine's `"draft": false` reply. Prompts run in 2,048-row chunks.
+Decoding is serial (no drafts): each step replays a captured CUDA graph, one per KDA parity in the dense window and
+one per pool bucket past it.
+
+On two GB10s, over the eight reference prompts (221 to 16,416 prompt tokens; 24,099 in all), it read prompts at
+435.0 tok/s and decoded at 24.7 tok/s. The Python engine decodes serially at 24.6 tok/s.
+
+Two commands check it, both built by `zig build -Dnvcc=/usr/local/cuda/bin/nvcc`:
+
+```bash
+# fixtures from the Python engine (run on a CUDA box with redsnow-v1 installed), then the kernels and layers
+python zig/tests/cuda/oracle/glm_oracle.py --model MODEL_DIR --out FIXTURES
+tf-cuda-test glm FIXTURES
+# the whole model at TP=2, a process a rank (rank 1 first), against recorded replies
+tf-cuda-test glm-model 1 HOST:PORT MODEL_DIR reference.json
+tf-cuda-test glm-model 0 HOST:PORT MODEL_DIR reference.json
+```
+
+`tf-cuda-test glm` holds every kernel against float64 references and the Python engine's bytes, then whole KDA
+and DSA layers of the real checkpoint at TP=2. `reference.json` is
+`{"cases": [{"name", "prompt_ids", "ids"}], "eos_ids"}`: the Python engine's greedy replies with `"draft": false`.
+A case passes when the Zig reply's ids equal `ids`. Rank 0 prints a PASS or FAIL line for each case and
+`RESULT glm.model.prefill_tps` / `decode_tps`.
+
 ### Draft policies
 
 For the affine checkpoint with its MTP head loaded, the default `auto` policy uses MTP for sampled requests.

zig/build/cuda.zig

@@ -41,6 +41,16 @@ 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 = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda.cu's device code (chain, replay, wide trio, conv shift)
+    .{ .name = "glm_kda_chunk", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda_chunk.cu's device code (32-row WY sub-chunks)
+    .{ .name = "glm_bmm", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda's bf16 matmul and router dots (qmm._bmm, glue._router)
+    .{ .name = "glm_glue", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda/glue.py's Triton kernels as device code
+    .{ .name = "glm_glue2", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glue.py's 1536/512 rmsnorms and the fp32-y combine (glm_glue.cu is at the file cap)
+    .{ .name = "glm_sparse", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // glm5_next/cuda/sparse.py's Triton kernels as device code
+    .{ .name = "glm_latent", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // latent.py's cache write and kv_b projections
+    .{ .name = "glm_latent_att", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // latent.py's chunk attention and merge (glm_latent.cu is at the file cap)
     .{ .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 },
@@ -77,7 +87,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti
     return cuda;
 }
 
-fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } {
+fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, glm: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } {
     const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true });
     const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true });
     core.addImport("tokenizer", tokenizer);
@@ -89,7 +99,10 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin
     nemotron.addImport("lanes", lanes);
     nemotron.addImport("heat", heat);
     nemotron.addImport("nemotron_draft_ids", draft_ids);
-    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat };
+    const glm = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    glm.addImport("cuda", cuda);
+    glm.addImport("core", core);
+    return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .glm = glm, .tokenizer = tokenizer, .heat = heat };
 }
 
 /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`.
@@ -137,6 +150,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     runner.addImport("exlr8_frozen", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen.zig"), .target = target, .optimize = optimize }));
     runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize }));
     for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module);
+    runner.addImport("glm", mods.glm);
     b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }));
     nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip;
 }
@@ -204,7 +218,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     step.dependOn(&server_test.step);
     b.step("test-server-cpu", "The HTTP server's unit tests (routes, templates, tool parsing) without a GPU").dependOn(&server_test.step);
     // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them
-    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step);
+    for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.glm, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step);
     const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true });
     cli.addImport("cuda", cuda);
     cli.addImport("core", mods.core);

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