dory: the family on CUDA, skeleton, f64 reference, MXFP8 and BF16 maps, and the forward for one stream (#1233)

docs/kernel-contract.md

@@ -149,6 +149,14 @@ CUDA also checks `window_attention`'s rows at head dims 256 and 512 (`window_row
 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).
+CUDA also checks Dory's nGPT rows: the L2 norm rows with and without their bf16 add, and the normalized lerp residual
+(`norms_l2.cuh`), each output within float64's range at its bf16 points (three controls, each dropping one point).
+CUDA also checks Dory's attention inputs (`qk_l2.cuh`): RoPE first, each head's L2 norm, q times its scale vector, v as is, at
+head dim 256 and positions past 199,800 (two controls, each dropping one rounding point).
+CUDA also checks Dory's activation rows (`act_l2.cuh`): the scaled sigmoid output gate over all heads and the scaled SwiGLU,
+each row then L2-normed, against float64's ranges at their bf16 points and expf's ulps (four controls, one point each).
+CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) times a per-column scale), bit-equal to its
+unscaled tile's outputs rounded and scaled on the host.
 HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out).
 
 RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then

zig/build/cuda.zig

@@ -102,7 +102,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module);
     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));
+    for ([_]struct { []const u8, *std.Build.Module }{ .{ "kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes) }, .{ "dory", @import("dory.zig").module(b, target, optimize, cuda, mods.core) } }) |m| runner.addImport(m[0], m[1]);
     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);

zig/build/cuda_family_host.zig

@@ -3,6 +3,7 @@ const std = @import("std");
 
 pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.Build.Step) void {
     const host = b.graph.host;
+    @import("dory.zig").host(b, @import("dory.zig").module(b, host, .debug, cuda, mods.core), step);
     const proposal = b.createModule(.{ .root_source_file = b.path("zig/proposal_router_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "core", .module = mods.core }, .{ .name = "lanes", .module = mods.lanes } } });
     b.step("test-proposal-router-host", "Optional common proposal routing, legacy parity and lifecycle").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = proposal, .filters = &.{"proposal router:"} })).step);
 

zig/build/cuda_row_contracts.zig

@@ -8,7 +8,7 @@ pub fn host(b: *std.Build, target: std.Build.ResolvedTarget, cuda: *std.Build.Mo
     step.dependOn(&radio_run.step);
     const stress = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_group_stress_host.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } });
     b.step("test-group-sums-stress-host", "Actual group-sum stress inputs distinguish reduced sums at every seed").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = stress })).step);
