cuda gemma4: every MLX affine width the family admits, and the 31B on CUDA (#1048 step 9)

docs/kernel-contract.md

@@ -40,7 +40,9 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho
 3. Shapes. Every served shape, then each tiled dimension at its least size, a tile minus a step, a tile, a tile plus a
    step and the largest served size. Each output must be within `|want| * outputs' half ulp + gamma(depth) * magnitude`
    of float64, or for a replacement kernel bit-equal to the kernel it replaces. Control: a row sum whose partial tail
-   loses its last input; it passes at whole tiles and fails at the edges.
+   loses its last input; it passes at whole tiles and fails at the edges. A kernel with bf16 rounding points between
+   its steps (MLX's op-by-op rows) checks each output against float64's value at those points, as a range over the
+   spread its fp32 steps allow, and its controls each drop one rounding point, which that check must catch.
 4. fp32 sums. The inputs are integers whose partial sums reach about 2^20 in any order and stay below 2^24, so fp32
    sums are exact whatever the kernel's order; the second half of each row mirrors the first, so the result is small.
    Before the kernel runs, the harness replays the sums on the CPU in fp16 and bf16 (one chain, contiguous halves, and
@@ -104,15 +106,24 @@ Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), a
 formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
 prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
 tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
-CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router logits (and their 128-row tiles,
+CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and
+the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles,
 bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and
 shared add over bf16 share sums (bit-equal to `tf_super_combine` and `tf_super_moe_out` over the fp32 shares), the
 Mamba scan of a verify window (its host twin's bytes), qmm_group's cluster tile, lane_gemv and its split-K form (both
 bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a
 cluster, fused, or by `reduce_kernel<true>`), the checkpoint math's FP8 128-row prompt tiles (both slice meetings,
 bit-equal to the 64-row tiles) and NVFP4 experts: the decode-sized forms and the prompt tiles (up to down's NVFP4
-rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), and the
-reduce-scatter's `tf_rank_sum`.
+rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), the
+reduce-scatter's `tf_rank_sum`, and `window_attention`'s rows at head dims 256 and 512 (`window_rows.cu`: decode chunks
+and their merge, prompt tiles with P as bf16 or in two halves, a row's last key, windows over a ring or none over a
+linear cache), each output within float64's range over the spread its fp32 steps and tensor-core sums allow
+(`core/attention_range.zig`; seven controls, each one planted fault), and the activation rows in MLX's bf16 steps (GeGLU, and
+SwiGLU as DFlash drafters round it; twelve controls, each dropping one rounding point; their 64-value sums bit-equal to
+the fp32 tree), and the MLX affine token rows (4 bits, and 2 to 6 or 8 bits times a bf16 scale, each
+output bit-equal to float64 at its rounding points; two controls, one dropping the rounding point, one folding the scale
+into the group's scale and bias), and qk_rows' attention inputs at head dims to 512 with a unit v norm and MLX's bf16
+rows (three controls, each dropping one option).
 
 ## CUDA records per nvcc release
 

zig/build/cuda.zig

@@ -115,6 +115,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat);
     runner.addImport("glm_cuda", glm_cuda);
     runner.addImport("kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes));
+    runner.addImport("gemma4_cuda", @import("gemma4_cuda.zig").module(b, target, optimize, cuda, mods.core, mods.lanes));
     const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner });
     b.installArtifact(runner_exe);
     // tf-tp-watch-check: the ranks' lines with no GPU or NCCL; a failing or dying rank stops every rank
@@ -204,6 +205,7 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti
             .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) },
             .{ .name = "qwen35", .module = qwen35 },
             .{ .name = "kolibri1", .module = kolibri1Module(b, target, optimize, cuda, core, lanes) },
+            .{ .name = "gemma4_cuda", .module = @import("gemma4_cuda.zig").module(b, target, optimize, cuda, core, lanes) },
         },
     });
     return .{ .api = api, .engines = mod };
@@ -292,6 +294,14 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
         .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } },
     }) }));
     step.dependOn(&writeback_contract.step);
