diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 75bd900cda..bc8e43ca52 100644 --- a/docs/kernel-contract.md +++ b/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 diff --git a/zig/build/cuda_row_contracts.zig b/zig/build/cuda_row_contracts.zig index 2bf71d7940..cc67382e91 100644 --- a/zig/build/cuda_row_contracts.zig +++ b/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, diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 91b3171bbf..7cb455c713 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/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(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 diff --git a/zig/kernels/cuda/norms.cu b/zig/kernels/cuda/norms.cu index fa214406c6..4426f43d6a 100644 --- a/zig/kernels/cuda/norms.cu +++ b/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(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(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(res, out, alpha, scale, h, width, eps); +} diff --git a/zig/kernels/cuda/norms_l2.cuh b/zig/kernels/cuda/norms_l2.cuh new file mode 100644 index 0000000000..8296101195 --- /dev/null +++ b/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 +#include + +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 +__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(&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(x + r). +template +__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(x + at)[c] : make_uint4(0, 0, 0, 0); + rp[j] = live && r != nullptr ? reinterpret_cast(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(__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(y + at)[c] = pack8(v[j]); + } +} + +// y = bf16(s * inv_l2(s)), s = x (r null) or point(x + r): a block a row of `width`. +template +__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(blockIdx.x) * width; + const int chunks = width / 8; + float v[MOST][8]; + load(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(l2(res)), b = point(l2(out)), lr = |alpha scale| (fp32 steps). +template +__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(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(__fmul_rn(a[j][i], ia)); + const float bv = point(__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 diff --git a/zig/src/core/contract/dory_cuda_records.zig b/zig/src/core/contract/dory_cuda_records.zig new file mode 100644 index 0000000000..955c80e94a --- /dev/null +++ b/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 } } }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 971e52151c..453ae2c721 100644 --- a/zig/src/core/contract/registry.zig +++ b/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 }; diff --git a/zig/src/core/l2_ref.zig b/zig/src/core/l2_ref.zig new file mode 100644 index 0000000000..d6fbbe2a32 --- /dev/null +++ b/zig/src/core/l2_ref.zig @@ -0,0 +1,151 @@ +//! nGPT rows in float64 (Dory's L2 norms and slerp residual), as ranges over an fp32 run at bf16 rounding points. +const std = @import("std"); +const bf = @import("bf16_round.zig"); + +/// A value's least and largest ends; NaN ends for a NaN. +pub const Iv = [2]f64; + +/// An fp32 run's stray: an inverse norm's relative spread, and each fp32 product or sum's (0, 0: float64's points). +pub const Stray = struct { spread: f64 = 0, u: f64 = 0 }; + +/// Sums of squares in float64; a test may pass its own `squares` (another accumulator, gate 4). +pub const Exact = struct { + pub fn squares(_: Exact, v: []const f64) f64 { + var s: f64 = 0; + for (v) |x| s += x * x; + return s; + } +}; + +pub fn point(v: f64) Iv { + return .{ v, v }; +} + +/// a and b in order; NaN in either gives NaN for both. +pub fn order(a: f64, b: f64) Iv { + if (std.math.isNan(a) or std.math.isNan(b)) return @splat(std.math.nan(f64)); + return .{ @min(a, b), @max(a, b) }; +} + +/// The ends rounded to bf16 (rounding is monotone). +pub fn bf16(r: Iv) Iv { + return order(bf.bf16(r[0]), bf.bf16(r[1])); +} + +pub fn bf16If(round: bool, r: Iv) Iv { + return if (round) bf16(r) else r; +} + +/// The ends of x * y over both intervals, widened by u of the largest magnitude (one fp32 product). +pub fn mul(x: Iv, y: Iv, u: f64) Iv { + var r = Iv{ std.math.inf(f64), -std.math.inf(f64) }; + for (x) |a| for (y) |b| { + const p = a * b; + if (std.math.isNan(p)) return @splat(std.math.nan(f64)); + r = .{ @min(r[0], p), @max(r[1], p) }; + }; + const m = @max(@abs(r[0]), @abs(r[1])) * u; + return .{ r[0] - m, r[1] + m }; +} + +/// 1 / sqrt(sum(v^2) + eps) over every v in its interval, and the root's and quotient's spread `sp`. +pub fn invL2(v: []const Iv, eps: f64, sums: anytype, sp: f64, small: []f64, big: []f64) Iv { + for (v, small, big) |x, *s, *b| { + const lo, const hi = order(@abs(x[0]), @abs(x[1])); + s.* = if (x[0] <= 0 and x[1] >= 0) 0 else lo; + b.* = hi; + } + return .{ (1 - sp) / @sqrt(sums.squares(big) + eps), (1 + sp) / @sqrt(sums.squares(small) + eps) }; +} + +/// Every row value times the inverse norm over v, in place; bf16 with `round`. +pub fn normalize(v: []Iv, eps: f64, sums: anytype, st: Stray, round: bool, small: []f64, big: []f64) void { + const inv = invL2(v, eps, sums, st.spread, small, big); + for (v) |*x| x.* = bf16If(round, mul(x.*, inv, st.u)); +} + +/// tf_l2_rows: y = bf16(l2(s)), s = x, or x + r (fp32) rounded to bf16 with `sum`; y as ranges. +pub fn l2Rows(x: []const f64, r: ?[]const f64, sum: bool, eps: f64, sums: anytype, st: Stray, y: []Iv, small: []f64, big: []f64) void { + for (y, x, 0..) |*o, v, i| { + const s: f64 = if (r) |rr| @as(f32, @floatCast(v + rr[i])) else v; // the fp32 add rounds the exact sum once + o.* = point(if (sum and r != null) bf.bf16(s) else s); + } + normalize(y, eps, sums, st, true, small, big); +} + +/// The slerp's rounding points: the normalized residual a, the normalized branch b. +pub const Points = struct { a: bool = true, b: bool = true }; + +/// tf_slerp_rows: h = bf16(l2(a + lr (b - a))), a = l2(res), b = l2(out) (bf16 at points), lr = |fp32(alpha scale)|. +pub fn slerp(res: []const f64, out: []const f64, alpha: []const f64, scale: f32, p: Points, eps: f64, sums: anytype, st: Stray, h: []Iv, scratch: []Iv, small: []f64, big: []f64) void { + const a = h; + const b = scratch; + for (a, b, res, out) |*x, *y, r, o| { + x.* = point(r); + y.* = point(o); + } + normalize(a, eps, sums, st, p.a, small, big); + normalize(b, eps, sums, st, p.b, small, big); + for (a, b, alpha) |*x, y, al| { + const lr = @abs(@as(f64, @as(f32, @floatCast(al * scale)))); + var z = Iv{ std.math.inf(f64), -std.math.inf(f64) }; + var nan = false; + for (x.*) |av| for (y) |bv| { // linear in a and b: the corners bound it + const m = lr * (bv - av); + const zc = av + m; + nan = nan or std.math.isNan(zc); + const e = if (std.math.isFinite(zc)) (@abs(zc) + 2 * @abs(m)) * st.u * 1.01 else 0; // b - a, lr d, a + m + z = .{ @min(z[0], zc - e), @max(z[1], zc + e) }; + }; + x.* = if (nan) @splat(std.math.nan(f64)) else z; + } + normalize(a, eps, sums, st, true, small, big); +} + +test "points: l2 rows are x / |x| in bf16, the sum rounds first, the slerp's ends meet at a and b" { + var x = [_]f64{ 3, 4, 0, 0 }; + var y: [4]Iv = undefined; + var s: [4]f64 = undefined; + var g: [4]f64 = undefined; + l2Rows(&x, null, true, 0, Exact{}, .{}, &y, &s, &g); + try std.testing.expectEqual(bf.bf16(0.6), y[0][0]); + try std.testing.expectEqual(bf.bf16(0.8), y[1][1]); + const r = [_]f64{ 0x1p-9, 0, 0, 0 }; + var one = [_]f64{ 1, 0, 0, 0 }; + l2Rows(&one, &r, true, 0, Exact{}, .{}, &y, &s, &g); + try std.testing.expectEqual(@as(f64, 1), y[0][0]); // bf16(1 + 2^-9) = 1 before the norm + var h: [4]Iv = undefined; + var t: [4]Iv = undefined; + const alpha = [_]f64{ 0, 0, 0, 0 }; + slerp(&x, &[_]f64{ 0, 0, 1, 0 }, &alpha, 1, .{}, 0, Exact{}, .{}, &h, &t, &s, &g); + try std.testing.expectEqual(bf.bf16(0.6), h[0][0]); // lr 0: the residual itself + const ones = [_]f64{ -1, 1, 1, 1 }; // |alpha scale| = 1: the branch + slerp(&x, &[_]f64{ 0, 0, 1, 0 }, &ones, 1, .{}, 0, Exact{}, .{}, &h, &t, &s, &g); + try std.testing.expectEqual(@as(f64, 1), h[2][0]); + try std.testing.expectEqual(@as(f64, 0), h[0][1]); +} + +test "an fp32 run's ranges hold the points and stay far inside a bf16 step" { + var prng = std.Random.DefaultPrng.init(11); + var res: [256]f64 = undefined; + var out: [256]f64 = undefined; + var alpha: [256]f64 = undefined; + for (&res, &out, &alpha) |*a, *b, *c| { + a.* = bf.bf16(prng.random().floatNorm(f64)); + b.* = bf.bf16(4 * prng.random().floatNorm(f64)); + c.* = bf.bf16(0.3 * prng.random().floatNorm(f64)); + } + var hp: [256]Iv = undefined; + var hr: [256]Iv = undefined; + var t: [256]Iv = undefined; + var s: [256]f64 = undefined; + var g: [256]f64 = undefined; + slerp(&res, &out, &alpha, 2.8248587, .{}, 1e-6, Exact{}, .{}, &hp, &t, &s, &g); + slerp(&res, &out, &alpha, 2.8248587, .{}, 1e-6, Exact{}, .{ .spread = 0x1p-20, .u = 0x1p-24 }, &hr, &t, &s, &g); + var ranged: usize = 0; + for (hp, hr) |p, r| { + try std.testing.expect(r[0] <= p[0] and p[0] <= r[1]); + if (r[0] != r[1]) ranged += 1; + } + try std.testing.expect(ranged < 32); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index eef4237f7f..59f6f2fcfb 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -30,6 +30,7 @@ pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); +pub const l2_ref = @import("l2_ref.zig"); pub const act_ref = @import("act_ref.zig"); pub const bf16_round = @import("bf16_round.zig"); pub const embed_ref = @import("embed_ref.zig"); @@ -70,6 +71,7 @@ test { _ = registry; _ = tuning; _ = norms_ref; + _ = l2_ref; _ = act_ref; _ = bf16_round; _ = embed_ref; diff --git a/zig/src/cuda/norms.zig b/zig/src/cuda/norms.zig index 4814b42240..4680ea92da 100644 --- a/zig/src/cuda/norms.zig +++ b/zig/src/cuda/norms.zig @@ -54,6 +54,8 @@ pub const Norms = struct { gated_fns: [3]Function, // head dim 64, 128, 256 rms_sums_fn: Function, tail_fns: [2]Function, // each norm rounded once, MLX's bf16 steps + l2_fn: Function, + slerp_fn: Function, pub fn load(d: *const Driver) !Norms { if (!kernels.available) return error.BuiltWithoutKernels; @@ -70,6 +72,8 @@ pub const Norms = struct { .gated_fns = .{ try mod.function("tf_gated_rms_rows_64"), try mod.function("tf_gated_rms_rows_128"), try mod.function("tf_gated_rms_rows_256") }, .rms_sums_fn = try mod.function("tf_rms_rows_sums"), .tail_fns = .{ try mod.function("tf_add_rms_rows_bf16h_tail"), try mod.function("tf_add_rms_rows_bf16h_tail_round") }, + .l2_fn = try mod.function("tf_l2_rows"), + .slerp_fn = try mod.function("tf_slerp_rows"), }; } @@ -168,6 +172,30 @@ pub const Norms = struct { try launch_.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); } + /// nGPT rows: y = bf16(s / sqrt(sum(s^2) + eps)), s = x, or bf16(x + r) with r; width % 64, at most 8,192. + pub fn l2(n: *const Norms, s: Stream, x: u64, r: ?u64, y: u64, rows: usize, width: usize, eps: f32) !void { + if (width % 64 != 0 or !vectorOk(width, &.{ x, r orelse 0, y })) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ x, r orelse 0, y }) |v| a.add(v); + a.add(@as(i32, @intCast(width))); + a.add(eps); + try launch_.launch(n.l2_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); + } + + /// The nGPT residual: h = bf16(l2(a + |alpha scale| (b - a))), a = bf16(l2(res)), b = bf16(l2(out)); bf16 alpha. + pub fn slerp(n: *const Norms, s: Stream, res: u64, out: u64, alpha: u64, scale: f32, h: u64, rows: usize, width: usize, eps: f32) !void { + if (width % 64 != 0 or !vectorOk(width, &.{ res, out, alpha, h })) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ res, out, alpha }) |v| a.add(v); + a.add(scale); + a.add(h); + a.add(@as(i32, @intCast(width))); + a.add(eps); + try launch_.launch(n.slerp_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); + } + /// Per head (dv 64/128/256): out = bf16(silu(z) * rms(y) * w), y/z rows at strides y_ld/z_ld; xs: 64-value sums. pub fn gatedRms(n: *const Norms, s: Stream, y: u64, y_ld: usize, z: u64, z_ld: usize, w: u64, out: u64, xs: ?u64, rows: usize, heads: usize, dv: usize, eps: f32) !void { const which: usize = switch (dv) { diff --git a/zig/tests/cuda/contract_dory.zig b/zig/tests/cuda/contract_dory.zig new file mode 100644 index 0000000000..43c282e223 --- /dev/null +++ b/zig/tests/cuda/contract_dory.zig @@ -0,0 +1,27 @@ +//! Dory's rows kernels on the CUDA contract runner: each registration and control, owned until the run ends. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const l2_k = @import("contract_l2.zig"); + +const g = core.contract.gates; + +pub const Dory = struct { + l2: [std.enums.values(l2_k.Which).len]l2_k.L2 = undefined, + l2s: usize = 0, + + /// Every Dory registration onto `kernels`; the controls' symbols from the contract_controls image `ctl`. + pub fn add(self: *Dory, gpu: check.Gpu, s: cuda.Stream, norms: *const cuda.norms.Norms, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + const a = gpu.gpa; + for (std.enums.values(l2_k.Which), &self.l2) |w, *t| { + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .norms = norms, .ctl = if (l2_k.control(w) != null) try ctl.function(l2_k.symbol(w)) else null }; + self.l2s += 1; + } + for (&self.l2) |*t| try kernels.append(a, t.kernel()); + } + + pub fn deinit(self: *Dory) void { + for (self.l2[0..self.l2s]) |*t| t.deinit(); + } +}; diff --git a/zig/tests/cuda/contract_l2.zig b/zig/tests/cuda/contract_l2.zig new file mode 100644 index 0000000000..d290b75788 --- /dev/null +++ b/zig/tests/cuda/contract_l2.zig @@ -0,0 +1,276 @@ +//! norms.cu's nGPT rows under the contract (tf_l2_rows with and without its add, tf_slerp_rows), and their controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const rows = core.contract.rows; +const ref = core.l2_ref; +const bfr = core.bf16_round; + +/// The registrations, and the controls dropping one rounding point each. +pub const Which = enum { l2, l2_add, slerp, drop_sum, drop_a, drop_b }; + +/// The norm's sum chain: 32 fma a thread over 16-byte loads, then 256 threads' tree (norms_tail_check's depth). +const chain = 42; + +fn depth(_: c.Shape) usize { + return chain; +} + +/// An fp32 inverse norm's relative spread from float64's: the sum's gamma halved by the root, the fp32 steps. +pub fn spread() f64 { + return c.gamma(chain, 0x1p-24) / 2 + 8 * 0x1p-24; +} + +const why = "an output follows its sums of squares as ss^-1/2 through bf16 rounding points, so a bf16 sum's scale error cancels in the row's norm, and an fp16 sum overflows to a zero output"; + +/// Dory's hidden rows: a decode window and a short verify (2560 wide). +const served = [_]c.Shape{ .{ .rows = 16, .n = 2560, .k = 2560 }, .{ .rows = 3, .n = 2560, .k = 2560 } }; + +pub fn spec(which: Which) c.Spec { + return .{ + .name = switch (which) { + .l2 => "cuda nGPT rows tf_l2_rows (a loop's exit: bf16(l2(h)))", + .l2_add => "cuda nGPT rows tf_l2_rows with its add (a loop's entry: bf16(l2(bf16(state + input))))", + .slerp => "cuda nGPT rows tf_slerp_rows (the residual: bf16(l2(a + lr (b - a))), a and b bf16(l2(.)))", + .drop_sum => "control: tf_l2_rows without the sum's rounding before the norm", + .drop_a => "control: tf_slerp_rows without the residual's normalized rounding", + .drop_b => "control: tf_slerp_rows without the branch's normalized rounding", + }, + .accum = .f32, + .depth = depth, + .ops = 12, + .mids = if (isSlerp(which)) 4 else 2, // a bf16 point (a, b) decided the other way by an fp32 inverse: two half-ulps + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &served, + .sums = .{ .f16 = 40, .bf16 = 0, .why = why }, + }; +} + +pub fn control(which: Which) ?g.Gate { + return switch (which) { + .l2, .l2_add, .slerp => null, + else => .shapes, + }; +} + +/// The kernel a registration launches: the served ones through cuda/norms.zig, the controls by symbol. +pub fn symbol(which: Which) [:0]const u8 { + return switch (which) { + .drop_sum => "tf_ctl_l2_drop_sum", + .drop_a => "tf_ctl_slerp_drop_a", + .drop_b => "tf_ctl_slerp_drop_b", + else => unreachable, + }; +} + +fn isSlerp(which: Which) bool { + return which == .slerp or which == .drop_a or which == .drop_b; +} + +/// fp32 alpha_init / std at Dory's 0.05 and 0.0177. +const alpha_scale: f32 = @floatCast(0.05 / 0.0177); + +/// Sums of squares in an accumulator and order (gate 4's simulation). +const Sim = struct { + acc: c.Accum, + order: rows.Order, + buf: []f64, + + pub fn squares(s: Sim, v: []const f64) f64 { + for (s.buf[0..v.len], v) |*t, x| t.* = x * x; + return rows.sumIn(s.acc, s.order, s.buf[0..v.len]); + } +}; + +/// got in [lo, hi]: its bits where the ends are one value, else between the ends (-0 below +0); NaN for NaN. +fn inRange(got: u16, r: ref.Iv) bool { + const v: f64 = bfr.value(got); + if (std.math.isNan(r[0]) or std.math.isNan(r[1])) return std.math.isNan(v); + if (@as(u64, @bitCast(r[0])) == @as(u64, @bitCast(r[1]))) return got == bfr.bits(@floatCast(r[0])); + return atMost(r[0], v) and atMost(v, r[1]); +} + +fn atMost(a: f64, b: f64) bool { + return a < b or (a == b and (std.math.signbit(a) or !std.math.signbit(b))); +} + +pub const L2 = struct { + rig: rig_.Rig, + which: Which, + norms: *const cuda.norms.Norms, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + n: usize = 0, + width: usize = 0, + x: []f64 = &.{}, // x or res + r: []f64 = &.{}, // r or out + alpha: []f64 = &.{}, + dev: [3]u64 = undefined, + out: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *L2) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(L2); + + pub fn kernel(self: *L2) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = &vt, .control = control(self.which) }; + } + + fn bfs(rng: std.Random, v: []f64, scale: f32) void { + for (v, 0..) |*x, i| x.* = bfr.value(bfr.bits((rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1))); + } + + /// Gate 4's rows: integers whose squares sum to about 2^20, exact in fp32 in any order. + fn integers(rng: std.Random, v: []f64) void { + for (v) |*x| x.* = @as(f64, @floatFromInt(rng.uintLessThan(u32, 65))) - 32; + } + + fn upload(self: *L2, v: []const f64) !u64 { + const a = self.arena.?.allocator(); + const b = try a.alloc(u16, v.len); + for (b, v) |*h, x| h.* = bfr.bits(@floatCast(x)); + return self.rig.input(std.mem.sliceAsBytes(b)); + } + + pub fn prepare(self: *L2, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const width = shape.k; + if (width % 64 != 0 or width > 8192 or shape.rows == 0) return false; + self.rig.reset(); + if (self.arena) |*a| a.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.width = width; + self.n = shape.rows * width; + self.x = try a.alloc(f64, self.n); + self.r = try a.alloc(f64, self.n); + self.alpha = try a.alloc(f64, width); + bfs(rng, self.x, 1); + bfs(rng, self.r, 4); + bfs(rng, self.alpha, 0.4); // |alpha scale| to about 1.1, lr past 1 on a few + if (fill == .nonfinite) { + const planted = [_]struct { usize, bool, usize, f64 }{ .{ 0, true, width / 3, std.math.nan(f64) }, .{ 1, false, width / 2, std.math.inf(f64) }, .{ 2, true, width - 1, -std.math.inf(f64) }, .{ 3, false, 0, std.math.inf(f64) }, .{ 3, false, width - 1, -std.math.inf(f64) } }; + for (planted) |p| if (p[0] < shape.rows) { + (if (p[1]) self.x else self.r)[p[0] * width + p[2]] = p[3]; + }; + } + if (fill == .stress) { + integers(rng, self.x); + integers(rng, self.r); + } + const take_r = self.which != .l2; + self.dev = .{ try self.upload(self.x), if (take_r) try self.upload(self.r) else 0, if (isSlerp(self.which)) try self.upload(self.alpha) else 0 }; + self.out = try self.rig.output(self.n * 2); + return true; + } + + pub fn run(self: *L2) !void { + const s = self.rig.s; + const rws = self.n / self.width; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + args.add(self.dev[0]); + args.add(self.dev[1]); + if (isSlerp(self.which)) { + args.add(self.dev[2]); + args.add(alpha_scale); + } + args.add(self.out); + args.add(@as(c_int, @intCast(self.width))); + args.add(@as(f32, 1e-6)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rws), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &args); + } else if (isSlerp(self.which)) { + try self.norms.slerp(s, self.dev[0], self.dev[1], self.dev[2], alpha_scale, self.out, rws, self.width, 1e-6); + } else try self.norms.l2(s, self.dev[0], if (self.which == .l2) null else self.dev[1], self.out, rws, self.width, 1e-6); + try s.synchronize(); + } + + pub fn count(self: *L2) usize { + return self.n; + } + + /// Each row's outputs as ranges under `st`, sums of squares by `sums`. + fn rowsOf(self: *L2, a: std.mem.Allocator, sums: anytype, st: ref.Stray, out: []ref.Iv) !void { + const w = self.width; + const scratch = try a.alloc(ref.Iv, w); + const small = try a.alloc(f64, w); + const big = try a.alloc(f64, w); + for (0..self.n / w) |row| { + const at = row * w; + const x = self.x[at..][0..w]; + const r = self.r[at..][0..w]; + const o = out[at..][0..w]; + switch (self.which) { + .l2 => ref.l2Rows(x, null, true, 1e-6, sums, st, o, small, big), + .l2_add, .drop_sum => ref.l2Rows(x, r, true, 1e-6, sums, st, o, small, big), + else => ref.slerp(x, r, self.alpha, alpha_scale, .{}, 1e-6, sums, st, o, scratch, small, big), + } + } + } + + fn points(self: *L2, sums: anytype, out: []f64) !void { + var arena: std.heap.ArenaAllocator = .init(self.rig.gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const iv = try a.alloc(ref.Iv, self.n); + try self.rowsOf(a, sums, .{}, iv); + for (out, iv) |*o, v| o.* = v[0]; + } + + pub fn reference(self: *L2, want: []f64, mag: []f64) !void { + try self.points(ref.Exact{}, want); + for (mag, want) |*m, v| m.* = @abs(v) / 2; // v ~ ss^-1/2 + } + + pub fn simulate(self: *L2, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try self.rig.gpu.gpa.alloc(f64, self.width); + defer self.rig.gpu.gpa.free(buf); + try self.points(Sim{ .acc = acc, .order = order, .buf = buf }, out); + } + + /// Every output in float64's range over the fp32 run's spread (a dropped point leaves it). + pub fn check(self: *L2, _: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const got = try a.alloc(u16, self.n); + @memcpy(std.mem.sliceAsBytes(got), try self.rig.bytes(self.out, self.n * 2)); + const iv = try a.alloc(ref.Iv, self.n); + try self.rowsOf(a, ref.Exact{}, .{ .spread = spread(), .u = 0x1p-24 }, iv); + for (got, iv, 0..) |h, r, i| if (!inRange(h, r)) + return try std.fmt.bufPrint(&self.msg, "row {d} out[{d}] = {e} outside float64's [{e}, {e}]", .{ i / self.width, i % self.width, bfr.value(h), r[0], r[1] }); + return null; + } +}; + +test "gate 4's rows: integer squares are exact in fp32, fp16 overflows and bf16 misses" { + var terms: [2560]f64 = undefined; + var prng = std.Random.DefaultPrng.init(0x6c32); + for (&terms) |*t| { + const v: f64 = @floatFromInt(@as(i32, @intCast(prng.random().uintLessThan(u32, 65))) - 32); + t.* = v * v; + } + var exact: f64 = 0; + for (terms) |t| exact += t; + try std.testing.expect(exact < 0x1p24 and exact > 0x1p19); + for (std.enums.values(rows.Order)) |o| { + try std.testing.expectEqual(exact, rows.sumIn(.f32, o, &terms)); + try std.testing.expect(std.math.isInf(rows.sumIn(.f16, o, &terms))); + try std.testing.expect(rows.sumIn(.bf16, o, &terms) != exact); + } +} + +test "the range check: a dropped bf16 point leaves the fp32 run's range" { + try std.testing.expect(spread() < 0x1p-18); + const r = ref.Iv{ 0.5, 0.5 }; + try std.testing.expect(inRange(bfr.bits(0.5), r) and !inRange(bfr.bits(0.50390625), r)); + try std.testing.expect(inRange(0x7fc0, @splat(std.math.nan(f64))) and !inRange(0, @splat(std.math.nan(f64)))); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 39a15df860..ec90c01756 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -7,6 +7,7 @@ const options = @import("contract_options"); const check = @import("check.zig"); const rows_k = @import("contract_rows.zig"); const tail_k = @import("contract_norms_tail.zig"); +const dory_k = @import("contract_dory.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const mx_k = @import("contract_mxfp8.zig"); @@ -139,6 +140,9 @@ pub fn run(gpu: check.Gpu) !void { for (tail_which, &tails) |w, *t| t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .norms = &norms, .ctl = if (tail_k.control(w) != null) try mods[5].function(tail_k.symbol(w)) else null }; defer for (&tails) |*t| t.deinit(); for (&tails) |*t| try kernels.append(a, t.kernel()); + var dory: dory_k.Dory = .{}; + defer dory.deinit(); + try dory.add(gpu, s, &norms, mods[5], &kernels); var logits = router_k.Logits{ .rig = .{ .gpu = gpu, .s = s }, .f = try mods[0].function("tf_nemo_router") }; defer logits.deinit(); var logits_tiles = router_k.Logits{ .rig = .{ .gpu = gpu, .s = s }, .f = logits.f, .tiles = try mods[0].function("tf_nemo_router_tiles") };