-    for ([_][]const u8{ "contract_norms_tail.zig", "contract_act_round.zig", "contract_embed.zig", "contract_qk_round.zig", "contract_window_rows.zig" }) |name| {
+    for ([_][]const u8{ "contract_norms_tail.zig", "contract_l2.zig", "contract_qk_l2.zig", "contract_act_l2.zig", "contract_act_round.zig", "contract_embed.zig", "contract_qk_round.zig", "contract_window_rows.zig" }) |name| {
         const m = b.createModule(.{
             .root_source_file = b.path(b.pathJoin(&.{ "zig/tests/cuda", name })),
             .target = target,

zig/build/dory.zig

@@ -0,0 +1,12 @@
+//! Dory family construction and its host tests (config, weight map, schedule, memory) without a GPU.
+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) *std.Build.Module {
+    return b.createModule(.{ .root_source_file = b.path("zig/src/families/dory/dory.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } });
+}
+
+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-dory-host", "Dory config, weight map, sublayer schedule and memory without a GPU").dependOn(&run.step);
+    all.dependOn(&run.step);
+}

zig/kernels/cuda/act_l2.cuh

@@ -0,0 +1,135 @@
+// act_rows.cu's nGPT activations (Dory): a scaled sigmoid gate and a scaled SwiGLU, each row then L2-normed.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+namespace tf_act_l2 {
+
+constexpr int NT = 256;
+constexpr int MOST = 4; // 16-byte loads a thread: rows of width % 64, at most 8,192, 16 B aligned
+
+// The rounding points: the gated row; SwiGLU's scaled gate, scaled up and their product.
+enum : int { GATED = 1, GS = 2, US = 4, PROD = 8 };
+
+__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); }
+
+template <int R, int P>
+__device__ __forceinline__ float point(float v) {
+    if constexpr ((R & P) != 0) return bf(v);
+    else return v;
+}
+
+__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;
+}
+
+// torch's fp32 sigmoid and silu: 1 / (1 + exp(-x)), x / (1 + exp(-x)), IEEE division.
+__device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); }
+__device__ __forceinline__ float silu(float x) { return __fdiv_rn(x, __fadd_rn(1.f, expf(-x))); }
+
+// 1 / sqrt(sum + eps) of the row's squares: a thread's fma chain, the warp's xor tree, the 8 warps' tree (IEEE).
+__device__ __forceinline__ float inv_l2(const float (&v)[MOST][8], float eps) {
+    __shared__ float warps[NT / 32];
+    float part = 0.f;
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) part = __fmaf_rn(v[j][i], v[j][i], part);
+    for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m));
+    if ((threadIdx.x & 31) == 0) warps[threadIdx.x >> 5] = part;
+    __syncthreads();
+    float s = (threadIdx.x & 31) < NT / 32 ? warps[threadIdx.x & 31] : 0.f;
+    for (int m = 4; m > 0; m >>= 1) s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, m));
+    s = __shfl_sync(0xFFFFFFFFu, s, 0);
+    return __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(s, eps)));
+}
+
+// One input row's chunks of 8 (zeros past the row), and a scale vector's (fp32, 2 x 16 B a chunk).
+__device__ __forceinline__ void rowOf(const __nv_bfloat16* __restrict__ x, int64_t at, int chunks, float (&v)[MOST][8]) {
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) {
+        const int c = threadIdx.x + j * NT;
+        unpack8(c < chunks ? reinterpret_cast<const uint4*>(x + at)[c] : make_uint4(0, 0, 0, 0), v[j]);
+    }
+}
+
+__device__ __forceinline__ void scaleOf(const float* __restrict__ s, int chunks, float (&v)[MOST][8]) {
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) {
+        const int c = threadIdx.x + j * NT;
+        const float4 a = c < chunks ? reinterpret_cast<const float4*>(s)[2 * c] : make_float4(0.f, 0.f, 0.f, 0.f);
+        const float4 b = c < chunks ? reinterpret_cast<const float4*>(s)[2 * c + 1] : make_float4(0.f, 0.f, 0.f, 0.f);
+        v[j][0] = a.x, v[j][1] = a.y, v[j][2] = a.z, v[j][3] = a.w;
+        v[j][4] = b.x, v[j][5] = b.y, v[j][6] = b.z, v[j][7] = b.w;
+    }
+}
+
+__device__ __forceinline__ void normStore(__nv_bfloat16* __restrict__ out, int64_t at, int chunks, float (&v)[MOST][8], float eps) {
+    const float inv = inv_l2(v, eps);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) {
+        const int c = threadIdx.x + j * NT;
+        float y[8];
+#pragma unroll
+        for (int i = 0; i < 8; ++i) y[i] = __fmul_rn(v[j][i], inv);
+        if (c < chunks) reinterpret_cast<uint4*>(out + at)[c] = pack8(y);
+    }
+}
+
+// out = bf16(l2(t)), t = point<GATED>(o * sigmoid(g * sg)) over the row's n values (all heads): a block a row.
+template <int R>
+__device__ __forceinline__ void gate_rows(const __nv_bfloat16* __restrict__ o, const __nv_bfloat16* __restrict__ g, long long g_ld,
+                                          const float* __restrict__ sg, __nv_bfloat16* __restrict__ out, int n, float eps) {
+    const int64_t r = blockIdx.x;
+    const int chunks = n / 8;
+    float ov[MOST][8], gv[MOST][8], sv[MOST][8];
+    rowOf(o, r * n, chunks, ov);
+    rowOf(g, r * g_ld, chunks, gv);
+    scaleOf(sg, chunks, sv);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) ov[j][i] = point<R, GATED>(__fmul_rn(ov[j][i], sigmoid(__fmul_rn(gv[j][i], sv[j][i]))));
+    normStore(out, r * n, chunks, ov, eps);
+}
+
+// out = bf16(l2(m)), m = point<PROD>(silu(gs) us), gs = point<GS>(g sug), us = point<US>(u suu): dory.py's MLP.
+template <int R>
+__device__ __forceinline__ void swiglu_rows(const __nv_bfloat16* __restrict__ g, long long g_ld, const __nv_bfloat16* __restrict__ u,
+                                            long long u_ld, const float* __restrict__ sug, const float* __restrict__ suu,
+                                            __nv_bfloat16* __restrict__ out, int n, float eps) {
+    const int64_t r = blockIdx.x;
+    const int chunks = n / 8;
+    float gv[MOST][8], uv[MOST][8], sv[MOST][8];
+    rowOf(g, r * g_ld, chunks, gv);
+    rowOf(u, r * u_ld, chunks, uv);
+    scaleOf(sug, chunks, sv);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) gv[j][i] = point<R, GS>(__fmul_rn(gv[j][i], sv[j][i]));
+    scaleOf(suu, chunks, sv);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) gv[j][i] = point<R, PROD>(__fmul_rn(silu(gv[j][i]), point<R, US>(__fmul_rn(uv[j][i], sv[j][i]))));
+    normStore(out, r * n, chunks, gv, eps);
+}
+
+}  // namespace tf_act_l2