+    const tail_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{
+        .root_source_file = b.path("zig/tests/cuda/contract_norms_tail.zig"),
+        .target = host,
+        .optimize = .debug,
+        .link_libc = true,
+        .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } },
+    }) }));
+    step.dependOn(&tail_contract.step);
     const utility_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{
         .root_source_file = b.path("zig/tests/cuda/contract_utility_ref.zig"),
         .target = host,
@@ -321,6 +331,16 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     }) }));
     step.dependOn(&qsa_contract.step);
     flash_contract.dependOn(&qsa_contract.step);
+    // the Gemma 4 operators' contract tests (#1048): window rows, MLX activations and soft cap, embed rows, qk rows
+    for ([_][]const u8{ "contract_window_rows.zig", "contract_act_round.zig", "contract_softcap.zig", "contract_embed.zig", "contract_qk_round.zig" }) |f| {
+        step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{
+            .root_source_file = b.path(b.fmt("zig/tests/cuda/{s}", .{f})),
+            .target = host,
+            .optimize = .debug,
+            .link_libc = true,
+            .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } },
+        }) })).step);
+    }
     const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)");
     const projection_step = b.step("test-cuda-projection-fixture", "Projection fixture dispatch controls with CPU-only CUDA calls");
     for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| {
@@ -471,6 +491,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat);
     @import("glm_family_host.zig").steps(b, cuda, glm, mods, step);
     @import("kolibri.zig").host(b, kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), native_modules.api, step);
+    @import("gemma4_cuda.zig").host(b, @import("gemma4_cuda.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step);
     @import("glm_contracts.zig").tests(b, glm, mods.lanes, cuda, mods.core);
     // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's)
     const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug });

zig/build/cuda_kernels.zig

@@ -70,7 +70,7 @@ pub const kernels = [_]Kernel{
     .{ .name = "ck_lane", .flags = &.{"-O3"} }, // nvfp4/lane4.cu's A8: FP8 rows times FP8 weights (checkpoint math)
     .{ .name = "ck_experts", .flags = &.{"-O3"}, .arch_specific = true }, // ours: NVFP4 x NVFP4 experts (sm_121a)
     .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig)
-    .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig)
+    .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: MLX affine and bf16 token rows (cuda/embed_rows.zig)
     .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows
     .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle
     .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig)
@@ -82,6 +82,7 @@ pub const kernels = [_]Kernel{
     .{ .name = "tree_attention", .flags = &.{"-O3"} }, // ours: committed keys and paths (cuda/tree_attention.zig)
     .{ .name = "block_attention", .flags = &.{"-O3"} }, // ours: a drafter's block over its context (cuda/block_attention.zig)
     .{ .name = "window_attention", .flags = &.{"-O3"} }, // ours: sliding-window attention (cuda/window_attention.zig)
+    .{ .name = "window_rows", .flags = &.{"-O3"} }, // ours: window attention at head dims 256 and 512 (cuda/window_attention.zig)
     .{ .name = "moe_glue", .flags = &.{"-O3"} }, // ours: sigmoid top-k and the weighted combine (cuda/moe_glue.zig)
     .{ .name = "dense_rows", .flags = &.{"-O3"} }, // ours: bf16 rows times weight rows into fp32 (cuda/dense_rows.zig)
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts

zig/build/gemma4_cuda.zig

@@ -0,0 +1,12 @@
+//! Gemma 4's CUDA family (#1048 step 8) over the maintained runtime, the core and the lane core; its host tests.
+const std = @import("std");
+
+pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module {
+    return b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } });
+}
+
+pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void {
+    const run = b.addRunArtifact(b.addTest(.{ .root_module = m }));
+    b.step("test-gemma4-cuda-host", "Gemma 4 CUDA config copies, RoPE tables, snapshot runs and refusals without a GPU").dependOn(&run.step);
+    all.dependOn(&run.step);
+}

zig/kernels/cuda/act_round.cuh

