cuda norms: nGPT rows, L2 norms with eps on the sum and the normalized lerp residual (Dory, #1233)

docs/kernel-contract.md

@@ -149,6 +149,8 @@ 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).
 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_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_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/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,21 @@ 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

zig/kernels/cuda/norms.cu

@@ -520,3 +520,19 @@ extern "C" __global__ void __launch_bounds__(tf_rms_tail::NT) tf_add_rms_rows_bf
     const __nv_bfloat16* w, __nv_bfloat16* h, __nv_bfloat16* y, float* xs, int width, float eps) {
     tf_rms_tail::tail<tf_rms_tail::MLX>(x, r, wb, sc, w, h, y, xs, width, eps);
 }
+
+// nGPT rows (norms_l2.cuh): L2 norms with eps on the sum, and the normalized lerp residual (Dory).
+#include "norms_l2.cuh"
+
+// y = bf16(l2(s)), s = x, or bf16(x + r) with r (a loop's entry): a block a row, width % 64, at most 8,192.
+extern "C" __global__ void __launch_bounds__(tf_l2::NT) tf_l2_rows(const __nv_bfloat16* x, const __nv_bfloat16* r,
+                                                                   __nv_bfloat16* y, int width, float eps) {
+    tf_l2::l2_rows<tf_l2::SUM>(x, r, y, width, eps);
+}
+
+// h = bf16(l2(a + |alpha scale| (b - a))), a = bf16(l2(res)), b = bf16(l2(out)): dory.py's slerp_norm.
+extern "C" __global__ void __launch_bounds__(tf_l2::NT) tf_slerp_rows(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<tf_l2::A | tf_l2::B>(res, out, alpha, scale, h, width, eps);
+}

zig/kernels/cuda/norms_l2.cuh

@@ -0,0 +1,139 @@
+// nGPT rows (norms.cu's tf_l2_rows, tf_slerp_rows): L2 norms with eps on the sum, and the normalized lerp residual.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+namespace tf_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 sum x + r before its norm; the slerp's normalized residual a and branch b.
+enum : int { SUM = 1, A = 2, B = 4 };
+
+__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); }
+
+// v rounded to bf16 where R holds the point P.
+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;
+}
+
+// 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);
+    __syncthreads(); // `warps` is written again by the row's next norm
+    return __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(s, eps)));
+}
+
+// A row's chunks of 8 into v (zeros past the row); with r, v = point<SUM>(x + r).
+template <int R>
+__device__ __forceinline__ void load(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ r,
+                                     int64_t at, int chunks, float (&v)[MOST][8]) {
+    uint4 xp[MOST], rp[MOST];
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) { // every load issued before any is used
+        const int c = threadIdx.x + j * NT;
+        const bool live = c < chunks;
+        xp[j] = live ? reinterpret_cast<const uint4*>(x + at)[c] : make_uint4(0, 0, 0, 0);
+        rp[j] = live && r != nullptr ? reinterpret_cast<const uint4*>(r + at)[c] : make_uint4(0, 0, 0, 0);
+    }
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) {
+        unpack8(xp[j], v[j]);
+        if (r != nullptr) {
+            float rv[8];
+            unpack8(rp[j], rv);
+#pragma unroll
+            for (int i = 0; i < 8; ++i) v[j][i] = point<R, SUM>(__fadd_rn(v[j][i], rv[i]));
+        }
+    }
+}
+
+__device__ __forceinline__ void store(__nv_bfloat16* __restrict__ y, int64_t at, int chunks, const float (&v)[MOST][8]) {
+#pragma unroll
+    for (int j = 0; j < MOST; ++j) {
+        const int c = threadIdx.x + j * NT;
+        if (c < chunks) reinterpret_cast<uint4*>(y + at)[c] = pack8(v[j]);
+    }
+}
+
+// y = bf16(s * inv_l2(s)), s = x (r null) or point<SUM>(x + r): a block a row of `width`.
+template <int R>
+__device__ __forceinline__ void l2_rows(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ r,
+                                        __nv_bfloat16* __restrict__ y, int width, float eps) {
+    const int64_t at = static_cast<int64_t>(blockIdx.x) * width;
+    const int chunks = width / 8;
+    float v[MOST][8];
+    load<R>(x, r, at, chunks, v);
+    const float inv = inv_l2(v, eps);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) v[j][i] = __fmul_rn(v[j][i], inv);
+    store(y, at, chunks, v);
+}
+
+// h = bf16(l2(a + lr (b - a))), a = point<A>(l2(res)), b = point<B>(l2(out)), lr = |alpha scale| (fp32 steps).
+template <int R>
+__device__ __forceinline__ void slerp_rows(const __nv_bfloat16* __restrict__ res, const __nv_bfloat16* __restrict__ out,
+                                           const __nv_bfloat16* __restrict__ alpha, float scale,
+                                           __nv_bfloat16* __restrict__ h, int width, float eps) {
+    const int64_t at = static_cast<int64_t>(blockIdx.x) * width;
+    const int chunks = width / 8;
+    float a[MOST][8], b[MOST][8], lr[MOST][8];
+    load<0>(res, nullptr, at, chunks, a);
+    load<0>(out, nullptr, at, chunks, b);
+    load<0>(alpha, nullptr, 0, chunks, lr);
+    const float ia = inv_l2(a, eps);
+    const float ib = inv_l2(b, eps);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) {
+            const float av = point<R, A>(__fmul_rn(a[j][i], ia));
+            const float bv = point<R, B>(__fmul_rn(b[j][i], ib));
+            const float l = fabsf(__fmul_rn(lr[j][i], scale));
+            a[j][i] = __fadd_rn(av, __fmul_rn(l, __fsub_rn(bv, av))); // past the row: 0 + |0 scale| (0 - 0) = 0
+        }
+    const float iz = inv_l2(a, eps);
+#pragma unroll
+    for (int j = 0; j < MOST; ++j)
+#pragma unroll
+        for (int i = 0; i < 8; ++i) a[j][i] = __fmul_rn(a[j][i], iz);
+    store(h, at, chunks, a);
+}
+
+}  // namespace tf_l2

zig/src/core/contract/dory_cuda_records.zig

@@ -0,0 +1,7 @@
+//! Dory's CUDA launch records (the nGPT rows in norms, qk_rows, act_rows and dense_rows) in registry.cuda.
+const CudaRecord = @import("registry.zig").CudaRecord;
+
+pub const cuda = [_]CudaRecord{
+    .{ .launch = .{ .name = "norms tf_l2_rows", .symbol = "tf_l2_rows", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 48, .static = 32, .bound = 256 }, .{ .regs = 48, .static = 32, .bound = 256 }, .{ .regs = 48, .static = 32, .bound = 256 }, .{ .regs = 64, .static = 32, .bound = 256 }, .{ .regs = 64, .static = 32, .bound = 256 } } },
+    .{ .launch = .{ .name = "norms tf_slerp_rows", .symbol = "tf_slerp_rows", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 113, .static = 32, .bound = 256 }, .{ .regs = 113, .static = 32, .bound = 256 }, .{ .regs = 113, .static = 32, .bound = 256 }, .{ .regs = 113, .static = 32, .bound = 256 }, .{ .regs = 113, .static = 32, .bound = 256 } } },
+};

zig/src/core/contract/registry.zig

@@ -194,7 +194,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{
     .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } },
     .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } },
-} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda;
+} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("dory_cuda_records.zig").cuda;
 
 /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins.
 pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 };