zig/kernels/cuda/act_rows.cu

@@ -164,3 +164,21 @@ extern "C" __global__ void __launch_bounds__(NT) tf_softcap_rows(const __nv_bflo
                                                                  float cap) {
     tf_act_round::softcap_rows(x, out, n, cap, tf_act_round::SOFTCAP_ALL);
 }
+
+// Dory's nGPT activations (act_l2.cuh): a scaled sigmoid gate and a scaled SwiGLU, each row then L2-normed.
+#include "act_l2.cuh"
+
+// out = bf16(l2(bf16(o * sigmoid(g * sg)))) over a row's n values (all heads), sg fp32 [n]: dory.py's output gate.
+extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) tf_gate_l2_rows(const __nv_bfloat16* o, const __nv_bfloat16* g,
+                                                                           long long g_ld, const float* sg, __nv_bfloat16* out,
+                                                                           int n, float eps) {
+    tf_act_l2::gate_rows<tf_act_l2::GATED>(o, g, g_ld, sg, out, n, eps);
+}
+
+// out = bf16(l2(bf16(silu(bf16(g * sug)) * bf16(u * suu)))), sug and suu fp32 [n]: dory.py's MLP before down.
+extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) tf_swiglu_l2_rows(const __nv_bfloat16* g, long long g_ld,
+                                                                             const __nv_bfloat16* u, long long u_ld,
+                                                                             const float* sug, const float* suu,
+                                                                             __nv_bfloat16* out, int n, float eps) {
+    tf_act_l2::swiglu_rows<tf_act_l2::GS | tf_act_l2::US | tf_act_l2::PROD>(g, g_ld, u, u_ld, sug, suu, out, n, eps);
+}

zig/kernels/cuda/contract_controls.cu

@@ -76,6 +76,8 @@ extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_ctl_qk_round(
 }
 
 // Gate 3: act_rows.cu's tf_softcap_rows keeping only the rounding points in `keep`; the exact check must go red.
+#include "act_round.cuh"
+
 extern "C" __global__ void __launch_bounds__(256) tf_ctl_softcap(const __nv_bfloat16* x, __nv_bfloat16* out, int n,
                                                                   float cap, int keep) {
     tf_act_round::softcap_rows(x, out, n, cap, keep);
@@ -93,8 +95,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_embed_affine(const int*
 
 
 // Gate 3: act_rows.cu's MLX activations keeping only the rounding points in `keep`; the exact check must go red.
-#include "act_round.cuh"
-
 extern "C" __global__ void __launch_bounds__(256) tf_ctl_act_round(const __nv_bfloat16* gate, long long g_ld,
                                                                     const __nv_bfloat16* up, long long u_ld,
                                                                     __nv_bfloat16* out, float* xs, int n, int geglu,
@@ -235,3 +235,55 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_vit_g4_pool(const __nv_
                                                                     float root, int keep) {
     tf_vit_g4::pool_std(x, bias, scale, out, gw, C, root, keep);
 }
+
+// Gate 3: norms.cu's nGPT rows with one rounding point dropped each; their range check must go red.
+#include "norms_l2.cuh"
+
+extern "C" __global__ void __launch_bounds__(tf_l2::NT) tf_ctl_l2_drop_sum(const __nv_bfloat16* x, const __nv_bfloat16* r,
+                                                                           __nv_bfloat16* y, int width, float eps) {
+    tf_l2::l2_rows<0>(x, r, y, width, eps);
+}
+
+#define TF_CTL_SLERP(NAME, KEEP)                                                                                   \
+    extern "C" __global__ void __launch_bounds__(tf_l2::NT) NAME(const __nv_bfloat16* res, const __nv_bfloat16* out, \
+                                                                 const __nv_bfloat16* alpha, float scale,          \
+                                                                 __nv_bfloat16* h, int width, float eps) {         \
+        tf_l2::slerp_rows<KEEP>(res, out, alpha, scale, h, width, eps);                                            \
+    }
+TF_CTL_SLERP(tf_ctl_slerp_drop_a, tf_l2::B)
+TF_CTL_SLERP(tf_ctl_slerp_drop_b, tf_l2::A)
+#undef TF_CTL_SLERP
+
+// Gate 3: qk_rows.cu's tf_qk_rows_l2 with one rounding point dropped each; its range check must go red.
+#include "qk_l2.cuh"
+
+#define TF_CTL_QK_L2(NAME, KEEP)                                                                                   \
+    extern "C" __global__ void __launch_bounds__(tf_qk_l2::WIDE) NAME(                                             \
+        const tf_qk::QkSrc in, const float* qs, const float* cosv, const float* sinv, const float* inv, int rot,   \
+        const int* pos, const int* slot, const tf_qk::QkDst dst, __nv_bfloat16* kc, __nv_bfloat16* vc,            \
+        long long len, int h, int hk, float eps) {                                                                 \
+        tf_qk_l2::rows<KEEP>(in, qs, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps);               \
+    }
+TF_CTL_QK_L2(tf_ctl_qk_l2_drop_rope, tf_qk_l2::NORM)
+TF_CTL_QK_L2(tf_ctl_qk_l2_drop_norm, tf_qk_l2::ROPE)
+#undef TF_CTL_QK_L2
+
+// Gate 3: act_rows.cu's nGPT activations with one rounding point dropped each; their range checks must go red.
+#include "act_l2.cuh"
+
+extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) tf_ctl_gate_l2_drop_gated(const __nv_bfloat16* o, const __nv_bfloat16* g,
+                                                                                     long long g_ld, const float* sg,
+                                                                                     __nv_bfloat16* out, int n, float eps) {
+    tf_act_l2::gate_rows<0>(o, g, g_ld, sg, out, n, eps);
+}
+
+#define TF_CTL_SWIGLU_L2(NAME, KEEP)                                                                               \
+    extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) NAME(                                              \
+        const __nv_bfloat16* g, long long g_ld, const __nv_bfloat16* u, long long u_ld, const float* sug,          \
+        const float* suu, __nv_bfloat16* out, int n, float eps) {                                                  \
+        tf_act_l2::swiglu_rows<KEEP>(g, g_ld, u, u_ld, sug, suu, out, n, eps);                                     \
+    }
+TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_gs, tf_act_l2::US | tf_act_l2::PROD)
+TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_us, tf_act_l2::GS | tf_act_l2::PROD)
+TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_prod, tf_act_l2::GS | tf_act_l2::US)
+#undef TF_CTL_SWIGLU_L2