@@ -0,0 +1,116 @@
+// MLX's bf16-op-by-op activations (act_rows.cu's tf_geglu_rows, tf_swiglu_rows_round); `keep`: the rounding points.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+namespace tf_act_round {
+
+// mlx_lm's gelu_approx(x) * up as MLX's bf16 ops (Metal gemma4/glue.metal's gemma_geglu): its rounding points.
+enum : int { CUBE = 1, CUBE_SCALED = 2, INNER = 4, INNER_SCALED = 8, TANH = 16, ONE_PLUS = 32, GELU = 64 };
+constexpr int GEGLU_ALL = 127;
+// tf_draft_swiglu's bf16 steps (Metal core/draft_ops.metal, gemma4's DFlash drafter): its rounding points.
+enum : int { EXP = 1, DENOM = 2, RECIP = 4, SIGMOID = 8, SILU = 16 };
+constexpr int SWIGLU_ALL = 31;
+// MLX's logit soft cap cap * tanh(x / cap) as bf16 ops (Metal gemma4/glue.metal's gemma_softcap): its rounding points.
+enum : int { CAP_DIV = 1, CAP_TANH = 2 };
+constexpr int SOFTCAP_ALL = 3;
+
+__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); }
+
+// v rounded to bf16 where `keep` holds the point p (a constant in the served kernels, so the test folds away).
+__device__ __forceinline__ float point(float v, int keep, int p) { return (keep & p) != 0 ? bf(v) : v; }
+
+// 0.5 x (1 + tanh(c0 (x + c1 x^3))) up, each step rounded to bf16 (x^3 exact in fp32 for a bf16 x; c0, c1 as bf16).
+__device__ __forceinline__ float geglu(float x, float up, int keep) {
+    const float c0 = bf(0.7978845608028654f), c1 = bf(0.044715f);
+    const float cube = point(__fmul_rn(__fmul_rn(x, x), x), keep, CUBE);
+    const float inner = point(__fadd_rn(x, point(__fmul_rn(c1, cube), keep, CUBE_SCALED)), keep, INNER);
+    const float th = point(tanhf(point(__fmul_rn(c0, inner), keep, INNER_SCALED)), keep, TANH);
+    const float half_x = bf(__fmul_rn(0.5f, x)); // exact for a bf16 x: no point to drop
+    const float gl = point(__fmul_rn(half_x, point(__fadd_rn(1.f, th), keep, ONE_PLUS)), keep, GELU);
+    return __fmul_rn(gl, up);
+}
+
+// silu(g) up as sigmoid from exp(|g|): 1 / (1 + e) for g < 0, else 1 - that; each step rounded to bf16.
+__device__ __forceinline__ float swiglu(float g, float up, int keep) {
+    const float e = point(expf(fabsf(g)), keep, EXP);
+    const float inv = point(__fdiv_rn(1.f, point(__fadd_rn(1.f, e), keep, DENOM)), keep, RECIP);
+    const float sig = g < 0.f ? inv : point(__fsub_rn(1.f, inv), keep, SIGMOID);
+    return __fmul_rn(point(__fmul_rn(g, sig), keep, SILU), up);
+}
+
+// bf16(bf16(tanh(bf16(x / cap))) * cap): the division and product correctly rounded, tanhf within its 2 ulps.
+__device__ __forceinline__ float softcap(float x, float cap, int keep) {
+    const float t = point(__fdiv_rn(x, cap), keep, CAP_DIV);
+    return __fmul_rn(point(tanhf(t), keep, CAP_TANH), cap);
+}
+
+__device__ __forceinline__ void unpack8(const uint4& u, float* f) {
+    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&u);
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+        const float2 t = __bfloat1622float2(h[i]);
+        f[2 * i] = t.x;
+        f[2 * i + 1] = t.y;
+    }
+}
+
+__device__ __forceinline__ uint4 pack8(const float* f) {
+    uint4 u;
+    __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&u);
+#pragma unroll
+    for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(f[2 * i], f[2 * i + 1]);
+    return u;
+}
+
+// act_rows.cu's group_sum8: a thread's 8 stored values in order, then the group's 8 lanes in a fixed tree.
+__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) {
+    float f[8];
+    unpack8(packed, f);
+    float s = 0.f;
+    if (live)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]);
+    s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4));
+    s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2));
+    return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1));
+}
+
+// out[r, i] = bf16(act(gate[r * g_ld + i], up[r * u_ld + i])) for i < n (n % 64 == 0), GEGLU or not; xs its sums.
+template <bool GEGLU>
+__device__ __forceinline__ void rows(const __nv_bfloat16* __restrict__ gate, long long g_ld,
+                                     const __nv_bfloat16* __restrict__ up, long long u_ld, __nv_bfloat16* __restrict__ out,
+                                     float* __restrict__ xs, int n, int keep) {
+    const int64_t r = blockIdx.x;
+    const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values
+    const bool live = 8 * c < n;
+    float g[8], u[8], o[8];
+    if (live) {
+        unpack8(*reinterpret_cast<const uint4*>(gate + r * g_ld + 8 * c), g);
+        unpack8(*reinterpret_cast<const uint4*>(up + r * u_ld + 8 * c), u);
+    }
+#pragma unroll
+    for (int i = 0; i < 8; ++i) o[i] = live ? (GEGLU ? geglu(g[i], u[i], keep) : swiglu(g[i], u[i], keep)) : 0.f;
+    const uint4 op = pack8(o);
+    if (live) *reinterpret_cast<uint4*>(out + r * n + 8 * c) = op;
+    if (xs != nullptr) {
+        const float s = group_sum8(op, live);
+        if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s;
+    }
+}
+
+// out[r, i] = softcap(x[r, i]) for i < n (n % 8 == 0), rows packed n wide (out may be x); 8 values a thread.
+__device__ __forceinline__ void softcap_rows(const __nv_bfloat16* x, __nv_bfloat16* out, int n, float cap, int keep) {
+    const int64_t r = blockIdx.x;
+    const int c = blockIdx.y * blockDim.x + threadIdx.x;
+    if (8 * c >= n) return;
+    float v[8];
+    unpack8(*reinterpret_cast<const uint4*>(x + r * n + 8 * c), v);
+#pragma unroll
+    for (int i = 0; i < 8; ++i) v[i] = softcap(v[i], cap, keep);
+    *reinterpret_cast<uint4*>(out + r * n + 8 * c) = pack8(v);
+}
+
+}  // namespace tf_act_round

zig/kernels/cuda/act_rows.cu

@@ -141,3 +141,26 @@ extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat
     }
     *reinterpret_cast<uint4*>(out + r * d + 8 * c) = pack8(y);
 }
+
+// MLX's bf16-op-by-op activations (act_round.cuh), the same rows and sums as tf_swiglu_rows.
+#include "act_round.cuh"
+
+// out = bf16(gelu_tanh(gate) * up) as mlx_lm's geglu in MLX's bf16 ops (Gemma 4's MLP); xs (or null) out's sums.
+extern "C" __global__ void __launch_bounds__(NT) tf_geglu_rows(const __nv_bfloat16* gate, long long g_ld,
+                                                               const __nv_bfloat16* up, long long u_ld,
+                                                               __nv_bfloat16* out, float* xs, int n) {
+    tf_act_round::rows<true>(gate, g_ld, up, u_ld, out, xs, n, tf_act_round::GEGLU_ALL);
+}
+
+// out = bf16(silu(gate) * up) in MLX's bf16 steps (tf_draft_swiglu's, Gemma 4's DFlash drafter); xs (or null) its sums.
+extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows_round(const __nv_bfloat16* gate, long long g_ld,
+                                                                      const __nv_bfloat16* up, long long u_ld,
+                                                                      __nv_bfloat16* out, float* xs, int n) {
+    tf_act_round::rows<false>(gate, g_ld, up, u_ld, out, xs, n, tf_act_round::SWIGLU_ALL);
+}
+
+// out = bf16(cap * tanh(x / cap)) in MLX's bf16 steps: a soft-capped head's logits (final_logit_softcapping).
+extern "C" __global__ void __launch_bounds__(NT) tf_softcap_rows(const __nv_bfloat16* x, __nv_bfloat16* out, int n,
+                                                                 float cap) {
+    tf_act_round::softcap_rows(x, out, n, cap, tf_act_round::SOFTCAP_ALL);
+}