zig/kernels/cuda/dense_rows.cu

@@ -21,9 +21,9 @@ __device__ __forceinline__ void unpack(const uint4& u, float* f) {
 }
 
 // out[r, n] = x[r] . w[n] for a tile's rows: lane l sums k = 8 l + 256 j in order, then the fixed shuffle tree.
-template <int RT>
+template <int RT, bool SCALED = false> // SCALED: out = float(bf16(x[r] . w[n])) * scale[n] (Dory's logits)
 __device__ __forceinline__ void tile(const __nv_bfloat16* x, int64_t x_stride, const __nv_bfloat16* w, float* out,
-                                     int64_t out_stride, int rows, int n, int k) {
+                                     int64_t out_stride, int rows, int n, int k, const float* scale = nullptr) {
     const int lane = threadIdx.x & 31;
     const int col = blockIdx.x * WARPS + (threadIdx.x >> 5);
     const int r0 = blockIdx.y * RT;
@@ -49,6 +49,7 @@ __device__ __forceinline__ void tile(const __nv_bfloat16* x, int64_t x_stride, c
     for (int r = 0; r < RT; ++r) {
         float v = acc[r];
         for (int m = 16; m > 0; m >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xFFFFFFFFu, v, m));
+        if constexpr (SCALED) v = __fmul_rn(__bfloat162float(__float2bfloat16_rn(v)), scale[col]);
         if (lane == 0 && r0 + r < rows) out[(r0 + r) * out_stride + col] = v;
     }
 }
@@ -68,3 +69,16 @@ TF_DENSE_ROWS(1)
 TF_DENSE_ROWS(4)
 TF_DENSE_ROWS(16)
 TF_DENSE_ROWS(32)
+
+// tf_dense_rows with the scaled epilogue: out = float(bf16(x[r] . w[n])) * scale[n], the same sums.
+#define TF_DENSE_ROWS_SCALED(RT)                                                                                       \
+    extern "C" __global__ void __launch_bounds__(WARPS * 32) tf_dense_rows_scaled_##RT(                              \
+        const __nv_bfloat16* x, long long x_stride, const __nv_bfloat16* w, const float* scale, float* out,           \
+        long long out_stride, int rows, int n, int k) {                                                                \
+        tile<RT, true>(x, x_stride, w, out, out_stride, rows, n, k, scale);                                            \
+    }
+
+TF_DENSE_ROWS_SCALED(1)
+TF_DENSE_ROWS_SCALED(4)
+TF_DENSE_ROWS_SCALED(16)
+TF_DENSE_ROWS_SCALED(32)