From 0d6ffa11cc82c544e8ea8c92d12a4171bc672bd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:01:56 +0200 Subject: [PATCH 1/9] cuda contract: the softcap control includes the act_round helpers it calls --- zig/kernels/cuda/contract_controls.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 91b3171bbf..fd66e8873c 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, From d2a0e1b2215620c510bb6829d89e903e07bc0a45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:25:30 +0200 Subject: [PATCH 2/9] cuda norms: nGPT rows, L2 norms with eps on the sum and the normalized lerp residual Dory's residual stream lives on the unit sphere, so its rows need two kernels norms.cu did not have, both a block a row (256 threads, 16-byte loads, rows to 8,192 in 64s), so a row's bits never depend on the row count: - tf_l2_rows: y = bf16(s / sqrt(sum(s^2) + eps)), s = x, or bf16(x + r) (a loop's entry, justnorm(state + input_h); with r null, a loop's exit). - tf_slerp_rows: dory.py's slerp_norm, h = bf16(l2(a + |alpha scale| (b - a))) with a = bf16(l2(res)) and b = bf16(l2(out)), the lerp in fp32 steps. Norms.l2 and Norms.slerp launch them; no existing kernel or launch changes. Both are registered with the kernel contract. Their check compares every output with float64's value at the bf16 rounding points as a range over the fp32 run's spread (core/l2_ref.zig); three controls each drop one rounding point (the sum before the norm, the slerp's a, its b) and go red at gate 3. Gate 4's integer rows make an fp16 sum overflow to a zero output (85x and 51x the bound); a bf16 sum's scale cancels in the norm, so that margin is stated, not proven, as for the RMS tails. Launch records from nvcc 13.0's sm_80/86/89/120/121 images. --- docs/kernel-contract.md | 2 + zig/build/cuda_row_contracts.zig | 2 +- zig/kernels/cuda/contract_controls.cu | 18 ++ zig/kernels/cuda/norms.cu | 16 ++ zig/kernels/cuda/norms_l2.cuh | 139 ++++++++++ zig/src/core/contract/dory_cuda_records.zig | 7 + zig/src/core/contract/registry.zig | 2 +- zig/src/core/l2_ref.zig | 151 +++++++++++ zig/src/core/root.zig | 2 + zig/src/cuda/norms.zig | 28 ++ zig/tests/cuda/contract_dory.zig | 27 ++ zig/tests/cuda/contract_l2.zig | 276 ++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 4 + 13 files changed, 672 insertions(+), 2 deletions(-) create mode 100644 zig/kernels/cuda/norms_l2.cuh create mode 100644 zig/src/core/contract/dory_cuda_records.zig create mode 100644 zig/src/core/l2_ref.zig create mode 100644 zig/tests/cuda/contract_dory.zig create mode 100644 zig/tests/cuda/contract_l2.zig 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 fd66e8873c..7cb455c713 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -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") }; From 8a2d188f9285eb8acf7ddb9820592d32971a9ceb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:52:38 +0200 Subject: [PATCH 3/9] cuda qk_rows: Dory's attention inputs, RoPE first, then each head's L2 norm and q's scale tf_qk_rows_l2 takes dory.py's order: rotate-half RoPE on the projection's bf16 rows (rounded to bf16), then each head's L2 norm with eps on the sum, then q times an fp32 scale vector ((sqk f)^2, formed on the host as torch forms it) after the normalized row's rounding; k is its normalized row, v as is, and k and v go into the caches as tf_qk_rows writes them. Shape.l2 and Shape.q_scale select it; every other call keeps its kernel. It is registered with the kernel contract: every output against float64's range at its rounding points (cosf's ulps on the fp32 angle pos * inv, the inverse's fp32 spread), at head dim 256, 16 query heads over 4, a third of the rows past position 199,800. Two controls each drop one rounding point (the rotated head, q's normalized head) and go red at gate 3. Gate 4's integer heads (no rotation) overflow an fp16 sum to a zero output. core/l2_ref.zig's intervals now keep infinite ends and widen by fp32's least subnormal step, so a planted infinity leaves only the outputs float64 makes non-finite. --- docs/kernel-contract.md | 2 + zig/build/cuda_row_contracts.zig | 2 +- zig/kernels/cuda/contract_controls.cu | 14 ++ zig/kernels/cuda/qk_l2.cuh | 76 +++++++ zig/kernels/cuda/qk_rows.cu | 12 + zig/src/core/contract/dory_cuda_records.zig | 1 + zig/src/core/l2_ref.zig | 36 ++- zig/src/cuda/qk_rows.zig | 41 +++- zig/tests/cuda/contract_dory.zig | 19 +- zig/tests/cuda/contract_qk_l2.zig | 236 ++++++++++++++++++++ 10 files changed, 431 insertions(+), 8 deletions(-) create mode 100644 zig/kernels/cuda/qk_l2.cuh create mode 100644 zig/tests/cuda/contract_qk_l2.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index bc8e43ca52..6317957b2c 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -151,6 +151,8 @@ output within float64's range over the spread its fp32 steps and tensor-core sum 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). 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 cc67382e91..b17a80b606 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_l2.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_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 7cb455c713..6c40bfcdb3 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -253,3 +253,17 @@ extern "C" __global__ void __launch_bounds__(tf_l2::NT) tf_ctl_l2_drop_sum(const 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(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 diff --git a/zig/kernels/cuda/qk_l2.cuh b/zig/kernels/cuda/qk_l2.cuh new file mode 100644 index 0000000000..3a19083d5b --- /dev/null +++ b/zig/kernels/cuda/qk_l2.cuh @@ -0,0 +1,76 @@ +// qk_rows.cu's tf_qk_rows_l2 (Dory): RoPE first, then each head's L2 norm (eps on the sum), q times a scale vector. + +#pragma once + +#include "qk_rows.cuh" + +namespace tf_qk_l2 { + +using namespace tf_qk; + +constexpr int WIDE = 512; // head_dim at most + +// The rounding points: the rotated head, and q's normalized head before its scale (k's is its output). +enum : int { ROPE = 1, NORM = 2 }; + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +template +__device__ __forceinline__ float point(float v) { + if constexpr ((R & P) != 0) return bf(v); + else return v; +} + +// 1 / sqrt(sum + eps) of a head's squares: each warp's xor tree, then the warps in order (tf_qk_rows' order). +__device__ __forceinline__ float inv_l2(float y, int dim, float eps) { + __shared__ float warps[WIDE / 32]; + float part = __fmul_rn(y, y); + 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 = 0.f; + for (int w = 0; w < dim / 32; ++w) s = __fadd_rn(s, warps[w]); + return __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(s, eps))); +} + +// A block a (row, head), a thread an element: q = bf16(qs * point(l2(point(rope(x))))), k without qs. +template +__device__ __forceinline__ void rows(const QkSrc& in, const float* qs, const float* cosv, const float* sinv, + const float* inv, int rot, const int* pos, const int* slot, const QkDst& dst, + __nv_bfloat16* kc, __nv_bfloat16* vc, long long len, int h, int hk, float eps) { + const int d = threadIdx.x; + const int dim = blockDim.x; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y; + const __nv_bfloat16* src = head_src(in, r, hd, h, hk); + const int p = pos[r]; + const int64_t pl = p < len ? p : p % len; + const int64_t at = slot != nullptr ? static_cast(slot[r]) * len + pl : pl; + if (hd >= h + hk) { // v: as is + const int k = hd - h - hk; + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + d] = src[d]; + if (vc != nullptr) vc[(at * hk + k) * dim + d] = src[d]; + return; + } + float y = __bfloat162float(src[d]); + const int half = rot / 2; + if ((cosv != nullptr || inv != nullptr) && d < rot) { + const float yp = __bfloat162float(src[d < half ? d + half : d - half]); + float c, sn; + rope_at(cosv, sinv, inv, p, rot, d, c, sn); + y = __fadd_rn(__fmul_rn(y, c), __fmul_rn(d < half ? -yp : yp, sn)); + } + y = point(y); + const float n = __fmul_rn(y, inv_l2(y, dim, eps)); + if (hd < h) { + const float out = qs != nullptr ? __fmul_rn(qs[hd * dim + d], point(n)) : n; + dst.q[r * dst.q_ld + hd * dst.q_hs + d] = __float2bfloat16_rn(out); + } else { + const int k = hd - h; + const __nv_bfloat16 out = __float2bfloat16_rn(n); + dst.k[r * dst.k_ld + k * dst.k_hs + d] = out; + if (kc != nullptr) kc[(at * hk + k) * dim + d] = out; + } +} + +} // namespace tf_qk_l2 diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu index b56c7ff642..01bda9e16a 100644 --- a/zig/kernels/cuda/qk_rows.cu +++ b/zig/kernels/cuda/qk_rows.cu @@ -156,3 +156,15 @@ extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_qk_rows_round tf_qk_round::rows(in, nw, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps, tf_qk_round::MLX | (vnorm ? tf_qk_round::VNORM : 0)); } + +// Dory's attention inputs (qk_l2.cuh): RoPE first, then each head's L2 norm, q times a scale vector, v as is. +#include "qk_l2.cuh" + +// q = bf16(qs * bf16(l2(bf16(rope(x))))), k = bf16(l2(bf16(rope(x)))), v as is; k and v also into kc/vc. +extern "C" __global__ void __launch_bounds__(tf_qk_l2::WIDE) tf_qk_rows_l2(const QkSrc in, const float* qs, const float* cosv, + const float* sinv, const float* inv, int rot, + const int* pos, const int* slot, const QkDst dst, + __nv_bfloat16* kc, __nv_bfloat16* vc, long long len, + int h, int hk, float eps) { + tf_qk_l2::rows(in, qs, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps); +} diff --git a/zig/src/core/contract/dory_cuda_records.zig b/zig/src/core/contract/dory_cuda_records.zig index 955c80e94a..dbeac45d2d 100644 --- a/zig/src/core/contract/dory_cuda_records.zig +++ b/zig/src/core/contract/dory_cuda_records.zig @@ -4,4 +4,5 @@ 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 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows_l2 at dim 256", .symbol = "tf_qk_rows_l2", .image = "qk_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 } } }, }; diff --git a/zig/src/core/l2_ref.zig b/zig/src/core/l2_ref.zig index d6fbbe2a32..33be0254df 100644 --- a/zig/src/core/l2_ref.zig +++ b/zig/src/core/l2_ref.zig @@ -44,10 +44,16 @@ pub fn mul(x: Iv, y: Iv, u: f64) Iv { 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 }; + return widen(r, u); } +/// r widened by u of its largest magnitude and, with u, fp32's least subnormal step; infinite ends stay as they are. +pub fn widen(r: Iv, u: f64) Iv { + const m = @max(@abs(r[0]), @abs(r[1])) * u + @as(f64, if (u > 0) 0x1p-149 else 0); + return .{ if (std.math.isFinite(r[0])) r[0] - m else r[0], if (std.math.isFinite(r[1])) r[1] + m else r[1] }; +} + + /// 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| { @@ -149,3 +155,29 @@ test "an fp32 run's ranges hold the points and stay far inside a bf16 step" { } try std.testing.expect(ranged < 32); } + +/// tf_qk_rows_l2's rounding points: the rotated head, q's normalized head before its scale. +pub const QkPoints = struct { rope: bool = true, norm: bool = true }; + +/// tf_qk_rows_l2's head as ranges: RoPE (cos, sin within `trig`), bf16; its L2 norm n; q bf16(qs n), k bf16(n). +pub fn qkHead(x: []const f64, cos: []const f64, sin: []const f64, rot: usize, qs: ?[]const f64, p: QkPoints, eps: f64, sums: anytype, st: Stray, trig: f64, y: []Iv, small: []f64, big: []f64) void { + const half = rot / 2; + for (y, x, 0..) |*o, v, i| { + if (i >= rot) { + o.* = point(v); + continue; + } + const xp = if (i < half) -x[i + half] else x[i - half]; + const cs = order(cos[i] - trig * @abs(cos[i]), cos[i] + trig * @abs(cos[i])); + const sn = order(sin[i] - trig * @abs(sin[i]), sin[i] + trig * @abs(sin[i])); + const a = mul(point(v), cs, st.u); + const b = mul(point(xp), sn, st.u); + o.* = bf16If(p.rope, widen(order(a[0] + b[0], a[1] + b[1]), st.u)); + } + const inv = invL2(y, eps, sums, st.spread, small, big); + for (y, 0..) |*o, i| { + const n = mul(o.*, inv, st.u); + o.* = if (qs) |q| bf16(mul(bf16If(p.norm, n), point(q[i]), st.u)) else bf16(n); + } +} + diff --git a/zig/src/cuda/qk_rows.zig b/zig/src/cuda/qk_rows.zig index 528baa233b..4826a3129d 100644 --- a/zig/src/cuda/qk_rows.zig +++ b/zig/src/cuda/qk_rows.zig @@ -30,6 +30,10 @@ pub const Shape = struct { round: bool = false, /// v by a unit RMS norm, bf16(x * inv); else v as is. v_norm: bool = false, + /// Dory's order (tf_qk_rows_l2): RoPE first, then each head's L2 norm (eps on the sum), v as is; no norm weights. + l2: bool = false, + /// With l2: q times this fp32 vector [h, dim] after its normalized row's rounding (null: q as k). + q_scale: ?u64 = null, }; /// qk_rows.cu's QkW. @@ -61,12 +65,14 @@ pub const QkRows = struct { /// tf_qk_rows_once and tf_qk_rows_round: dims past 256, a unit v norm, MLX's rounding points. once_fn: Function, round_fn: Function, + /// tf_qk_rows_l2: Dory's attention inputs. + l2_fn: Function, pub fn load(d: *const Driver) !QkRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qk_rows); errdefer mod.unload(); - return .{ .mod = mod, .fn_ = try mod.function("tf_qk_rows"), .warp_fn = try mod.function("tf_qk_rows_w"), .once_fn = try mod.function("tf_qk_rows_once"), .round_fn = try mod.function("tf_qk_rows_round") }; + return .{ .mod = mod, .fn_ = try mod.function("tf_qk_rows"), .warp_fn = try mod.function("tf_qk_rows_w"), .once_fn = try mod.function("tf_qk_rows_once"), .round_fn = try mod.function("tf_qk_rows_round"), .l2_fn = try mod.function("tf_qk_rows_l2") }; } pub fn unload(g: *QkRows) void { @@ -94,6 +100,7 @@ pub const QkRows = struct { if (roped and shape.rot == 0) return error.HeadShape; if (src.v == 0 and out.vc != null) return error.HeadShape; if (n == 0) return; + if (shape.l2) return g.rowsL2(f, s, src, hs, shape, pos, slot, out, n); var a: launch_.Args = .{}; a.add(src); const wsize: usize = if (shape.norms_bf16) 2 else 4; @@ -130,4 +137,34 @@ pub const QkRows = struct { const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast(heads), .z = 1 }; try launch_.launch(g.fn_, .{ .grid = grid, .block = .{ .x = @intCast(hs.dim) } }, s, &a); } -}; + + /// tf_qk_rows_l2 (or `f`, a control with its arguments) on the shape's RoPE, q scale and caches. + fn rowsL2(g: *const QkRows, f: ?Function, s: Stream, src: Src, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + if (shape.norms != null or shape.round or shape.v_norm) return error.HeadShape; + var a: launch_.Args = .{}; + a.add(src); + a.add(shape.q_scale orelse 0); + a.add(shape.cos orelse 0); + a.add(shape.sin orelse 0); + a.add(shape.inv orelse 0); + const roped = shape.cos != null or shape.inv != null; + a.add(@as(i32, @intCast(if (roped) shape.rot else 0))); + a.add(pos); + a.add(slot orelse 0); + const dim: i64 = @intCast(hs.dim); + const qld: i64 = if (out.q_ld != 0) @intCast(out.q_ld) else @as(i64, @intCast(hs.h)) * dim; + const qhs: i64 = if (out.q_hs != 0) @intCast(out.q_hs) else dim; + const kld: i64 = if (out.kv_ld != 0) @intCast(out.kv_ld) else @as(i64, @intCast(hs.hk)) * dim; + const khs: i64 = if (out.kv_hs != 0) @intCast(out.kv_hs) else dim; + a.add(Dst{ .q = out.q, .q_ld = qld, .q_hs = qhs, .k = out.k, .k_ld = kld, .k_hs = khs, .v = out.v, .v_ld = kld, .v_hs = khs }); + a.add(out.kc orelse 0); + a.add(out.vc orelse 0); + a.add(@as(i64, @intCast(out.len))); + a.add(@as(i32, @intCast(hs.h))); + a.add(@as(i32, @intCast(hs.hk))); + a.add(shape.eps); + const heads = hs.h + (if (src.v != 0) 2 * hs.hk else hs.hk); + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast(heads), .z = 1 }; + try launch_.launch(f orelse g.l2_fn, .{ .grid = grid, .block = .{ .x = @intCast(hs.dim) } }, s, &a); + } +}; \ No newline at end of file diff --git a/zig/tests/cuda/contract_dory.zig b/zig/tests/cuda/contract_dory.zig index 43c282e223..045c65f2be 100644 --- a/zig/tests/cuda/contract_dory.zig +++ b/zig/tests/cuda/contract_dory.zig @@ -4,24 +4,37 @@ const cuda = @import("cuda"); const core = @import("core"); const check = @import("check.zig"); const l2_k = @import("contract_l2.zig"); +const qk_k = @import("contract_qk_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, + qk_rows: ?cuda.qk_rows.QkRows = null, + qk: [qk_k.forms.len]qk_k.Qk = undefined, + qks: 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 { + /// Every Dory registration onto `kernels`; the controls' symbols from the contract_controls image `ctl_mod`. + pub fn add(self: *Dory, gpu: check.Gpu, s: cuda.Stream, norms: *const cuda.norms.Norms, ctl_mod: 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 }; + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .norms = norms, .ctl = if (l2_k.control(w) != null) try ctl_mod.function(l2_k.symbol(w)) else null }; self.l2s += 1; } for (&self.l2) |*t| try kernels.append(a, t.kernel()); + self.qk_rows = try cuda.qk_rows.QkRows.load(gpu.d); + for (qk_k.forms, &self.qk) |f, *t| { + const ctl = if (f.keep) |keep| try ctl_mod.function(if (keep == 2) "tf_ctl_qk_l2_drop_rope" else "tf_ctl_qk_l2_drop_norm") else null; + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .form = f, .k = &self.qk_rows.?, .ctl = ctl, .arena = .init(a) }; + self.qks += 1; + } + for (&self.qk) |*t| try kernels.append(a, t.kernel()); } pub fn deinit(self: *Dory) void { for (self.l2[0..self.l2s]) |*t| t.deinit(); + for (self.qk[0..self.qks]) |*t| t.deinit(); + if (self.qk_rows) |*q| q.unload(); } }; diff --git a/zig/tests/cuda/contract_qk_l2.zig b/zig/tests/cuda/contract_qk_l2.zig new file mode 100644 index 0000000000..4389d5aaf9 --- /dev/null +++ b/zig/tests/cuda/contract_qk_l2.zig @@ -0,0 +1,236 @@ +//! qk_rows' tf_qk_rows_l2 (Dory: RoPE, then each head's L2 norm, q by a scale vector, v as is) and its 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 bfr = core.bf16_round; +const ref = core.l2_ref; +const qk = cuda.qk_rows; + +/// A registration; a control keeps the rounding points in `keep` (qk_l2.cuh's ROPE 1, NORM 2). +pub const Form = struct { name: []const u8, keep: ?i32 = null }; + +pub const forms = [_]Form{ + .{ .name = "cuda qk_rows tf_qk_rows_l2 (Dory: dim 256, rotation from inv_freq, L2 heads, q by (sqk f)^2, v as is)" }, + .{ .name = "control: tf_qk_rows_l2 without the rotated head's rounding", .keep = 2 }, + .{ .name = "control: tf_qk_rows_l2 without q's normalized rounding before its scale", .keep = 1 }, +}; + +const dim = 256; +const group = 4; +/// cosf and sinf within 2 ulps of their value either way. +const trig = 0x1p-22; + +/// The sum of squares' chain: a warp's xor tree (5), then the head's warps in order. +fn depth(_: c.Shape) usize { + return 5 + dim / 32; +} + +/// 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(5 + dim / 32, 0x1p-24) / 2 + 4 * 0x1p-24; +} + +const why = "an output follows its head's sum of squares as ss^-1/2 through bf16 rounding points: a bf16 sum's scale cancels in the norm, and an fp16 sum overflows to a zero output"; + +/// Shape: rows x query heads (n) x head dim (k); Dory's 16 query heads over 4 key heads. +pub fn spec(f: Form) c.Spec { + return .{ + .name = f.name, + .accum = .f32, + .depth = depth, + .ops = 12, + .mids = 4, + .out = .bf16, + .tiles = .{ .{ .size = 16 }, .{}, .{} }, + .served = &.{ .{ .rows = 1, .n = 16, .k = dim }, .{ .rows = 300, .n = 16, .k = dim } }, + .sums = .{ .f16 = 40, .bf16 = 0, .why = why }, + }; +} + +const len = 600; + +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]); + } +}; + +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 (r[0] < v or r[0] == v) and (v < r[1] or v == r[1]); +} + +pub const Qk = struct { + rig: rig_.Rig, + form: Form, + k: *const qk.QkRows, + ctl: ?cuda.Function = null, + arena: std.heap.ArenaAllocator, + r: usize = 0, + h: usize = 0, + hk: usize = 0, + x: []f64 = &.{}, // [r, (h + 2 hk) dim] + qs: []f64 = &.{}, + cs: []f64 = &.{}, // [r, dim]: cos and sin of the fp32 angle pos * inv + sn: []f64 = &.{}, + dx: u64 = 0, + sh: qk.Shape = .{}, + dpos: u64 = 0, + dslot: u64 = 0, + out: qk.Out = undefined, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Qk) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Qk) g.Kernel { + return .{ .spec = spec(self.form), .ctx = self, .vt = &vt, .control = if (self.form.keep != null) .shapes else null }; + } + + const vt = g.hooks(Qk); + + fn heads(self: *const Qk) usize { + return self.h + 2 * self.hk; + } + + pub fn prepare(self: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + if (shape.rows == 0 or shape.k != dim or shape.n % group != 0 or shape.n == 0) return false; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.r = shape.rows; + self.h = shape.n; + self.hk = shape.n / group; + const ld = self.heads() * dim; + const xb = try a.alloc(u16, self.r * ld); + for (xb, 0..) |*h, i| h.* = bfr.bits(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const qsf = try a.alloc(f32, self.h * dim); + for (qsf) |*v| v.* = 0.5 + 30 * rng.float(f32); // (sqk f)^2 + const inv = try a.alloc(f32, dim / 2); // vLLM's fp32 inv_freq at theta 1e4 + for (inv, 0..) |*v, j| v.* = 1 / std.math.pow(f32, 1e4, @as(f32, @floatFromInt(2 * j)) / dim); + const pos = try a.alloc(i32, self.r); + const slot = try a.alloc(i32, self.r); + for (pos, slot, 0..) |*p, *sl, i| { // distinct (slot, pos % len) a row; a third of them past 199,800 + p.* = @intCast((i / 2) % len + len * (i % 5) + (if (i % 3 == 0) @as(usize, 199_800) else 0)); + sl.* = @intCast(i % 2); + } + if (fill == .nonfinite) { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + const at = [_][2]usize{ .{ 0, ld / 3 }, .{ 1, ld / 2 }, .{ 2, ld - 1 }, .{ 3, 0 }, .{ 3, dim - 1 } }; + for (at, put) |p, v| if (p[0] < self.r) { + xb[p[0] * ld + p[1]] = bfr.bits(v); + }; + } + if (fill == .stress) { // integers 32 to 64 of either sign: squares sum past fp16's range, exactly in fp32; no rotation + for (xb) |*h| h.* = bfr.bits(@as(f32, @floatFromInt(32 + rng.uintLessThan(u32, 33))) * (if (rng.boolean()) @as(f32, 1) else -1)); + @memset(inv, 0); + } + self.x = try a.alloc(f64, xb.len); + for (self.x, xb) |*v, h| v.* = bfr.value(h); + self.qs = try a.alloc(f64, qsf.len); + for (self.qs, qsf) |*v, q| v.* = q; + self.cs = try a.alloc(f64, self.r * dim); + self.sn = try a.alloc(f64, self.r * dim); + for (0..self.r) |row| for (0..dim) |i| { + const ang: f64 = @as(f32, @floatFromInt(pos[row])) * inv[i % (dim / 2)]; // the kernel's fp32 angle + self.cs[row * dim + i] = @cos(ang); + self.sn[row * dim + i] = @sin(ang); + }; + self.dx = try self.rig.input(std.mem.sliceAsBytes(xb)); + self.sh = .{ .l2 = true, .q_scale = try self.rig.input(std.mem.sliceAsBytes(qsf)), .inv = try self.rig.input(std.mem.sliceAsBytes(inv)), .rot = dim }; + self.dpos = try self.rig.input(std.mem.sliceAsBytes(pos)); + self.dslot = try self.rig.input(std.mem.sliceAsBytes(slot)); + const kv = self.r * self.hk * dim * 2; + const cache = 2 * len * self.hk * dim * 2; + self.out = .{ .q = try self.rig.output(self.r * self.h * dim * 2), .k = try self.rig.output(kv), .v = try self.rig.output(kv), .kc = try self.rig.scratch(cache), .vc = try self.rig.scratch(cache), .len = len }; + return true; + } + + pub fn run(self: *Qk) !void { + const ld: i64 = @intCast(self.heads() * dim); + const kp = self.dx + self.h * dim * 2; + const src: qk.Src = .{ .q = self.dx, .q_ld = ld, .q_hs = dim, .k = kp, .k_ld = ld, .k_hs = dim, .v = kp + self.hk * dim * 2, .v_ld = ld, .v_hs = dim }; + try self.k.rowsVia(self.ctl, 0, self.rig.s, src, .{ .h = self.h, .hk = self.hk, .dim = dim }, self.sh, self.dpos, self.dslot, self.out, self.r); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Qk) usize { + return self.r * self.heads() * dim; + } + + /// The output index of head hd's first value: q [r, h, dim], then k and v [r, hk, dim]. + fn outAt(self: *const Qk, row: usize, hd: usize) usize { + if (hd < self.h) return (row * self.h + hd) * dim; + const base = self.r * self.h * dim; + if (hd < self.h + self.hk) return base + (row * self.hk + hd - self.h) * dim; + return base + self.r * self.hk * dim + (row * self.hk + hd - self.h - self.hk) * dim; + } + + /// Every head's outputs as ranges under `st` (trig: cosf's and sinf's stray), sums of squares by `sums`. + fn ranges(self: *Qk, sums: anytype, st: ref.Stray, tr: f64, out: []ref.Iv) !void { + const a = self.arena.allocator(); + const small = try a.alloc(f64, dim); + const big = try a.alloc(f64, dim); + const pts: ref.QkPoints = .{}; // the served kernel's points: a control's dropped one leaves the ranges + for (0..self.r) |row| for (0..self.heads()) |hd| { + const x = self.x[(row * self.heads() + hd) * dim ..][0..dim]; + const o = out[self.outAt(row, hd)..][0..dim]; + if (hd >= self.h + self.hk) { + for (o, x) |*e, v| e.* = ref.point(v); + continue; + } + const qs: ?[]const f64 = if (hd < self.h) self.qs[hd * dim ..][0..dim] else null; + ref.qkHead(x, self.cs[row * dim ..][0..dim], self.sn[row * dim ..][0..dim], dim, qs, pts, 1e-6, sums, st, tr, o, small, big); + }; + } + + fn points(self: *Qk, sums: anytype, want: []f64) !void { + const iv = try self.arena.allocator().alloc(ref.Iv, want.len); + try self.ranges(sums, .{}, 0, iv); + for (want, iv) |*w, v| w.* = v[0]; + } + + pub fn reference(self: *Qk, want: []f64, mag: []f64) !void { + try self.points(ref.Exact{}, want); + for (mag, want) |*m, v| m.* = @abs(v) / 2; + } + + pub fn simulate(self: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const buf = try self.arena.allocator().alloc(f64, dim); + try self.points(Sim{ .acc = acc, .order = order, .buf = buf }, out); + } + + /// Every output in float64's range over the fp32 run (a dropped point leaves it); v as stored. + pub fn check(self: *Qk, _: g.Gate) !?[]const u8 { + const n = self.count(); + const a = self.arena.allocator(); + const got = try a.alloc(u16, n); + @memcpy(std.mem.sliceAsBytes(got), try self.rig.outputBytes()); + const iv = try a.alloc(ref.Iv, n); + try self.ranges(ref.Exact{}, .{ .spread = spread(), .u = 0x1p-24 }, trig, iv); + for (got, iv, 0..) |h, r, i| if (!inRange(h, r)) + return try std.fmt.bufPrint(&self.msg, "out[{d}] = {e} outside float64's [{e}, {e}]", .{ i, bfr.value(h), r[0], r[1] }); + return null; + } +}; + +test "the inverse's spread stays far below a bf16 ulp, and each control drops exactly one point" { + try std.testing.expect(spread() < 0x1p-18); + for (forms) |f| if (f.keep) |k| try std.testing.expectEqual(@as(u32, 1), @popCount(3 & ~@as(u32, @intCast(k)))); +} From 50b11cf704fb26ea0ecb08e5ed79db91a3c6e1b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:54:24 +0200 Subject: [PATCH 4/9] cuda act_rows: Dory's scaled output gate and scaled SwiGLU, each row then L2-normed Two activation rows for Dory, a block a row, 16-byte loads, fp32 scale vectors formed on the host as torch forms them: - tf_gate_l2_rows: out = bf16(l2(bf16(o * sigmoid(g * sg)))) over all heads of a row (dory.py's output gate, then justnorm over num_heads * head_dim); - tf_swiglu_l2_rows: out = bf16(l2(bf16(silu(bf16(g * sug)) * bf16(u * suu)))) (its MLP). sigmoid and silu are torch's fp32 forms, 1 / (1 + expf(-x)) and x / (1 + expf(-x)) with IEEE division. ActRows.gateL2 and ActRows.swigluL2 launch them; no existing launch changes. Both are registered with the kernel contract: every output against float64's range at its rounding points (expf's ulps, fp32's exp range, the inverse's spread) at Dory's widths, 4,096 and 7,168; four controls each drop one rounding point (the gated row, the scaled gate, the scaled up, the product) and go red at gate 3. Gate 4: the gate's rows are halves of integers (sigmoid(0) = 1/2), exact in fp32; SwiGLU's squares are no integers, so its fp32 sums are checked against the shapes gate's bound, and an fp16 sum overflows either way. --- docs/kernel-contract.md | 2 + zig/build/cuda_row_contracts.zig | 2 +- zig/kernels/cuda/act_l2.cuh | 135 ++++++++++ zig/kernels/cuda/act_rows.cu | 18 ++ zig/kernels/cuda/contract_controls.cu | 20 ++ zig/src/core/contract/dory_cuda_records.zig | 2 + zig/src/core/l2_ref.zig | 62 +++++ zig/src/cuda/act_rows.zig | 36 ++- zig/tests/cuda/contract_act_l2.zig | 266 ++++++++++++++++++++ zig/tests/cuda/contract_dory.zig | 12 + 10 files changed, 553 insertions(+), 2 deletions(-) create mode 100644 zig/kernels/cuda/act_l2.cuh create mode 100644 zig/tests/cuda/contract_act_l2.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 6317957b2c..38ee1bc48e 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -153,6 +153,8 @@ CUDA also checks Dory's nGPT rows: the L2 norm rows with and without their bf16 (`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). 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 b17a80b606..065b3ecea1 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_l2.zig", "contract_qk_l2.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, diff --git a/zig/kernels/cuda/act_l2.cuh b/zig/kernels/cuda/act_l2.cuh new file mode 100644 index 0000000000..5b347c12f8 --- /dev/null +++ b/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 +#include + +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 +__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; +} + +// 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(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(s)[2 * c] : make_float4(0.f, 0.f, 0.f, 0.f); + const float4 b = c < chunks ? reinterpret_cast(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(out + at)[c] = pack8(y); + } +} + +// out = bf16(l2(t)), t = point(o * sigmoid(g * sg)) over the row's n values (all heads): a block a row. +template +__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(__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(silu(gs) us), gs = point(g sug), us = point(u suu): dory.py's MLP. +template +__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(__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(__fmul_rn(silu(gv[j][i]), point(__fmul_rn(uv[j][i], sv[j][i])))); + normStore(out, r * n, chunks, gv, eps); +} + +} // namespace tf_act_l2 diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index 3563f1f021..0eb5e75147 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/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(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(g, g_ld, u, u_ld, sug, suu, out, n, eps); +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 6c40bfcdb3..9c43d2281a 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -267,3 +267,23 @@ TF_CTL_SLERP(tf_ctl_slerp_drop_b, tf_l2::A) 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(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 diff --git a/zig/src/core/contract/dory_cuda_records.zig b/zig/src/core/contract/dory_cuda_records.zig index dbeac45d2d..d81e34e4e1 100644 --- a/zig/src/core/contract/dory_cuda_records.zig +++ b/zig/src/core/contract/dory_cuda_records.zig @@ -5,4 +5,6 @@ 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 } } }, .{ .launch = .{ .name = "qk_rows tf_qk_rows_l2 at dim 256", .symbol = "tf_qk_rows_l2", .image = "qk_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 } } }, + .{ .launch = .{ .name = "act_rows tf_gate_l2_rows", .symbol = "tf_gate_l2_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 107, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_l2_rows", .symbol = "tf_swiglu_l2_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 118, .static = 32, .bound = 256 }, .{ .regs = 118, .static = 32, .bound = 256 }, .{ .regs = 121, .static = 32, .bound = 256 }, .{ .regs = 121, .static = 32, .bound = 256 } } }, }; diff --git a/zig/src/core/l2_ref.zig b/zig/src/core/l2_ref.zig index 33be0254df..3e267e3bbb 100644 --- a/zig/src/core/l2_ref.zig +++ b/zig/src/core/l2_ref.zig @@ -181,3 +181,65 @@ pub fn qkHead(x: []const f64, cos: []const f64, sin: []const f64, rot: usize, qs } } +/// exp in fp32's range: past its largest finite value, infinity (expf overflows there). +fn exp32(x: f64) f64 { + const e = @exp(x); + return if (e > std.math.floatMax(f32)) std.math.inf(f64) else e; +} + +/// exp's ends over x's interval, each within `e` relative (expf's ulps); monotone. +fn expIv(x: Iv, e: f64) Iv { + return .{ exp32(x[0]) * (1 - e), exp32(x[1]) * (1 + e) }; +} + +/// 1 / (1 + exp(-x)) over x's interval in fp32 steps (expf within `e`, the sum and quotient by u): increasing in x. +pub fn sigmoid(x: Iv, e: f64, u: f64) Iv { + const ex = expIv(.{ -x[1], -x[0] }, e); + return .{ 1 / ((1 + ex[1]) * (1 + u)) * (1 - u), 1 / ((1 + ex[0]) * (1 - u)) * (1 + u) }; +} + +/// x / (1 + exp(-x)) at one value (a point x), its fp32 steps within e and u. +pub fn silu(x: f64, e: f64, u: f64) Iv { + const ex = exp32(-x); + const lo = x / ((1 + ex * (1 + e)) * (1 + u)); + const hi = x / ((1 + ex * (1 - e)) * (1 - u)); + return mul(order(lo, hi), point(1), u); +} + +/// tf_gate_l2_rows' rounding point: the gated row before its norm. +pub const GatePoints = struct { gated: bool = true }; + +/// tf_gate_l2_rows: out = bf16(l2(t)), t = o sigmoid(g sg) (bf16 at `gated`), over the row's values; expf within e. +pub fn gateL2(o: []const f64, g: []const f64, sg: []const f64, p: GatePoints, eps: f64, sums: anytype, st: Stray, e: f64, y: []Iv, small: []f64, big: []f64) void { + for (y, o, g, sg) |*out, ov, gv, sv| { + const t = mul(point(gv), point(sv), st.u); + if (std.math.isNan(t[0])) { + out.* = @splat(std.math.nan(f64)); + continue; + } + out.* = bf16If(p.gated, mul(point(ov), sigmoid(t, e, st.u), st.u)); + } + normalize(y, eps, sums, st, true, small, big); +} + +/// tf_swiglu_l2_rows' rounding points: the scaled gate and up, their product. +pub const SwigluPoints = struct { gs: bool = true, us: bool = true, prod: bool = true }; + +/// tf_swiglu_l2_rows: out = bf16(l2(m)), m = silu(gs) us, gs = g sug, us = u suu (bf16 at the points); expf within e. +pub fn swigluL2(g: []const f64, u: []const f64, sug: []const f64, suu: []const f64, p: SwigluPoints, eps: f64, sums: anytype, st: Stray, e: f64, y: []Iv, small: []f64, big: []f64) void { + for (y, g, u, sug, suu) |*out, gv, uv, sgv, suv| { + const gs = bf16If(p.gs, mul(point(gv), point(sgv), st.u)); + const us = bf16If(p.us, mul(point(uv), point(suv), st.u)); + if (std.math.isNan(gs[0]) or std.math.isNan(us[0])) { + out.* = @splat(std.math.nan(f64)); + continue; + } + // silu is not monotone below -1.28: its ends over gs's interval from both ends and the minimum between + var s = silu(gs[0], e, st.u); + const s2 = silu(gs[1], e, st.u); + s = .{ @min(s[0], s2[0]), @max(s[1], s2[1]) }; + if (gs[0] < -1.2784645 and gs[1] > -1.2784645) s[0] = @min(s[0], -0.2784645 * (1 + 4 * st.u)); + out.* = bf16If(p.prod, mul(s, us, st.u)); + } + normalize(y, eps, sums, st, true, small, big); +} diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig index e23223e7d5..ee448c6a16 100644 --- a/zig/src/cuda/act_rows.zig +++ b/zig/src/cuda/act_rows.zig @@ -31,12 +31,14 @@ pub const ActRows = struct { geglu_fn: Function, swiglu_round_fn: Function, softcap_fn: Function, + gate_l2_fn: Function, + swiglu_l2_fn: Function, pub fn load(d: *const Driver) !ActRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.act_rows); errdefer mod.unload(); - return .{ .mod = mod, .swiglu_fn = try mod.function("tf_swiglu_rows"), .gate_fn = try mod.function("tf_gate_rows"), .sums_fn = try mod.function("tf_group_sums_rows"), .dconv_fn = try mod.function("tf_dconv_rows"), .geglu_fn = try mod.function("tf_geglu_rows"), .swiglu_round_fn = try mod.function("tf_swiglu_rows_round"), .softcap_fn = try mod.function("tf_softcap_rows") }; + return .{ .mod = mod, .swiglu_fn = try mod.function("tf_swiglu_rows"), .gate_fn = try mod.function("tf_gate_rows"), .sums_fn = try mod.function("tf_group_sums_rows"), .dconv_fn = try mod.function("tf_dconv_rows"), .geglu_fn = try mod.function("tf_geglu_rows"), .swiglu_round_fn = try mod.function("tf_swiglu_rows_round"), .softcap_fn = try mod.function("tf_softcap_rows"), .gate_l2_fn = try mod.function("tf_gate_l2_rows"), .swiglu_l2_fn = try mod.function("tf_swiglu_l2_rows") }; } pub fn unload(g: *ActRows) void { @@ -103,6 +105,38 @@ pub const ActRows = struct { try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d, t), .block = .{ .x = t } }, s, &a); } + /// Dory's gate: out [rows, n] = bf16(l2(bf16(o * sigmoid(g * sg)))), o packed, g rows g_ld apart, sg fp32 [n]. + pub fn gateL2(g: *const ActRows, s: Stream, o: u64, gate_: u64, g_ld: usize, sg: u64, out: u64, rows: usize, n: usize, eps: f32) !void { + if (n % 64 != 0 or n > 8192 or !aligned(&.{ o, gate_, out, sg }, &.{g_ld})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(o); + a.add(gate_); + a.add(@as(i64, @intCast(g_ld))); + a.add(sg); + a.add(out); + a.add(@as(i32, @intCast(n))); + a.add(eps); + try launch_.launch(g.gate_l2_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// Dory's MLP: out = bf16(l2(bf16(silu(bf16(g * sug)) * bf16(u * suu)))), sug and suu fp32 [n]. + pub fn swigluL2(g: *const ActRows, s: Stream, gate_: u64, g_ld: usize, up: u64, u_ld: usize, sug: u64, suu: u64, out: u64, rows: usize, n: usize, eps: f32) !void { + if (n % 64 != 0 or n > 8192 or !aligned(&.{ gate_, up, out, sug, suu }, &.{ g_ld, u_ld })) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(gate_); + a.add(@as(i64, @intCast(g_ld))); + a.add(up); + a.add(@as(i64, @intCast(u_ld))); + a.add(sug); + a.add(suu); + a.add(out); + a.add(@as(i32, @intCast(n))); + a.add(eps); + try launch_.launch(g.swiglu_l2_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + /// xs [rows, k / 64]: x's rows (stride ld) summed over each 64 values. k % 64 == 0. pub fn groupSums(g: *const ActRows, s: Stream, x: u64, ld: usize, xs: u64, rows: usize, k: usize) !void { if (k % 64 != 0 or !aligned(&.{x}, &.{ld})) return error.RowShape; diff --git a/zig/tests/cuda/contract_act_l2.zig b/zig/tests/cuda/contract_act_l2.zig new file mode 100644 index 0000000000..a5f5ab45ea --- /dev/null +++ b/zig/tests/cuda/contract_act_l2.zig @@ -0,0 +1,266 @@ +//! act_rows' nGPT activations under the contract (tf_gate_l2_rows, tf_swiglu_l2_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 { gate, swiglu, drop_gated, drop_gs, drop_us, drop_prod }; + +const chain = 42; // 32 fma a thread over 16-byte loads, then 256 threads' tree + +fn depth(_: c.Shape) usize { + return chain; +} + +pub fn spread() f64 { + return c.gamma(chain, 0x1p-24) / 2 + 8 * 0x1p-24; +} + +/// expf within 2 ulps of its value either way. +const exp_stray = 0x1p-22; + +fn isGate(which: Which) bool { + return which == .gate or which == .drop_gated; +} + +const why = "an output follows its row's sum of squares as ss^-1/2 through bf16 rounding points, so a bf16 sum's scale cancels in the norm, and an fp16 sum overflows to a zero output"; + +/// Dory's attention rows (16 heads of 256) and MLP rows (7,168), a decode window and a short verify. +const gate_served = [_]c.Shape{ .{ .rows = 16, .n = 4096, .k = 4096 }, .{ .rows = 3, .n = 4096, .k = 4096 } }; +const mlp_served = [_]c.Shape{ .{ .rows = 16, .n = 7168, .k = 7168 }, .{ .rows = 3, .n = 7168, .k = 7168 } }; + +pub fn spec(which: Which) c.Spec { + return .{ + .name = switch (which) { + .gate => "cuda act_rows tf_gate_l2_rows (Dory's output gate: bf16(l2(bf16(o sigmoid(g sg)))) over all heads)", + .swiglu => "cuda act_rows tf_swiglu_l2_rows (Dory's MLP: bf16(l2(bf16(silu(bf16(g sug)) bf16(u suu)))))", + .drop_gated => "control: tf_gate_l2_rows without the gated row's rounding", + .drop_gs => "control: tf_swiglu_l2_rows without the scaled gate's rounding", + .drop_us => "control: tf_swiglu_l2_rows without the scaled up's rounding", + .drop_prod => "control: tf_swiglu_l2_rows without the product's rounding", + }, + .accum = .f32, + .depth = depth, + .ops = 16, + .mids = if (isGate(which)) 2 else 4, // a bf16 point decided the other way by fp32 steps: two half-ulps a point + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = if (isGate(which)) &gate_served else &mlp_served, + // SwiGLU's squares are no integers: its fp32 sums carry their own gamma, the shapes gate's bound + .sums = .{ .f16 = 40, .bf16 = 0, .why = why, .exact = isGate(which) }, + }; +} + +pub fn control(which: Which) ?g.Gate { + return switch (which) { + .gate, .swiglu => null, + else => .shapes, + }; +} + +pub fn symbol(which: Which) [:0]const u8 { + return switch (which) { + .drop_gated => "tf_ctl_gate_l2_drop_gated", + .drop_gs => "tf_ctl_swiglu_l2_drop_gs", + .drop_us => "tf_ctl_swiglu_l2_drop_us", + .drop_prod => "tf_ctl_swiglu_l2_drop_prod", + else => unreachable, + }; +} + +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]); + } +}; + +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 (r[0] < v or r[0] == v) and (v < r[1] or v == r[1]); +} + +pub const Act = struct { + rig: rig_.Rig, + which: Which, + k: *const cuda.act_rows.ActRows, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + n: usize = 0, + width: usize = 0, + a: []f64 = &.{}, // o (gate) or g (swiglu) + b: []f64 = &.{}, // g (gate) or u (swiglu) + sa: []f64 = &.{}, // sg (gate) or sug + sb: []f64 = &.{}, // suu + dev: [5]u64 = @splat(0), + out: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Act) void { + if (self.arena) |*x| x.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Act); + + pub fn kernel(self: *Act) 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, 8) else 1))); + } + + fn integers(rng: std.Random, v: []f64) void { + for (v) |*x| x.* = @as(f64, @floatFromInt(rng.uintLessThan(u32, 65))) - 32; + } + + fn bf16In(self: *Act, v: []const f64) !u64 { + const b = try self.arena.?.allocator().alloc(u16, v.len); + for (b, v) |*h, x| h.* = bfr.bits(@floatCast(x)); + return self.rig.input(std.mem.sliceAsBytes(b)); + } + + fn f32In(self: *Act, v: []const f64) !u64 { + const b = try self.arena.?.allocator().alloc(f32, v.len); + for (b, v) |*h, x| h.* = @floatCast(x); + return self.rig.input(std.mem.sliceAsBytes(b)); + } + + pub fn prepare(self: *Act, 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) |*x| x.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const al = self.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.width = width; + self.n = shape.rows * width; + self.a = try al.alloc(f64, self.n); + self.b = try al.alloc(f64, self.n); + self.sa = try al.alloc(f64, width); + self.sb = try al.alloc(f64, width); + bfs(rng, self.a, 2); + bfs(rng, self.b, 2); + // the scale vectors as fp32: Dory's (s_gate f)^2 and suv sqrt(hidden) ranges + for (self.sa, self.sb) |*x, *y| { + x.* = @as(f32, if (isGate(self.which)) 1.5 + 30 * rng.float(f32) else 0.3 + 3 * rng.float(f32)); + y.* = @as(f32, 0.3 + 3 * rng.float(f32)); + } + 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.a else self.b)[p[0] * width + p[2]] = p[3]; + }; + } + if (fill == .stress) { + integers(rng, self.a); + if (isGate(self.which)) @memset(self.b, 0) else integers(rng, self.b); // sigmoid(0) = 1/2: o / 2 exactly + @memset(self.sa, 1); + @memset(self.sb, 1); + } + self.dev = .{ try self.bf16In(self.a), try self.bf16In(self.b), try self.f32In(self.sa), try self.f32In(self.sb), 0 }; + self.out = try self.rig.output(self.n * 2); + return true; + } + + pub fn run(self: *Act) !void { + const s = self.rig.s; + const rws = self.n / self.width; + const w = self.width; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + args.add(self.dev[0]); + if (!isGate(self.which)) args.add(@as(i64, @intCast(w))); + args.add(self.dev[1]); + args.add(@as(i64, @intCast(w))); + args.add(self.dev[2]); + if (!isGate(self.which)) args.add(self.dev[3]); + args.add(self.out); + args.add(@as(c_int, @intCast(w))); + 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 (isGate(self.which)) { + try self.k.gateL2(s, self.dev[0], self.dev[1], w, self.dev[2], self.out, rws, w, 1e-6); + } else try self.k.swigluL2(s, self.dev[0], w, self.dev[1], w, self.dev[2], self.dev[3], self.out, rws, w, 1e-6); + try s.synchronize(); + } + + pub fn count(self: *Act) usize { + return self.n; + } + + fn ranges(self: *Act, al: std.mem.Allocator, sums: anytype, st: ref.Stray, e: f64, out: []ref.Iv) !void { + const w = self.width; + const small = try al.alloc(f64, w); + const big = try al.alloc(f64, w); + for (0..self.n / w) |row| { + const at = row * w; + const o = out[at..][0..w]; + switch (self.which) { + .gate => ref.gateL2(self.a[at..][0..w], self.b[at..][0..w], self.sa, .{}, 1e-6, sums, st, e, o, small, big), + else => ref.swigluL2(self.a[at..][0..w], self.b[at..][0..w], self.sa, self.sb, .{}, 1e-6, sums, st, e, o, small, big), + } + } + } + + fn points(self: *Act, sums: anytype, out: []f64) !void { + var arena: std.heap.ArenaAllocator = .init(self.rig.gpu.gpa); + defer arena.deinit(); + const iv = try arena.allocator().alloc(ref.Iv, self.n); + try self.ranges(arena.allocator(), sums, .{}, 0, iv); + for (out, iv) |*o, v| o.* = v[0]; + } + + pub fn reference(self: *Act, want: []f64, mag: []f64) !void { + try self.points(ref.Exact{}, want); + for (mag, want) |*m, v| m.* = @abs(v) / 2; + } + + pub fn simulate(self: *Act, 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); + } + + pub fn check(self: *Act, _: g.Gate) !?[]const u8 { + const al = self.arena.?.allocator(); + const got = try al.alloc(u16, self.n); + @memcpy(std.mem.sliceAsBytes(got), try self.rig.bytes(self.out, self.n * 2)); + const iv = try al.alloc(ref.Iv, self.n); + try self.ranges(al, ref.Exact{}, .{ .spread = spread(), .u = 0x1p-24 }, exp_stray, 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 gate rows: halves of integers square exactly in fp32, and fp16 overflows" { + var terms: [4096]f64 = undefined; + var prng = std.Random.DefaultPrng.init(0x6761); + for (&terms) |*t| { + const v: f64 = (@as(f64, @floatFromInt(prng.random().uintLessThan(u32, 65))) - 32) / 2; + t.* = v * v; + } + for (std.enums.values(rows.Order)) |o| { + var exact: f64 = 0; + for (terms) |t| exact += t; + try std.testing.expectEqual(exact, rows.sumIn(.f32, o, &terms)); + try std.testing.expect(std.math.isInf(rows.sumIn(.f16, o, &terms))); + } +} diff --git a/zig/tests/cuda/contract_dory.zig b/zig/tests/cuda/contract_dory.zig index 045c65f2be..c9acd23b15 100644 --- a/zig/tests/cuda/contract_dory.zig +++ b/zig/tests/cuda/contract_dory.zig @@ -5,6 +5,7 @@ const core = @import("core"); const check = @import("check.zig"); const l2_k = @import("contract_l2.zig"); const qk_k = @import("contract_qk_l2.zig"); +const act_k = @import("contract_act_l2.zig"); const g = core.contract.gates; @@ -14,6 +15,9 @@ pub const Dory = struct { qk_rows: ?cuda.qk_rows.QkRows = null, qk: [qk_k.forms.len]qk_k.Qk = undefined, qks: usize = 0, + act_rows: ?cuda.act_rows.ActRows = null, + act: [std.enums.values(act_k.Which).len]act_k.Act = undefined, + acts: usize = 0, /// Every Dory registration onto `kernels`; the controls' symbols from the contract_controls image `ctl_mod`. pub fn add(self: *Dory, gpu: check.Gpu, s: cuda.Stream, norms: *const cuda.norms.Norms, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { @@ -30,11 +34,19 @@ pub const Dory = struct { self.qks += 1; } for (&self.qk) |*t| try kernels.append(a, t.kernel()); + self.act_rows = try cuda.act_rows.ActRows.load(gpu.d); + for (std.enums.values(act_k.Which), &self.act) |w, *t| { + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .k = &self.act_rows.?, .ctl = if (act_k.control(w) != null) try ctl_mod.function(act_k.symbol(w)) else null }; + self.acts += 1; + } + for (&self.act) |*t| try kernels.append(a, t.kernel()); } pub fn deinit(self: *Dory) void { for (self.l2[0..self.l2s]) |*t| t.deinit(); for (self.qk[0..self.qks]) |*t| t.deinit(); if (self.qk_rows) |*q| q.unload(); + for (self.act[0..self.acts]) |*t| t.deinit(); + if (self.act_rows) |*q| q.unload(); } }; From 0346bdc405b39ba81130ea1479ec790608d8c06e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:55:16 +0200 Subject: [PATCH 5/9] cuda dense_rows: a per-column scaled epilogue for Dory's logits Dory's logits are the head's bf16 output times logit_scale / std per vocabulary column (dory.py: logits.float() * (logit_scale.float() / initializer_range)). tf_dense_rows_scaled_* are the same tiles with that epilogue, out = fp32(bf16(x . w[n])) * scale[n] in one fp32 product; DenseRows.rowsScaled launches them. tf_dense_rows_* keep their code: their records from nvcc 13.0 are unchanged. Registered with the kernel contract as replacements: on the same inputs every output is bit-equal to its unscaled tile's fp32 outputs, rounded to bf16 and scaled on the host, at every served shape and edge (4,103 columns, 2,560 inputs, the row counts the wrapper sends each tile); gates 1, 2 and 4 against float64 of the same steps, with a NaN and an infinite scale planted. --- docs/kernel-contract.md | 2 + zig/kernels/cuda/dense_rows.cu | 18 ++- zig/src/core/contract/dory_cuda_records.zig | 4 + zig/src/cuda/dense_rows.zig | 21 ++- zig/tests/cuda/contract_dense_scaled.zig | 158 ++++++++++++++++++++ zig/tests/cuda/contract_dory.zig | 12 ++ 6 files changed, 210 insertions(+), 5 deletions(-) create mode 100644 zig/tests/cuda/contract_dense_scaled.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 38ee1bc48e..32597e615b 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -155,6 +155,8 @@ CUDA also checks Dory's attention inputs (`qk_l2.cuh`): RoPE first, each head's 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 diff --git a/zig/kernels/cuda/dense_rows.cu b/zig/kernels/cuda/dense_rows.cu index acfc61efe3..bb6a33c351 100644 --- a/zig/kernels/cuda/dense_rows.cu +++ b/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 +template // 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(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) diff --git a/zig/src/core/contract/dory_cuda_records.zig b/zig/src/core/contract/dory_cuda_records.zig index d81e34e4e1..a56d1d8782 100644 --- a/zig/src/core/contract/dory_cuda_records.zig +++ b/zig/src/core/contract/dory_cuda_records.zig @@ -7,4 +7,8 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "qk_rows tf_qk_rows_l2 at dim 256", .symbol = "tf_qk_rows_l2", .image = "qk_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 30, .static = 64, .bound = 512 } } }, .{ .launch = .{ .name = "act_rows tf_gate_l2_rows", .symbol = "tf_gate_l2_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 107, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 109, .static = 32, .bound = 256 } } }, .{ .launch = .{ .name = "act_rows tf_swiglu_l2_rows", .symbol = "tf_swiglu_l2_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 109, .static = 32, .bound = 256 }, .{ .regs = 118, .static = 32, .bound = 256 }, .{ .regs = 118, .static = 32, .bound = 256 }, .{ .regs = 121, .static = 32, .bound = 256 }, .{ .regs = 121, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_scaled_1", .symbol = "tf_dense_rows_scaled_1", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_scaled_4", .symbol = "tf_dense_rows_scaled_4", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 45, .static = 0, .bound = 256 }, .{ .regs = 45, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_scaled_16", .symbol = "tf_dense_rows_scaled_16", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_scaled_32", .symbol = "tf_dense_rows_scaled_32", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 } } }, }; diff --git a/zig/src/cuda/dense_rows.zig b/zig/src/cuda/dense_rows.zig index 9241a4596e..125d932c4a 100644 --- a/zig/src/cuda/dense_rows.zig +++ b/zig/src/cuda/dense_rows.zig @@ -17,14 +17,19 @@ const tiles = [_]usize{ 1, 4, 16, 32 }; pub const DenseRows = struct { mod: Module, fns: [tiles.len]Function, + scaled_fns: [tiles.len]Function, pub fn load(d: *const Driver) !DenseRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.dense_rows); errdefer mod.unload(); var fns: [tiles.len]Function = undefined; - inline for (tiles, 0..) |t, i| fns[i] = try mod.function(std.fmt.comptimePrint("tf_dense_rows_{d}", .{t})); - return .{ .mod = mod, .fns = fns }; + var scaled: [tiles.len]Function = undefined; + inline for (tiles, 0..) |t, i| { + fns[i] = try mod.function(std.fmt.comptimePrint("tf_dense_rows_{d}", .{t})); + scaled[i] = try mod.function(std.fmt.comptimePrint("tf_dense_rows_scaled_{d}", .{t})); + } + return .{ .mod = mod, .fns = fns, .scaled_fns = scaled }; } pub fn unload(g: *DenseRows) void { @@ -33,6 +38,15 @@ pub const DenseRows = struct { /// out [m, n] fp32 = x [m, k] bf16 times w [n, k] bf16 transposed; strides in elements, k % 8 == 0, 16 B aligned. pub fn rows(g: *const DenseRows, s: Stream, x: u64, x_stride: usize, w: u64, out: u64, out_stride: usize, m: usize, n: usize, k: usize) !void { + return g.rowsWith(null, s, x, x_stride, w, out, out_stride, m, n, k); + } + + /// rows with the scaled epilogue: out = fp32(bf16(x . w[n])) * scale[n], scale fp32 [n] (Dory's logits). + pub fn rowsScaled(g: *const DenseRows, s: Stream, x: u64, x_stride: usize, w: u64, scale: u64, out: u64, out_stride: usize, m: usize, n: usize, k: usize) !void { + return g.rowsWith(scale, s, x, x_stride, w, out, out_stride, m, n, k); + } + + fn rowsWith(g: *const DenseRows, scale: ?u64, s: Stream, x: u64, x_stride: usize, w: u64, out: u64, out_stride: usize, m: usize, n: usize, k: usize) !void { if (k % 8 != 0 or x % 16 != 0 or w % 16 != 0 or x_stride % 8 != 0) return error.Misaligned; if (m == 0 or n == 0) return; const i: usize = if (m == 1) 0 else if (m <= 4) 1 else if (m <= 16) 2 else 3; @@ -40,12 +54,13 @@ pub const DenseRows = struct { a.add(x); a.add(@as(i64, @intCast(x_stride))); a.add(w); + if (scale) |v| a.add(v); a.add(out); a.add(@as(i64, @intCast(out_stride))); a.add(@as(i32, @intCast(m))); a.add(@as(i32, @intCast(n))); a.add(@as(i32, @intCast(k))); const grid: launch_.Dim3 = .{ .x = @intCast((n + warps - 1) / warps), .y = @intCast((m + tiles[i] - 1) / tiles[i]), .z = 1 }; - try launch_.launch(g.fns[i], .{ .grid = grid, .block = .{ .x = warps * 32 } }, s, &a); + try launch_.launch(if (scale != null) g.scaled_fns[i] else g.fns[i], .{ .grid = grid, .block = .{ .x = warps * 32 } }, s, &a); } }; diff --git a/zig/tests/cuda/contract_dense_scaled.zig b/zig/tests/cuda/contract_dense_scaled.zig new file mode 100644 index 0000000000..e0c05dd068 --- /dev/null +++ b/zig/tests/cuda/contract_dense_scaled.zig @@ -0,0 +1,158 @@ +//! dense_rows' scaled epilogue (tf_dense_rows_scaled_*) under the contract: bit-equal to its unscaled tile, scaled. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const moe_k = @import("contract_moe.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const bfr = core.bf16_round; + +pub const tiles = moe_k.dense_tiles; + +fn chain(shape: c.Shape) usize { + return moe_k.denseSpec(0).depth(shape); +} + +/// Dory's head (131,072 x 2,560) cut to an odd 4,103 columns, at the row counts the wrapper sends each tile. +pub fn spec(i: usize) c.Spec { + const names = [_][]const u8{ "cuda dense rows tf_dense_rows_scaled_1 (Dory's logits)", "cuda dense rows tf_dense_rows_scaled_4", "cuda dense rows tf_dense_rows_scaled_16", "cuda dense rows tf_dense_rows_scaled_32" }; + const m = [_][2]usize{ .{ 1, 1 }, .{ 4, 3 }, .{ 16, 9 }, .{ 32, 40 } }; + const served = [_][2]c.Shape{ + .{ .{ .rows = m[0][0], .n = 4103, .k = 2560 }, .{ .rows = m[0][1], .n = 384, .k = 2560 } }, + .{ .{ .rows = m[1][0], .n = 4103, .k = 2560 }, .{ .rows = m[1][1], .n = 384, .k = 2560 } }, + .{ .{ .rows = m[2][0], .n = 4103, .k = 2560 }, .{ .rows = m[2][1], .n = 384, .k = 2560 } }, + .{ .{ .rows = m[3][0], .n = 4103, .k = 2560 }, .{ .rows = m[3][1], .n = 384, .k = 2560 } }, + }; + return .{ + .name = names[i], + .accum = .f32, + .depth = chain, + .ops = 2, // the scale's product, after the bf16 rounding (mids) + .mids = 2, + .out = .f32, + .replaces = "its tf_dense_rows tile's fp32 outputs, rounded to bf16 and times the scale in fp32 on the host", + .tiles = .{ .{ .size = tiles[i] }, .{ .size = 8 }, .{ .size = 256, .step = 8, .min = 8 } }, + .served = &served[i], + }; +} + +pub const Scaled = struct { + rig: rig_.Rig, + tile: usize, + g: *const cuda.dense_rows.DenseRows, + d: ?Dense = null, + scale: []f32 = &.{}, + dev: [4]u64 = @splat(0), + base: std.ArrayList(u8) = .empty, + + pub fn deinit(self: *Scaled) void { + if (self.d) |d| d.deinit(); + self.rig.gpu.gpa.free(self.scale); + self.base.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Scaled) g.Kernel { + return .{ .spec = spec(self.tile), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Scaled); + + pub fn prepare(self: *Scaled, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.k % 8 != 0 or shape.k == 0 or shape.n == 0 or shape.rows == 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 1); + const bits = try a.alloc(u16, d.w.len); + defer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed ^ 0x77); + const rng = prng.random(); + const k = shape.k; + for (bits) |*b| b.* = c.toBf16(if (fill == .stress) 0 else rng.float(f32) * 2 - 1); + if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + bits[j * k + i] = c.toBf16(v); + bits[j * k + k / 2 + i] = c.toBf16(-v); + }; + for (d.w, bits) |*w, b| w.* = c.fromBf16(b); + a.free(self.scale); + self.scale = try a.alloc(f32, shape.n); + // logit_scale / std: Dory's runs about 20 to 120; a NaN and an infinity with the non-finite fill + for (self.scale) |*v| v.* = 20 + 100 * rng.float(f32); + if (fill == .nonfinite and shape.n > 7) { + self.scale[3] = std.math.nan(f32); + self.scale[7] = std.math.inf(f32); + } + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(bits)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(self.scale)); + self.dev[3] = try self.rig.output(d.outputs() * 4); + return true; + } + + /// The wrapper's arguments at this entry's row tile, through `f` (scale `with`: the scaled kernel's). + fn launch(self: *Scaled, f: cuda.Function, with: bool, out: u64) !void { + const sh = self.d.?.shape; + var a: cuda.Args = .{}; + a.add(self.dev[0]); + a.add(@as(i64, @intCast(sh.k))); + a.add(self.dev[1]); + if (with) a.add(self.dev[2]); + a.add(out); + a.add(@as(i64, @intCast(sh.n))); + for ([_]usize{ sh.rows, sh.n, sh.k }) |v| a.add(@as(i32, @intCast(v))); + const rt = tiles[self.tile]; + const grid: cuda.Dim3 = .{ .x = @intCast((sh.n + 7) / 8), .y = @intCast((sh.rows + rt - 1) / rt), .z = 1 }; + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn run(self: *Scaled) !void { + try self.launch(self.g.scaled_fns[self.tile], true, self.dev[3]); + } + + pub fn count(self: *Scaled) usize { + return self.d.?.outputs(); + } + + fn scaled(self: *Scaled, out: []f64) void { + const n = self.d.?.shape.n; + for (out, 0..) |*v, i| v.* = @as(f32, @floatCast(bfr.bf16(v.*))) * self.scale[i % n]; + } + + /// float64's sums, rounded to bf16 and scaled in fp32 (the epilogue's steps on an exact sum). + pub fn reference(self: *Scaled, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + const n = self.d.?.shape.n; + for (mag, 0..) |*m, i| m.* *= @abs(self.scale[i % n]); + self.scaled(want); + } + + pub fn simulate(self: *Scaled, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + self.scaled(out); + } + + /// The unscaled tile on the same inputs, its fp32 outputs rounded to bf16 and scaled on the host. + pub fn replaced(self: *Scaled) ![]const u8 { + const a = self.rig.gpu.gpa; + const n = self.count(); + var tmp = try cuda.DeviceBuffer.alloc(self.rig.gpu.d, n * 4); + defer tmp.free(); + try self.launch(self.g.fns[self.tile], false, tmp.ptr); + try self.base.resize(a, n * 4); + try tmp.download(0, self.base.items); + const v = std.mem.bytesAsSlice(f32, @as([]align(1) u8, self.base.items)); + const cols = self.d.?.shape.n; + for (0..n) |i| v[i] = bfr.round32(v[i]) * self.scale[i % cols]; + return self.base.items; + } +}; diff --git a/zig/tests/cuda/contract_dory.zig b/zig/tests/cuda/contract_dory.zig index c9acd23b15..6735c08847 100644 --- a/zig/tests/cuda/contract_dory.zig +++ b/zig/tests/cuda/contract_dory.zig @@ -6,6 +6,7 @@ const check = @import("check.zig"); const l2_k = @import("contract_l2.zig"); const qk_k = @import("contract_qk_l2.zig"); const act_k = @import("contract_act_l2.zig"); +const dense_k = @import("contract_dense_scaled.zig"); const g = core.contract.gates; @@ -18,6 +19,9 @@ pub const Dory = struct { act_rows: ?cuda.act_rows.ActRows = null, act: [std.enums.values(act_k.Which).len]act_k.Act = undefined, acts: usize = 0, + dense_rows: ?cuda.dense_rows.DenseRows = null, + dense: [dense_k.tiles.len]dense_k.Scaled = undefined, + denses: usize = 0, /// Every Dory registration onto `kernels`; the controls' symbols from the contract_controls image `ctl_mod`. pub fn add(self: *Dory, gpu: check.Gpu, s: cuda.Stream, norms: *const cuda.norms.Norms, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { @@ -40,6 +44,12 @@ pub const Dory = struct { self.acts += 1; } for (&self.act) |*t| try kernels.append(a, t.kernel()); + self.dense_rows = try cuda.dense_rows.DenseRows.load(gpu.d); + for (&self.dense, 0..) |*t, i| { + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .tile = i, .g = &self.dense_rows.? }; + self.denses += 1; + } + for (&self.dense) |*t| try kernels.append(a, t.kernel()); } pub fn deinit(self: *Dory) void { @@ -48,5 +58,7 @@ pub const Dory = struct { if (self.qk_rows) |*q| q.unload(); for (self.act[0..self.acts]) |*t| t.deinit(); if (self.act_rows) |*q| q.unload(); + for (self.dense[0..self.denses]) |*t| t.deinit(); + if (self.dense_rows) |*q| q.unload(); } }; From 94577eb5d2736a07e3410fc1801314309855924e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 09:52:55 +0200 Subject: [PATCH 6/9] dory: family skeleton, config, weight map, sublayer schedule and memory (host only) Nemotron-Labs Dory 1.0 is a looped nGPT stack: input sublayers once, the recurrent block n_recurrent_loops times, output sublayers once, attention and MLP sublayers alternating. - config.zig parses config.json and generation_config.json's end tokens: the block counts, sliding/full/mlp kinds, both RoPE thetas, the nGPT constants and their fp32 factors, the loop count (and withLoops, 1 to 4, for a load option). It refuses last_loop KV (chunk- dependent by construction), quantizations other than ModelOpt W4A16_NVFP4 at group 16 or plain bf16, scaled or partial RoPE, tied heads, biases and shapes the kernels do not take. - names.zig maps every tensor of the NVFP4 (or bf16) checkpoint with dtype and shape checks and refuses leftovers; the pinned checkpoint (TF_DORY_DIR) maps its 1,083 tensors. - schedule.zig gives the virtual sublayers for L loops with a per_loop cache per attention step: 8 + 20 L + 8 caches, matching llama.cpp #30306's physical[] and cache_layer[] and vLLM #60149's loop * layers + layer indices. - memory.zig counts weight bytes as stored (equal to the pinned file's tensor data) and KV bytes by loops, context and sliding ring: 36, 56 and 96 KiB a token at 1, 2 and 4 loops. --- zig/build/cuda_family_host.zig | 1 + zig/build/dory.zig | 12 + zig/src/families/dory/config.zig | 477 +++++++++++++++++++++++++++++ zig/src/families/dory/dory.zig | 14 + zig/src/families/dory/memory.zig | 72 +++++ zig/src/families/dory/names.zig | 210 +++++++++++++ zig/src/families/dory/schedule.zig | 154 ++++++++++ 7 files changed, 940 insertions(+) create mode 100644 zig/build/dory.zig create mode 100644 zig/src/families/dory/config.zig create mode 100644 zig/src/families/dory/dory.zig create mode 100644 zig/src/families/dory/memory.zig create mode 100644 zig/src/families/dory/names.zig create mode 100644 zig/src/families/dory/schedule.zig diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index 372b08ff22..489ca052fc 100644 --- a/zig/build/cuda_family_host.zig +++ b/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); diff --git a/zig/build/dory.zig b/zig/build/dory.zig new file mode 100644 index 0000000000..b3396b21cd --- /dev/null +++ b/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); +} diff --git a/zig/src/families/dory/config.zig b/zig/src/families/dory/config.zig new file mode 100644 index 0000000000..a096a0fc1f --- /dev/null +++ b/zig/src/families/dory/config.zig @@ -0,0 +1,477 @@ +//! Nemotron-Labs Dory's dimensions from config.json: a looped nGPT stack, its loop count, and the layouts refused. +const std = @import("std"); +const eos_ids = @import("core").tokenizer.eos; + +pub const max_sublayers = 128; +pub const max_loops = 4; // llama.cpp's limit, and the card's largest benchmarked count +pub const max_eos = 4; +pub const max_width = 8192; // the rows kernels' widest row: 4 sixteen-byte loads a thread, 256 threads + +/// A sublayer: attention over a sliding window, attention over every key before, or a SwiGLU MLP. +pub const Kind = enum { sliding, full, mlp }; + +/// How the projections are stored: ModelOpt NVFP4 W4A16 (group 16) or plain bf16. +pub const Quant = enum { nvfp4, bf16 }; + +/// Why a check refused the checkpoint, kept for the caller's log line (tests read it instead). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +pub const Config = struct { + hidden: usize, + intermediate: usize, + vocab: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + sublayers: usize, // num_hidden_layers: physical sublayers, input then recurrent then output + input: usize, + recurrent: usize, + output: usize, + loops: usize, // n_recurrent_loops, or the load option that replaces it + kinds: [max_sublayers]Kind = @splat(.mlp), + window: usize, // a sliding sublayer's keys, the row's own included + theta_sliding: f64, + theta_full: f64, + std: f64, // initializer_range: every nGPT scale is stored divided by it + alpha_init: f64, + gate_init: f64, // attention_output_gate_logit_scale_init, sqrt(hidden) when null + positions: usize, // max_position_embeddings + quant: Quant, + eos: [max_eos]u32 = @splat(0), + eos_count: usize = 0, + + pub fn isEos(self: Config, token: u32) bool { + for (self.eos[0..self.eos_count]) |e| if (e == token) return true; + return false; + } + + pub fn qDim(self: Config) usize { + return self.heads * self.head_dim; + } + + pub fn kvDim(self: Config) usize { + return self.kv_heads * self.head_dim; + } + + /// Which block sublayer `i` is in: 0 input, 1 recurrent, 2 output. + pub fn block(self: Config, i: usize) u2 { + return if (i < self.input) 0 else if (i < self.input + self.recurrent) 1 else 2; + } + + /// Sublayers of `kind` in [from, to). + pub fn count(self: Config, kind: Kind, from: usize, to: usize) usize { + var n: usize = 0; + for (self.kinds[from..to]) |k| n += @intFromBool(k == kind); + return n; + } + + /// The same model run with `loops` passes over the recurrent block (a load option, 1 to 4). + pub fn withLoops(self: Config, loops: usize) !Config { + if (loops == 0 or loops > max_loops) return error.BadLoops; + var c = self; + c.loops = loops; + return c; + } + + /// The nGPT factors in fp32: float64 quotients rounded once, as torch casts a Python float for an fp32 multiply. + pub fn factors(self: Config) Factors { + const d: f64 = @floatFromInt(self.head_dim); + const h: f64 = @floatFromInt(self.hidden); + return .{ + .alpha = @floatCast(self.alpha_init / self.std), + .sqk = @floatCast(@sqrt(d) / self.std), + .gate = @floatCast(@sqrt(self.gate_init) / self.std), + .suv = @floatCast(@sqrt(h)), + .logit = 1 / @as(f32, @floatCast(self.std)), // torch divides by a scalar as a multiply by its fp32 reciprocal + .attn = @floatCast(1 / @sqrt(d)), + }; + } + + /// End tokens from generation_config.json's text, after config.json's own: replies end on either. + pub fn addEos(self: *Config, allocator: std.mem.Allocator, generation: []const u8) !void { + var parsed = try std.json.parseFromSlice(std.json.Value, allocator, generation, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadEos; + const ids = try eos_ids.merge(allocator, self.eos[0..self.eos_count], parsed.value, self.vocab); + defer allocator.free(ids); + try setEos(self, ids); + } + + /// `dir`/config.json, checked for the kernels' shapes, and generation_config.json's end tokens when it is there. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(text); + var why: Why = .{}; + var c = parse(gpa, text, &why) catch |e| { + if (why.len > 0) std.log.err("{s}: {s}", .{ path, why.text() }); + return e; + }; + checkShapes(c, &why) catch |e| { + std.log.err("{s}: {s}", .{ path, why.text() }); + return e; + }; + const gen_path = try std.fs.path.join(gpa, &.{ dir, "generation_config.json" }); + defer gpa.free(gen_path); + const gen = std.Io.Dir.cwd().readFileAlloc(io, gen_path, gpa, .limited(1 << 20)) catch |e| switch (e) { + error.FileNotFound => return c, + else => return e, + }; + defer gpa.free(gen); + try c.addEos(gpa, gen); + return c; + } +}; + +/// The fp32 factors the scale vectors are multiplied by (Config.factors). +pub const Factors = struct { alpha: f32, sqk: f32, gate: f32, suv: f32, logit: f32, attn: f32 }; + +fn int(obj: std.json.ObjectMap, key: []const u8, why: *Why) !usize { + const v = obj.get(key) orelse { + why.set("config.json has no {s}", .{key}); + return error.BadConfig; + }; + if (v != .integer or v.integer <= 0) { + why.set("config.json's {s} is not a positive integer", .{key}); + return error.BadConfig; + } + return @intCast(v.integer); +} + +fn number(obj: std.json.ObjectMap, key: []const u8, why: *Why) !f64 { + const v = obj.get(key) orelse { + why.set("config.json has no {s}", .{key}); + return error.BadConfig; + }; + const f: f64 = switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => 0, + }; + if (!(f > 0) or !std.math.isFinite(f)) { + why.set("config.json's {s} is not a positive number", .{key}); + return error.BadConfig; + } + return f; +} + +fn flag(obj: std.json.ObjectMap, key: []const u8) bool { + const v = obj.get(key) orelse return false; + return v == .bool and v.bool; +} + +fn setEos(c: *Config, ids: []const u32) !void { + if (ids.len > max_eos) return error.TooManyEos; + @memcpy(c.eos[0..ids.len], ids); + c.eos_count = ids.len; +} + +/// rope_parameters[kind]: a default (unscaled) rotation of every dimension; its theta. +fn rope(o: std.json.ObjectMap, kind: []const u8, why: *Why) !f64 { + const all = o.get("rope_parameters") orelse { + why.set("config.json has no rope_parameters", .{}); + return error.BadConfig; + }; + if (all != .object) return error.BadConfig; + const p = all.object.get(kind) orelse { + why.set("rope_parameters has no {s}", .{kind}); + return error.BadConfig; + }; + if (p != .object) return error.BadConfig; + if (p.object.get("rope_type")) |t| if (t != .string or !std.mem.eql(u8, t.string, "default")) { + why.set("{s}'s rope_type is not default; Dory on CUDA rotates without scaling", .{kind}); + return error.UnsupportedRope; + }; + if (p.object.get("partial_rotary_factor")) |f| { + const v: f64 = switch (f) { + .float => |x| x, + .integer => |x| @floatFromInt(x), + else => 0, + }; + if (v != 1) { + why.set("{s} rotates part of a head; Dory rotates all of it", .{kind}); + return error.UnsupportedRope; + } + } + return number(p.object, "rope_theta", why); +} + +/// The quantization config.json declares: none (bf16) or ModelOpt's W4A16_NVFP4 at group 16. +fn quantOf(o: std.json.ObjectMap, why: *Why) !Quant { + const q = o.get("quantization_config") orelse return .bf16; + if (q != .object) return error.BadConfig; + const method = q.object.get("quant_method"); + const algo = q.object.get("quant_algo"); + const modelopt = method != null and method.? == .string and std.mem.eql(u8, method.?.string, "modelopt"); + const w4a16 = algo != null and algo.? == .string and std.mem.eql(u8, algo.?.string, "W4A16_NVFP4"); + if (!modelopt or !w4a16) { + why.set("Dory on CUDA reads ModelOpt's W4A16_NVFP4 checkpoint or plain bf16; this one is quantized otherwise", .{}); + return error.UnsupportedQuantization; + } + if (q.object.get("config_groups")) |groups| if (groups == .object) { + var it = groups.object.iterator(); + while (it.next()) |g| { + const w = if (g.value_ptr.* == .object) g.value_ptr.object.get("weights") else null; + const gs = if (w != null and w.? == .object) w.?.object.get("group_size") else null; + if (gs == null or gs.? != .integer or gs.?.integer != 16) { + why.set("NVFP4 group {s} is not 16 weights a scale", .{g.key_ptr.*}); + return error.UnsupportedQuantization; + } + } + }; + return .nvfp4; +} + +/// config.json's text; a refusal's reason goes to `why`. +pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadConfig; + const o = parsed.value.object; + const mt = o.get("model_type") orelse return error.BadConfig; + if (mt != .string or !std.mem.eql(u8, mt.string, "dory")) { + why.set("model_type {f} is not dory", .{std.json.fmt(mt, .{})}); + return error.NotDory; + } + if (flag(o, "tie_word_embeddings") or flag(o, "mlp_bias") or flag(o, "attention_bias")) { + why.set("Dory has an untied head and no biases; this config ties or adds them", .{}); + return error.UnsupportedLayer; + } + if (o.get("hidden_act")) |a| if (a != .string or !std.mem.eql(u8, a.string, "silu")) { + why.set("hidden_act is not silu", .{}); + return error.UnsupportedLayer; + }; + // per_loop (vLLM's default): a cache per recurrent attention sublayer and loop; last_loop depends on chunking + if (o.get("recurrent_kv_cache_mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "per_loop")) { + why.set("recurrent_kv_cache_mode {f}: only per_loop gives the same tokens however a prompt is cut", .{std.json.fmt(m, .{})}); + return error.UnsupportedKvMode; + }; + var c = Config{ + .hidden = try int(o, "hidden_size", why), + .intermediate = try int(o, "intermediate_size", why), + .vocab = try int(o, "vocab_size", why), + .heads = try int(o, "num_attention_heads", why), + .kv_heads = try int(o, "num_key_value_heads", why), + .head_dim = try int(o, "head_dim", why), + .sublayers = try int(o, "num_hidden_layers", why), + .input = try int(o, "n_input_layers", why), + .recurrent = try int(o, "n_recurrent_layers", why), + .output = try int(o, "n_output_layers", why), + .loops = try int(o, "n_recurrent_loops", why), + .window = try int(o, "sliding_window", why), + .theta_sliding = try rope(o, "sliding_attention", why), + .theta_full = try rope(o, "full_attention", why), + .std = try number(o, "initializer_range", why), + .alpha_init = try number(o, "ngpt_alpha_init", why), + .gate_init = 0, + .positions = try int(o, "max_position_embeddings", why), + .quant = try quantOf(o, why), + }; + c.gate_init = if (o.get("attention_output_gate_logit_scale_init")) |g| (if (g == .null) @sqrt(@as(f64, @floatFromInt(c.hidden))) else try number(o, "attention_output_gate_logit_scale_init", why)) else @sqrt(@as(f64, @floatFromInt(c.hidden))); + if (c.sublayers > max_sublayers or c.input + c.recurrent + c.output != c.sublayers) { + why.set("{d} sublayers against {d} input, {d} recurrent and {d} output (at most {d})", .{ c.sublayers, c.input, c.recurrent, c.output, max_sublayers }); + return error.BadConfig; + } + if (c.loops > max_loops) { + why.set("{d} recurrent loops; Dory runs 1 to {d}", .{ c.loops, max_loops }); + return error.BadLoops; + } + const types = o.get("layer_types") orelse { + why.set("config.json has no layer_types", .{}); + return error.BadConfig; + }; + if (types != .array or types.array.items.len != c.sublayers) { + why.set("layer_types does not name each of the {d} sublayers", .{c.sublayers}); + return error.BadConfig; + } + for (types.array.items, 0..) |t, i| { + if (t != .string) return error.BadConfig; + const named = std.meta.stringToEnum(enum { sliding_attention, full_attention, mlp }, t.string) orelse { + why.set("sublayer {d} is {s}; Dory has sliding_attention, full_attention and mlp sublayers", .{ i, t.string }); + return error.UnsupportedLayer; + }; + c.kinds[i] = switch (named) { + .sliding_attention => .sliding, + .full_attention => .full, + .mlp => .mlp, + }; + } + const ids = try eos_ids.resolve(allocator, parsed.value, .null, c.vocab); + defer allocator.free(ids); + try setEos(&c, ids); + return c; +} + +/// What the shared kernels take: head dim 256 in whole GQA groups, rows to 8192 wide in 64s, a window of 2 or more. +pub fn checkShapes(c: Config, why: *Why) !void { + if (c.head_dim != 256 or c.kv_heads == 0 or c.heads % c.kv_heads != 0 or c.heads / c.kv_heads > 16) { + why.set("head dim {d}, {d} query heads over {d} key heads: this engine takes head dim 256 and up to 16 query heads a key head", .{ c.head_dim, c.heads, c.kv_heads }); + return error.UnsupportedShapes; + } + for ([_]usize{ c.hidden, c.intermediate, c.qDim() }) |w| if (w % 64 != 0 or w > max_width) { + why.set("hidden {d}, intermediate {d}, attention {d}: the rows kernels take multiples of 64 up to {d}", .{ c.hidden, c.intermediate, c.qDim(), max_width }); + return error.UnsupportedShapes; + }; + if (c.window < 2) { + why.set("sliding window {d}: the window kernels take 2 or more keys", .{c.window}); + return error.UnsupportedShapes; + } + if (c.count(.mlp, 0, c.sublayers) == c.sublayers) { + why.set("no attention sublayer", .{}); + return error.UnsupportedShapes; + } +} + +/// The NVFP4 checkpoint's config.json at the pinned revision (41f02335), its layer_types written as the pattern. +pub const pinned = + \\{"architectures": ["DoryForCausalLM"], "model_type": "dory", "vocab_size": 131072, "hidden_size": 2560, + \\ "num_hidden_layers": 72, "intermediate_size": 7168, "num_attention_heads": 16, "num_key_value_heads": 4, + \\ "max_position_embeddings": 262144, "tie_word_embeddings": false, "ngpt_alpha_init": 0.05, + \\ "attention_output_gate_logit_scale_init": null, "mlp_bias": false, "attention_bias": false, + \\ "n_input_layers": 16, "n_recurrent_layers": 40, "n_output_layers": 16, "n_recurrent_loops": 2, + \\ "quantization_config": {"config_groups": {"group_0": {"weights": {"dynamic": false, "num_bits": 4, + \\ "type": "float", "group_size": 16}, "targets": ["Linear"]}}, "ignore": ["backbone.embeddings", "lm_head"], + \\ "quant_algo": "W4A16_NVFP4", "producer": {"name": "modelopt", "version": "0.0.0"}, "quant_method": "modelopt"}, + \\ "head_dim": 256, "hidden_act": "silu", "initializer_range": 0.0177, "dtype": "bfloat16", "sliding_window": 1024, + \\ "layer_types": [S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, + \\ S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, S, M, S, M, S, M, F, M, + \\ S, M, S, M, S, M, F, M], + \\ "rope_parameters": {"sliding_attention": {"rope_type": "default", "rope_theta": 10000, "partial_rotary_factor": 1.0}, + \\ "full_attention": {"rope_type": "default", "rope_theta": 10000000, "partial_rotary_factor": 1.0}}} +; + +/// The pinned generation_config.json: replies end on or <|im_end|>. +pub const pinned_generation = + \\{"_from_model_config": true, "bos_token_id": 1, "do_sample": true, "eos_token_id": [2, 11], "pad_token_id": 11, + \\ "top_p": 0.95, "temperature": 1.0, "transformers_version": "4.57.6"} +; + +const testing = std.testing; + +/// `pinned` with its pattern letters spelled out, as config.json has them. +pub fn pinnedText(a: std.mem.Allocator) ![]u8 { + const s = try std.mem.replaceOwned(u8, a, pinned, "S,", "\"sliding_attention\","); + defer a.free(s); + const f = try std.mem.replaceOwned(u8, a, s, "F,", "\"full_attention\","); + defer a.free(f); + const m = try std.mem.replaceOwned(u8, a, f, "M,", "\"mlp\","); + defer a.free(m); + return std.mem.replaceOwned(u8, a, m, "M]", "\"mlp\"]"); +} + +fn replaced(from: []const u8, to: []const u8) ![]u8 { + const t = try pinnedText(testing.allocator); + defer testing.allocator.free(t); + return std.mem.replaceOwned(u8, testing.allocator, t, from, to); +} + +fn refused(text: []const u8, err: anyerror, word: []const u8) !void { + var why: Why = .{}; + try testing.expectError(err, parse(testing.allocator, text, &why)); + if (std.mem.indexOf(u8, why.text(), word) == null) { + std.debug.print("refusal \"{s}\" does not name \"{s}\"\n", .{ why.text(), word }); + return error.TestUnexpectedResult; + } +} + +test "parse the pinned config: blocks, sublayer kinds, RoPE, nGPT factors and end tokens" { + const text = try pinnedText(testing.allocator); + defer testing.allocator.free(text); + var why: Why = .{}; + var c = try parse(testing.allocator, text, &why); + try checkShapes(c, &why); + try testing.expectEqual(@as(usize, 72), c.sublayers); + try testing.expectEqual(@as(usize, 2), c.loops); + try testing.expectEqual(Quant.nvfp4, c.quant); + try testing.expectEqual(Kind.sliding, c.kinds[0]); + try testing.expectEqual(Kind.mlp, c.kinds[1]); + for ([_]usize{ 6, 14, 22, 30, 38, 46, 54, 62, 70 }) |i| try testing.expectEqual(Kind.full, c.kinds[i]); + try testing.expectEqual(@as(usize, 2), c.count(.full, 0, 16)); + try testing.expectEqual(@as(usize, 5), c.count(.full, 16, 56)); + try testing.expectEqual(@as(usize, 15), c.count(.sliding, 16, 56)); + try testing.expectEqual(@as(usize, 36), c.count(.mlp, 0, 72)); + try testing.expectEqual(@as(f64, 1e4), c.theta_sliding); + try testing.expectEqual(@as(f64, 1e7), c.theta_full); + try testing.expectEqual(@sqrt(@as(f64, 2560)), c.gate_init); + const f = c.factors(); + try testing.expectEqual(@as(f32, @floatCast(0.05 / 0.0177)), f.alpha); + try testing.expectApproxEqRel(@as(f32, 903.9548), f.sqk, 1e-6); + try testing.expectApproxEqRel(@as(f32, 401.9), f.gate, 1e-4); + try testing.expectApproxEqRel(@as(f32, 50.596443), f.suv, 1e-7); + try testing.expectEqual(@as(f32, 1.0 / 16.0), f.attn); + try testing.expectEqual(@as(usize, 0), c.eos_count); // config.json declares none + try c.addEos(testing.allocator, pinned_generation); + try testing.expectEqualSlices(u32, &.{ 2, 11 }, c.eos[0..c.eos_count]); + try testing.expect(c.isEos(11) and !c.isEos(1)); + try testing.expectEqual(@as(usize, 4), (try c.withLoops(4)).loops); + try testing.expectError(error.BadLoops, c.withLoops(0)); + try testing.expectError(error.BadLoops, c.withLoops(5)); +} + +test "refusals: not dory, last_loop KV, other quantization, scaled RoPE, ties, other sublayers, counts" { + const not = try replaced("\"model_type\": \"dory\"", "\"model_type\": \"llama\""); + defer testing.allocator.free(not); + try refused(not, error.NotDory, "dory"); + const last = try replaced("\"n_recurrent_loops\": 2,", "\"n_recurrent_loops\": 2, \"recurrent_kv_cache_mode\": \"last_loop\","); + defer testing.allocator.free(last); + try refused(last, error.UnsupportedKvMode, "per_loop"); + const mx = try replaced("\"quant_algo\": \"W4A16_NVFP4\"", "\"quant_algo\": \"MXFP8\""); + defer testing.allocator.free(mx); + try refused(mx, error.UnsupportedQuantization, "W4A16_NVFP4"); + const group = try replaced("\"group_size\": 16", "\"group_size\": 32"); + defer testing.allocator.free(group); + try refused(group, error.UnsupportedQuantization, "16"); + const yarn = try replaced("{\"rope_type\": \"default\", \"rope_theta\": 10000000", "{\"rope_type\": \"yarn\", \"rope_theta\": 10000000"); + defer testing.allocator.free(yarn); + try refused(yarn, error.UnsupportedRope, "full_attention"); + const partial = try replaced("\"rope_theta\": 10000, \"partial_rotary_factor\": 1.0", "\"rope_theta\": 10000, \"partial_rotary_factor\": 0.5"); + defer testing.allocator.free(partial); + try refused(partial, error.UnsupportedRope, "part"); + const tied = try replaced("\"tie_word_embeddings\": false", "\"tie_word_embeddings\": true"); + defer testing.allocator.free(tied); + try refused(tied, error.UnsupportedLayer, "untied"); + const moe = try replaced("\"mlp\"]", "\"moe\"]"); + defer testing.allocator.free(moe); + try refused(moe, error.UnsupportedLayer, "moe"); + const counts = try replaced("\"n_output_layers\": 16", "\"n_output_layers\": 15"); + defer testing.allocator.free(counts); + try refused(counts, error.BadConfig, "15 output"); + const loops = try replaced("\"n_recurrent_loops\": 2", "\"n_recurrent_loops\": 5"); + defer testing.allocator.free(loops); + try refused(loops, error.BadLoops, "5 recurrent loops"); +} + +test "bf16 checkpoints parse without a quantization_config; shapes outside the kernels are refused" { + const t = try pinnedText(testing.allocator); + defer testing.allocator.free(t); + const at = std.mem.indexOf(u8, t, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, t, "\"head_dim\"").?; + const plain = try std.mem.concat(testing.allocator, u8, &.{ t[0..at], t[end..] }); + defer testing.allocator.free(plain); + var why: Why = .{}; + const c = try parse(testing.allocator, plain, &why); + try testing.expectEqual(Quant.bf16, c.quant); + var narrow = c; + narrow.head_dim = 128; + try testing.expectError(error.UnsupportedShapes, checkShapes(narrow, &why)); + var wide = c; + wide.intermediate = 9216; + try testing.expectError(error.UnsupportedShapes, checkShapes(wide, &why)); + var odd = c; + odd.kv_heads = 3; + try testing.expectError(error.UnsupportedShapes, checkShapes(odd, &why)); +} diff --git a/zig/src/families/dory/dory.zig b/zig/src/families/dory/dory.zig new file mode 100644 index 0000000000..567003f03d --- /dev/null +++ b/zig/src/families/dory/dory.zig @@ -0,0 +1,14 @@ +//! Nemotron-Labs Dory 1.0 (a looped nGPT stack) on CUDA: config, weight map, sublayer schedule and memory. + +pub const config = @import("config.zig"); +pub const Config = config.Config; +pub const names = @import("names.zig"); +pub const schedule = @import("schedule.zig"); +pub const memory = @import("memory.zig"); + +test { + _ = @import("config.zig"); + _ = @import("names.zig"); + _ = @import("schedule.zig"); + _ = @import("memory.zig"); +} diff --git a/zig/src/families/dory/memory.zig b/zig/src/families/dory/memory.zig new file mode 100644 index 0000000000..c038654e3e --- /dev/null +++ b/zig/src/families/dory/memory.zig @@ -0,0 +1,72 @@ +//! What Dory holds on the GPU: weight bytes as stored, KV bytes by loops, context and streams (per_loop caches). +const std = @import("std"); +const config = @import("config.zig"); +const schedule = @import("schedule.zig"); + +const Config = config.Config; + +/// A projection's bytes as stored: NVFP4 codes, e4m3 scales a 16 and the fp32 scale, or bf16. +pub fn projBytes(c: Config, n: usize, k: usize) usize { + return switch (c.quant) { + .nvfp4 => n * k / 2 + n * k / 16 + 4, + .bf16 => n * k * 2, + }; +} + +/// Every tensor's bytes, as the checkpoint stores them (the loader's packs pad rows to 128 on top). +pub fn weightBytes(c: Config) usize { + const d = c.hidden; + var total: usize = 2 * (2 * c.vocab * d) + 2 * c.vocab; // embeddings, head, logit_scale + for (c.kinds[0..c.sublayers]) |k| total += 2 * d + if (k == .mlp) + 2 * projBytes(c, c.intermediate, d) + projBytes(c, d, c.intermediate) + 2 * 2 * c.intermediate + else + 2 * projBytes(c, c.qDim(), d) + 2 * projBytes(c, c.kvDim(), d) + projBytes(c, d, c.qDim()) + 2 * 2 * c.qDim(); + return total; +} + +/// Parameters, counting each NVFP4 or bf16 weight once (scales not counted). +pub fn parameters(c: Config) usize { + const d = c.hidden; + var total: usize = 2 * c.vocab * d + c.vocab; + for (c.kinds[0..c.sublayers]) |k| total += d + if (k == .mlp) 3 * c.intermediate * d + 2 * c.intermediate else 2 * c.qDim() * d + 2 * c.kvDim() * d + d * c.qDim() + 2 * c.qDim(); + return total; +} + +/// bf16 K and V of one token in one cache. +pub fn tokenBytes(c: Config) usize { + return 2 * c.kvDim() * 2; +} + +/// One stream's KV at `c.loops`: full caches hold `context` positions, sliding caches a ring of `ring` keys. +pub fn streamBytes(c: Config, context: usize, ring: usize) usize { + const s = schedule.build(c); + return (s.cacheCount(.full) * context + s.cacheCount(.sliding) * ring) * tokenBytes(c); +} + +const testing = std.testing; + +fn pinned() !Config { + const text = try config.pinnedText(testing.allocator); + defer testing.allocator.free(text); + var why: config.Why = .{}; + return config.parse(testing.allocator, text, &why); +} + +test "3.97B parameters; the NVFP4 tensors' bytes equal the pinned file's data (3,202,498,032 bytes less its header)" { + const c = try pinned(); + try testing.expectEqual(@as(usize, 3_975_229_440), parameters(c)); + // model.safetensors at 41f02335: 8 bytes of length, a 121,192-byte header, then the tensors + try testing.expectEqual(@as(usize, 3_202_498_032 - 8 - 121_192), weightBytes(c)); +} + +test "KV a token: 36, 56, 96 KiB at 1, 2, 4 loops; 7 GiB at 128K and two loops, plus fixed sliding rings" { + const c = try pinned(); + try testing.expectEqual(@as(usize, 4096), tokenBytes(c)); + for ([_][2]usize{ .{ 1, 36 }, .{ 2, 56 }, .{ 4, 96 } }) |row| { + const cl = try c.withLoops(row[0]); + try testing.expectEqual(row[1] * 1024, streamBytes(cl, 1, 0)); + } + // 14 full caches over 128K keys, 42 sliding rings of 1,024 keys: 7 GiB + 168 MiB + try testing.expectEqual((@as(usize, 7) << 30) + (168 << 20), streamBytes(try c.withLoops(2), 128 << 10, 1024)); + try testing.expectEqual((@as(usize, 24) << 30) + (288 << 20), streamBytes(try c.withLoops(4), 256 << 10, 1024)); +} diff --git a/zig/src/families/dory/names.zig b/zig/src/families/dory/names.zig new file mode 100644 index 0000000000..c3fca61e0d --- /dev/null +++ b/zig/src/families/dory/names.zig @@ -0,0 +1,210 @@ +//! Dory's checkpoint tensors by name, each checked for dtype and shape before any upload; nothing may be left over. +const std = @import("std"); +const core = @import("core"); +const config = @import("config.zig"); + +const Config = config.Config; +const Tensor = core.checkpoint.Tensor; +const Checkpoint = core.checkpoint.Checkpoint; + +/// A projection [n, k] as stored: NVFP4 codes [n, k/2] with e4m3 scales a 16 and an fp32 scale, or bf16 [n, k]. +pub const Proj = struct { + w: Tensor, + blocks: ?Tensor = null, // NVFP4's weight_scale + scale: f32 = 1, // NVFP4's weight_scale_2 + n: usize, + k: usize, +}; + +pub const Attn = struct { alpha: Tensor, q: Proj, k: Proj, v: Proj, g: Proj, o: Proj, sqk: Tensor, s_gate: Tensor }; +pub const Mlp = struct { alpha: Tensor, gate: Proj, up: Proj, down: Proj, suv_gate: Tensor, suv_up: Tensor }; +pub const Sub = union(enum) { attn: Attn, mlp: Mlp }; + +pub const Map = struct { + gpa: std.mem.Allocator, + embed: Tensor, // bf16 [V, D] + head: Tensor, // bf16 [V, D] + logit_scale: Tensor, // bf16 [V] + subs: []Sub, + + pub fn deinit(m: *Map) void { + m.gpa.free(m.subs); + m.* = undefined; + } +}; + +const Names = struct { + ck: *Checkpoint, + quant: config.Quant, + buf: [160]u8 = undefined, + + fn name(self: *Names, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(&self.buf, fmt, args); + } + + fn bf16(self: *Names, shape: []const usize, comptime fmt: []const u8, args: anytype) !Tensor { + return self.ck.expect(try self.name(fmt, args), .bf16, shape); + } + + fn proj(self: *Names, n: usize, k: usize, comptime fmt: []const u8, args: anytype) !Proj { + if (self.quant == .bf16) return .{ .w = try self.bf16(&.{ n, k }, fmt ++ ".weight", args), .n = n, .k = k }; + const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .u8, &.{ n, k / 2 }); + const b = try self.ck.expect(try self.name(fmt ++ ".weight_scale", args), .f8_e4m3, &.{ n, k / 16 }); + const s = try self.ck.expect(try self.name(fmt ++ ".weight_scale_2", args), .f32, &.{}); + const v: f32 = @bitCast(std.mem.readInt(u32, s.bytes[0..4], .little)); + if (!(v > 0) or !std.math.isFinite(v)) { + std.log.err("{s}: weight_scale_2 {e} is not a positive finite scale", .{ try self.name(fmt, args), v }); + return error.UnexpectedTensor; + } + return .{ .w = w, .blocks = b, .scale = v, .n = n, .k = k }; + } +}; + +/// Every tensor the loader reads, sublayer by sublayer; a missing, mistyped or leftover tensor refuses the checkpoint. +pub fn map(gpa: std.mem.Allocator, ck: *Checkpoint, c: Config) !Map { + if (c.hidden % 16 != 0 or c.intermediate % 16 != 0 or c.qDim() % 16 != 0) return error.UnsupportedShapes; + var nm: Names = .{ .ck = ck, .quant = c.quant }; + const d = c.hidden; + const f = c.intermediate; + const q = c.qDim(); + const kv = c.kvDim(); + var m: Map = .{ .gpa = gpa, .embed = undefined, .head = undefined, .logit_scale = undefined, .subs = try gpa.alloc(Sub, c.sublayers) }; + errdefer gpa.free(m.subs); + m.embed = try nm.bf16(&.{ c.vocab, d }, "backbone.embeddings.weight", .{}); + for (m.subs, 0..) |*s, i| { + const pre = "backbone.layers.{d}."; + s.* = if (c.kinds[i] == .mlp) .{ .mlp = .{ + .alpha = try nm.bf16(&.{d}, pre ++ "mlp_alpha", .{i}), + .gate = try nm.proj(f, d, pre ++ "mixer.gate_proj", .{i}), + .up = try nm.proj(f, d, pre ++ "mixer.up_proj", .{i}), + .down = try nm.proj(d, f, pre ++ "mixer.down_proj", .{i}), + .suv_gate = try nm.bf16(&.{f}, pre ++ "mixer.suv_gate", .{i}), + .suv_up = try nm.bf16(&.{f}, pre ++ "mixer.suv_up", .{i}), + } } else .{ .attn = .{ + .alpha = try nm.bf16(&.{d}, pre ++ "attn_alpha", .{i}), + .q = try nm.proj(q, d, pre ++ "mixer.q_proj", .{i}), + .k = try nm.proj(kv, d, pre ++ "mixer.k_proj", .{i}), + .v = try nm.proj(kv, d, pre ++ "mixer.v_proj", .{i}), + .g = try nm.proj(q, d, pre ++ "mixer.g_proj", .{i}), + .o = try nm.proj(d, q, pre ++ "mixer.o_proj", .{i}), + .sqk = try nm.bf16(&.{q}, pre ++ "mixer.sqk", .{i}), + .s_gate = try nm.bf16(&.{q}, pre ++ "mixer.s_gate", .{i}), + } }; + } + m.head = try nm.bf16(&.{ c.vocab, d }, "lm_head.weight", .{}); + m.logit_scale = try nm.bf16(&.{c.vocab}, "lm_head.logit_scale", .{}); + if (ck.unused() != 0) return error.UnusedTensors; + return m; +} + +const testing = std.testing; +const io = testing.io; + +/// A tiny Dory: 2 input, 4 recurrent and 2 output sublayers, attention and MLP alternating, one of each full. +pub const tiny: Config = blk: { + var c: Config = .{ .hidden = 64, .intermediate = 96, .vocab = 32, .heads = 4, .kv_heads = 2, .head_dim = 16, .sublayers = 8, .input = 2, .recurrent = 4, .output = 2, .loops = 2, .window = 3, .theta_sliding = 1e4, .theta_full = 1e6, .std = 0.0177, .alpha_init = 0.05, .gate_init = 8, .positions = 64, .quant = .nvfp4 }; + for (0..8) |i| c.kinds[i] = if (i % 2 == 1) .mlp else if (i == 4) .full else .sliding; + break :blk c; +}; + +/// A safetensors file with each tensor `c` names, zero-filled, every NVFP4 scale 1. +pub fn writeFile(a: std.mem.Allocator, dir: std.Io.Dir, c: Config) !void { + const E = struct { name: []const u8, dtype: []const u8, shape: []const usize }; + var entries: std.ArrayList(E) = .empty; + const Add = struct { + fn bf(list: *std.ArrayList(E), al: std.mem.Allocator, nm: []const u8, shape: []const usize) !void { + try list.append(al, .{ .name = nm, .dtype = "BF16", .shape = try al.dupe(usize, shape) }); + } + fn proj(list: *std.ArrayList(E), al: std.mem.Allocator, cc: Config, pre: []const u8, n: usize, k: usize) !void { + if (cc.quant == .bf16) return bf(list, al, try std.fmt.allocPrint(al, "{s}.weight", .{pre}), &.{ n, k }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "U8", .shape = try al.dupe(usize, &.{ n, k / 2 }) }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k / 16 }) }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_2", .{pre}), .dtype = "F32", .shape = &.{} }); + } + }; + const d = c.hidden; + try Add.bf(&entries, a, "backbone.embeddings.weight", &.{ c.vocab, d }); + try Add.bf(&entries, a, "lm_head.weight", &.{ c.vocab, d }); + try Add.bf(&entries, a, "lm_head.logit_scale", &.{c.vocab}); + for (0..c.sublayers) |i| { + const p = try std.fmt.allocPrint(a, "backbone.layers.{d}.", .{i}); + if (c.kinds[i] == .mlp) { + try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}mlp_alpha", .{p}), &.{d}); + for ([_][]const u8{ "gate_proj", "up_proj" }) |n| try Add.proj(&entries, a, c, try std.fmt.allocPrint(a, "{s}mixer.{s}", .{ p, n }), c.intermediate, d); + try Add.proj(&entries, a, c, try std.fmt.allocPrint(a, "{s}mixer.down_proj", .{p}), d, c.intermediate); + for ([_][]const u8{ "suv_gate", "suv_up" }) |n| try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}mixer.{s}", .{ p, n }), &.{c.intermediate}); + } else { + try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}attn_alpha", .{p}), &.{d}); + for ([_][]const u8{ "q_proj", "g_proj" }) |n| try Add.proj(&entries, a, c, try std.fmt.allocPrint(a, "{s}mixer.{s}", .{ p, n }), c.qDim(), d); + for ([_][]const u8{ "k_proj", "v_proj" }) |n| try Add.proj(&entries, a, c, try std.fmt.allocPrint(a, "{s}mixer.{s}", .{ p, n }), c.kvDim(), d); + try Add.proj(&entries, a, c, try std.fmt.allocPrint(a, "{s}mixer.o_proj", .{p}), d, c.qDim()); + for ([_][]const u8{ "sqk", "s_gate" }) |n| try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}mixer.{s}", .{ p, n }), &.{c.qDim()}); + } + } + var header: std.Io.Writer.Allocating = .init(a); + try header.writer.writeAll("{"); + var at: usize = 0; + for (entries.items, 0..) |e, j| { + var n: usize = if (e.dtype[0] == 'B') 2 else if (e.dtype[1] == '3') 4 else 1; + for (e.shape) |x| n *= x; + try header.writer.print("{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ if (j == 0) "" else ",", e.name, e.dtype }); + for (e.shape, 0..) |x, i| try header.writer.print("{s}{d}", .{ if (i == 0) "" else ",", x }); + try header.writer.print("],\"data_offsets\":[{d},{d}]}}", .{ at, at + n }); + at += n; + } + try header.writer.writeAll("}"); + const h = header.written(); + const file = try a.alloc(u8, 8 + h.len + at); + @memset(file, 0); + std.mem.writeInt(u64, file[0..8], h.len, .little); + @memcpy(file[8..][0..h.len], h); + // every weight_scale_2 one: the map refuses a zero scale + var off: usize = 0; + for (entries.items) |e| { + var n: usize = if (e.dtype[0] == 'B') 2 else if (e.dtype[1] == '3') 4 else 1; + for (e.shape) |x| n *= x; + if (std.mem.endsWith(u8, e.name, "weight_scale_2")) std.mem.writeInt(u32, file[8 + h.len + off ..][0..4], @bitCast(@as(f32, 1)), .little); + off += n; + } + try dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file }); +} + +fn tmpPath(a: std.mem.Allocator, tmp: testing.TmpDir) ![]u8 { + const cwd = try std.process.currentPathAlloc(io, a); + return std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); +} + +test "every tensor of a tiny NVFP4 and a tiny bf16 checkpoint maps" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for ([_]config.Quant{ .nvfp4, .bf16 }) |quant| { + var c = tiny; + c.quant = quant; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeFile(a, tmp.dir, c); + var ck = try Checkpoint.openModel(testing.allocator, io, try tmpPath(a, tmp)); + defer ck.close(); + var m = try map(testing.allocator, &ck, c); + defer m.deinit(); + try testing.expectEqual(@as(usize, 8), m.subs.len); + try testing.expect(m.subs[1] == .mlp and m.subs[4] == .attn); + const o = m.subs[4].attn.o; + try testing.expectEqual(@as(usize, 64), o.n); + try testing.expectEqual(@as(usize, 64), o.k); + try testing.expectEqual(quant == .nvfp4, o.blocks != null); + } +} + +test "the pinned NVFP4 checkpoint (TF_DORY_DIR): 1,083 tensors, none left over" { + const dir = @import("sys").getenv("TF_DORY_DIR") orelse return error.SkipZigTest; + const c = try Config.read(testing.allocator, io, dir); + var ck = try Checkpoint.openModel(testing.allocator, io, dir); + defer ck.close(); + var m = try map(testing.allocator, &ck, c); + defer m.deinit(); + try testing.expectEqual(@as(usize, 72), m.subs.len); + try testing.expect(m.subs[71] == .mlp and m.subs[70] == .attn); + try testing.expectEqualSlices(u32, &.{ 2, 11 }, c.eos[0..c.eos_count]); +} diff --git a/zig/src/families/dory/schedule.zig b/zig/src/families/dory/schedule.zig new file mode 100644 index 0000000000..935745c9ae --- /dev/null +++ b/zig/src/families/dory/schedule.zig @@ -0,0 +1,154 @@ +//! Dory's sublayer order: input, the recurrent block `loops` times, output; per_loop KV caches, one an attention step. +const std = @import("std"); +const config = @import("config.zig"); + +const Config = config.Config; +const Kind = config.Kind; + +pub const max_steps = config.max_sublayers * config.max_loops; + +/// One virtual sublayer: its weights (`sub`), loop, cache (attention only) and the loop's entry and exit norms. +pub const Step = struct { + sub: u16, + kind: Kind, + loop: u8 = 0, // 0 outside the recurrent block + cache: ?u16 = null, + /// h = l2(state + input_h) before the sublayer; at loop 0 input_h = state = h first. + enter: bool = false, + /// state = l2(h) after the sublayer: the loop's last. + leave: bool = false, +}; + +pub const Schedule = struct { + steps: [max_steps]Step = undefined, + len: usize = 0, + caches: usize = 0, + cache_kind: [max_steps]Kind = undefined, // sliding or full, by cache + + pub fn slice(s: *const Schedule) []const Step { + return s.steps[0..s.len]; + } + + /// Caches of `kind` (sliding or full). + pub fn cacheCount(s: *const Schedule, kind: Kind) usize { + var n: usize = 0; + for (s.cache_kind[0..s.caches]) |k| n += @intFromBool(k == kind); + return n; + } + + /// vLLM #60149's attention layer index of a step: loop * sublayers + sub (its per_loop virtual layers). + pub fn vllmLayer(c: Config, st: Step) usize { + return @as(usize, st.loop) * c.sublayers + st.sub; + } +}; + +/// The steps of `c` at `c.loops` loops: I + R * loops + O, caches numbered in the order the steps run. +pub fn build(c: Config) Schedule { + var s: Schedule = .{}; + const add = struct { + fn one(sch: *Schedule, cc: Config, sub: usize, loop: usize) void { + var st: Step = .{ .sub = @intCast(sub), .kind = cc.kinds[sub], .loop = @intCast(loop) }; + if (st.kind != .mlp) { + st.cache = @intCast(sch.caches); + sch.cache_kind[sch.caches] = st.kind; + sch.caches += 1; + } + sch.steps[sch.len] = st; + sch.len += 1; + } + }.one; + for (0..c.input) |i| add(&s, c, i, 0); + for (0..c.loops) |l| { + for (c.input..c.input + c.recurrent) |i| add(&s, c, i, l); + s.steps[s.len - c.recurrent].enter = true; + s.steps[s.len - 1].leave = true; + } + for (c.input + c.recurrent..c.sublayers) |i| add(&s, c, i, 0); + return s; +} + +const testing = std.testing; + +fn pinnedConfig() !Config { + const text = try config.pinnedText(testing.allocator); + defer testing.allocator.free(text); + var why: config.Why = .{}; + return config.parse(testing.allocator, text, &why); +} + +/// llama.cpp #30306's `physical` and per_loop `cache_layer` (llama_model_dory::load_arch_hparams at 3834b4ae). +fn llamaCpp(c: Config, physical: []u32, cache_layer: []u32) usize { + var n: usize = 0; + for (0..c.input) |i| { + physical[n] = @intCast(i); + n += 1; + } + for (0..c.loops) |_| for (0..c.recurrent) |i| { + physical[n] = @intCast(c.input + i); + n += 1; + }; + for (0..c.output) |i| { + physical[n] = @intCast(c.input + c.recurrent + i); + n += 1; + } + for (0..n) |il| cache_layer[il] = @intCast(il); + return n; +} + +test "the schedule at 1 to 4 loops: llama.cpp's physical sublayers, a cache an attention step, 8 + 20 L + 8 caches" { + const base = try pinnedConfig(); + for (1..5) |loops| { + const c = try base.withLoops(loops); + const s = build(c); + try testing.expectEqual(16 + 40 * loops + 16, s.len); + try testing.expectEqual(8 + 20 * loops + 8, s.caches); + try testing.expectEqual(4 + 5 * loops, s.cacheCount(.full)); + try testing.expectEqual(12 + 15 * loops, s.cacheCount(.sliding)); + var physical: [max_steps]u32 = undefined; + var layer: [max_steps]u32 = undefined; + try testing.expectEqual(s.len, llamaCpp(c, &physical, &layer)); + // llama.cpp numbers its caches by virtual sublayer and skips the MLP ones: the same order as ours + var seen: usize = 0; + for (s.slice(), 0..) |st, il| { + try testing.expectEqual(physical[il], st.sub); + if (st.kind == .mlp) { + try testing.expectEqual(@as(?u16, null), st.cache); + continue; + } + try testing.expectEqual(@as(?u16, @intCast(seen)), st.cache); + try testing.expect(seen == 0 or layer[il] > layer[il - 1]); + seen += 1; + } + var enters: usize = 0; + var leaves: usize = 0; + for (s.slice(), 0..) |st, i| { + if (st.enter) { + enters += 1; + try testing.expectEqual(@as(u16, 16), st.sub); + try testing.expectEqual(16 + 40 * (enters - 1), i); + } + if (st.leave) { + leaves += 1; + try testing.expectEqual(@as(u16, 55), st.sub); + try testing.expectEqual(st.loop, @as(u8, @intCast(leaves - 1))); + } + } + try testing.expectEqual(loops, enters); + try testing.expectEqual(loops, leaves); + } +} + +test "vLLM's per_loop virtual layers are unique per (loop, sublayer) and keep the physical index at loop 0" { + const c = try (try pinnedConfig()).withLoops(4); + const s = build(c); + var used: [config.max_sublayers * config.max_loops]bool = @splat(false); + for (s.slice()) |st| { + if (st.kind == .mlp) continue; + const v = Schedule.vllmLayer(c, st); + try testing.expect(!used[v]); + used[v] = true; + if (st.loop == 0) try testing.expectEqual(@as(usize, st.sub), v); + if (st.loop > 0) try testing.expect(st.sub >= 16 and st.sub < 56); + } + try testing.expectEqual(@as(usize, 22 + 72), Schedule.vllmLayer(c, .{ .sub = 22, .kind = .full, .loop = 1 })); +} From aba44377d37835c9e32490da692d4dae041ca666 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 09:56:56 +0200 Subject: [PATCH 7/9] dory: the whole forward in float64, per_loop chunk-invariant to the bit ref.zig runs Dory's forward in float64 over prompt chunks with a stream's KV caches, following vLLM #60149's dory.py (ce9dd0ed): the nGPT slerp residual, q and k L2-normed after RoPE with q scaled by (sqk f)^2, the scaled sigmoid output gate and the row L2 over all heads, the scaled SwiGLU and its row L2, the loop entry l2(state + input_h) and exit l2(h), and the logit scale. Options give dory.py's bf16 rounding points with torch's fp32 scale vectors, per-step taps, and last_loop KV. ref_test.zig checks, on a tiny random config (2 + 4 + 2 sublayers, window 3): - per_loop: one chunk, several chunkings and one-token decode give the same logits to the bit, at 1 to 4 loops, in float64 and at the bf16 points; - last_loop: equal to per_loop at one loop and within one chunk, but at two loops the rows after a cut differ, the chunk dependence that keeps it out of this family; - every logit within 1e-12 of dory.py itself run in float64 (its fp32 casts as float64) on the same weights at 1, 2 and 4 loops (ref_fixture.zig; the oracle is a dev-time script, not part of the tree). --- zig/src/families/dory/dory.zig | 5 +- zig/src/families/dory/ref.zig | 398 ++++++++++++++++++++++++++ zig/src/families/dory/ref_fixture.zig | 83 ++++++ zig/src/families/dory/ref_test.zig | 118 ++++++++ 4 files changed, 603 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/dory/ref.zig create mode 100644 zig/src/families/dory/ref_fixture.zig create mode 100644 zig/src/families/dory/ref_test.zig diff --git a/zig/src/families/dory/dory.zig b/zig/src/families/dory/dory.zig index 567003f03d..50fc70624c 100644 --- a/zig/src/families/dory/dory.zig +++ b/zig/src/families/dory/dory.zig @@ -1,14 +1,17 @@ -//! Nemotron-Labs Dory 1.0 (a looped nGPT stack) on CUDA: config, weight map, sublayer schedule and memory. +//! Nemotron-Labs Dory 1.0 (a looped nGPT stack) on CUDA: config, weights, schedule, memory, float64 reference. pub const config = @import("config.zig"); pub const Config = config.Config; pub const names = @import("names.zig"); pub const schedule = @import("schedule.zig"); pub const memory = @import("memory.zig"); +/// The float64 forward the kernels and the CUDA model are checked against. +pub const ref = @import("ref.zig"); test { _ = @import("config.zig"); _ = @import("names.zig"); _ = @import("schedule.zig"); _ = @import("memory.zig"); + _ = @import("ref.zig"); } diff --git a/zig/src/families/dory/ref.zig b/zig/src/families/dory/ref.zig new file mode 100644 index 0000000000..076cd41283 --- /dev/null +++ b/zig/src/families/dory/ref.zig @@ -0,0 +1,398 @@ +//! Dory's whole forward in float64 (vLLM #60149's dory.py semantics), over prompt chunks with per_loop KV caches. +const std = @import("std"); +const core = @import("core"); +const config = @import("config.zig"); +const schedule = @import("schedule.zig"); + +const Config = config.Config; +const bfr = core.bf16_round; + +pub const eps = 1e-6; // justnorm's: added to the sum of squares, not its mean + +/// per_loop: a cache per attention step; last_loop: one per attention sublayer, each loop overwriting the chunk's keys. +pub const Mode = enum { per_loop, last_loop }; + +pub const Options = struct { + /// dory.py's bf16 points (progress log D2) and the fp32 scale vectors (D3); false: float64 throughout. + round: bool = false, + mode: Mode = .per_loop, + /// h after each step, [steps][rows][hidden], when set. + taps: ?[]f64 = null, +}; + +/// A projection [n, k], row-major. +pub const Proj = struct { w: []f64, n: usize, k: usize }; +pub const Attn = struct { alpha: []f64, q: Proj, k: Proj, v: Proj, g: Proj, o: Proj, sqk: []f64, s_gate: []f64 }; +pub const Mlp = struct { alpha: []f64, gate: Proj, up: Proj, down: Proj, suv_gate: []f64, suv_up: []f64 }; +pub const Sub = union(enum) { attn: Attn, mlp: Mlp }; + +pub const Weights = struct { + arena: std.heap.ArenaAllocator, + embed: []f64, // [V, D] + head: []f64, // [V, D] + logit_scale: []f64, // [V] + subs: []Sub, + + pub fn deinit(w: *Weights) void { + w.arena.deinit(); + } +}; + +/// splitmix64: the fixture weights' generator (the dev-time oracle draws the same stream). +pub const Gen = struct { + s: u64, + + pub fn next(g: *Gen) u64 { + g.s +%= 0x9E3779B97F4A7C15; + var z = g.s; + z = (z ^ (z >> 30)) *% 0xBF58476D1CE4E5B9; + z = (z ^ (z >> 27)) *% 0x94D049BB133111EB; + return z ^ (z >> 31); + } + + /// A code in [-128, 127], from the draw's top byte. + pub fn code(g: *Gen) i64 { + return @as(i64, @intCast(g.next() >> 56)) - 128; + } +}; + +/// n signed values code / 128 * 2^e, or with `positive` (128 + |code|) / 128 * 2^e: all exact in bf16. +fn fill(a: std.mem.Allocator, g: *Gen, n: usize, e: i32, positive: bool) ![]f64 { + const v = try a.alloc(f64, n); + for (v) |*x| { + const m = g.code(); + const q: f64 = @floatFromInt(if (positive) 128 + @as(i64, @intCast(@abs(m))) else m); + x.* = std.math.ldexp(q / 128, e); + } + return v; +} + +fn proj(a: std.mem.Allocator, g: *Gen, n: usize, k: usize) !Proj { + const e: i32 = -@as(i32, std.math.log2_int(usize, k)) >> 1; // about 1 / sqrt(k), a power of two + return .{ .w = try fill(a, g, n * k, e, false), .n = n, .k = k }; +} + +/// Random weights for `c` in a fixed order (embeddings, each sublayer's tensors as stored, head, logit_scale). +pub fn random(gpa: std.mem.Allocator, c: Config, seed: u64) !Weights { + var w: Weights = .{ .arena = .init(gpa), .embed = &.{}, .head = &.{}, .logit_scale = &.{}, .subs = &.{} }; + errdefer w.arena.deinit(); + const a = w.arena.allocator(); + var g: Gen = .{ .s = seed }; + const d = c.hidden; + w.embed = try fill(a, &g, c.vocab * d, 0, false); + w.subs = try a.alloc(Sub, c.sublayers); + for (w.subs, 0..) |*s, i| { + if (c.kinds[i] == .mlp) { + const alpha = try fill(a, &g, d, -2, false); + const gate = try proj(a, &g, c.intermediate, d); + const up = try proj(a, &g, c.intermediate, d); + const down = try proj(a, &g, d, c.intermediate); + s.* = .{ .mlp = .{ .alpha = alpha, .gate = gate, .up = up, .down = down, .suv_gate = try fill(a, &g, c.intermediate, -3, true), .suv_up = try fill(a, &g, c.intermediate, -3, true) } }; + } else { + const alpha = try fill(a, &g, d, -2, false); + const q = try proj(a, &g, c.qDim(), d); + const k = try proj(a, &g, c.kvDim(), d); + const v = try proj(a, &g, c.kvDim(), d); + const gp = try proj(a, &g, c.qDim(), d); + const o = try proj(a, &g, d, c.qDim()); + s.* = .{ .attn = .{ .alpha = alpha, .q = q, .k = k, .v = v, .g = gp, .o = o, .sqk = try fill(a, &g, c.qDim(), -8, true), .s_gate = try fill(a, &g, c.qDim(), -7, true) } }; + } + } + const e: i32 = -@as(i32, std.math.log2_int(usize, d)) >> 1; + w.head = try fill(a, &g, c.vocab * d, e, false); + w.logit_scale = try fill(a, &g, c.vocab, -6, true); + return w; +} + +/// The scale vectors dory.py builds from the stored ones: in float64, or as torch's fp32 steps (D3) with `round`. +const Scales = struct { + round: bool, + f: config.Factors, + c: Config, + + fn f32of(v: f64) f64 { + return @as(f32, @floatCast(v)); + } + + /// |alpha * alpha_init / std| + fn lr(s: Scales, alpha: f64) f64 { + return if (s.round) @abs(f32of(alpha * s.f.alpha)) else @abs(alpha * (s.c.alpha_init / s.c.std)); + } + + /// (v * factor)^2, factor sqrt(head_dim) / std (sqk) or sqrt(gate_init) / std (s_gate). + fn square(s: Scales, v: f64, gate: bool) f64 { + const d: f64 = @floatFromInt(s.c.head_dim); + if (s.round) { + const t = f32of(v * (if (gate) s.f.gate else s.f.sqk)); + return f32of(t * t); + } + const t = v * (if (gate) @sqrt(s.c.gate_init) else @sqrt(d)) / s.c.std; + return t * t; + } + + /// suv * sqrt(hidden) + fn suv(s: Scales, v: f64) f64 { + return if (s.round) f32of(v * s.f.suv) else v * @sqrt(@as(f64, @floatFromInt(s.c.hidden))); + } + + /// logit_scale / std (torch: times the fp32 reciprocal) + fn logit(s: Scales, v: f64) f64 { + return if (s.round) f32of(v * s.f.logit) else v / s.c.std; + } + + fn attn(s: Scales) f64 { + return if (s.round) s.f.attn else 1 / @sqrt(@as(f64, @floatFromInt(s.c.head_dim))); + } +}; + +/// A stream's KV caches, keys and values [positions][kv_heads][head_dim] a cache. +pub const State = struct { + arena: std.heap.ArenaAllocator, + c: Config, + sched: schedule.Schedule, + mode: Mode, + k: [][]f64, + v: [][]f64, + positions: usize, + len: usize = 0, // positions written + + pub fn init(gpa: std.mem.Allocator, c: Config, positions: usize, mode: Mode) !State { + var st: State = .{ .arena = .init(gpa), .c = c, .sched = schedule.build(c), .mode = mode, .k = &.{}, .v = &.{}, .positions = positions }; + errdefer st.arena.deinit(); + const a = st.arena.allocator(); + st.k = try a.alloc([]f64, st.sched.caches); + st.v = try a.alloc([]f64, st.sched.caches); + for (st.k, st.v) |*k, *v| { + k.* = try a.alloc(f64, positions * c.kvDim()); + v.* = try a.alloc(f64, positions * c.kvDim()); + @memset(k.*, 0); + @memset(v.*, 0); + } + return st; + } + + pub fn deinit(st: *State) void { + st.arena.deinit(); + } + + /// The cache a step reads and writes: its own (per_loop), or loop 0's of the same sublayer (last_loop). + fn cacheOf(st: *const State, step: schedule.Step) usize { + if (st.mode == .per_loop or step.loop == 0) return step.cache.?; + for (st.sched.slice()) |s| if (s.sub == step.sub and s.loop == 0) return s.cache.?; + unreachable; + } +}; + +fn bfIf(round: bool, v: f64) f64 { + return if (round) bfr.bf16(v) else v; +} + +/// y = x / sqrt(sum(x^2) + eps), in place. +fn l2(x: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const inv = 1 / @sqrt(ss + eps); + for (x) |*v| v.* *= inv; +} + +fn roundAll(round: bool, x: []f64) void { + if (round) for (x) |*v| { + v.* = bfr.bf16(v.*); + }; +} + +/// out [rows, n] = x [rows, k] times w transposed. +fn matmul(x: []const f64, p: Proj, out: []f64, rows: usize, round: bool) void { + for (0..rows) |r| for (0..p.n) |n| { + var s: f64 = 0; + for (x[r * p.k ..][0..p.k], p.w[n * p.k ..][0..p.k]) |a, b| s += a * b; + out[r * p.n + n] = bfIf(round, s); + }; +} + +/// Rotate-half RoPE of one head at position p. +fn rope(x: []f64, p: usize, theta: f64) void { + const half = x.len / 2; + for (0..half) |j| { + const inv = 1 / std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(x.len))); + const ang = @as(f64, @floatFromInt(p)) * inv; + const cs = @cos(ang); + const sn = @sin(ang); + const a = x[j]; + const b = x[j + half]; + x[j] = a * cs - b * sn; + x[j + half] = b * cs + a * sn; + } +} + +/// a = l2(res), b = l2(out), res = l2(a + lr (b - a)) a row at a time. +fn slerp(res: []f64, out: []f64, alpha: []const f64, sc: Scales, rows: usize, d: usize) void { + for (0..rows) |r| { + const a = res[r * d ..][0..d]; + const b = out[r * d ..][0..d]; + l2(a); + roundAll(sc.round, a); + l2(b); + roundAll(sc.round, b); + for (a, b, alpha) |*x, y, al| x.* = x.* + sc.lr(al) * (y - x.*); + l2(a); + roundAll(sc.round, a); + } +} + +const Ctx = struct { a: std.mem.Allocator, c: Config, sc: Scales, rows: usize, p0: usize, round: bool }; + +fn attention(x: Ctx, st: *State, step: schedule.Step, w: Attn, h: []const f64, out: []f64) !void { + const c = x.c; + const rows = x.rows; + const dh = c.head_dim; + const qd = c.qDim(); + const kvd = c.kvDim(); + const a = x.a; + const q = try a.alloc(f64, rows * qd); + const k = try a.alloc(f64, rows * kvd); + const v = try a.alloc(f64, rows * kvd); + const g = try a.alloc(f64, rows * qd); + const o = try a.alloc(f64, rows * qd); + matmul(h, w.q, q, rows, x.round); + matmul(h, w.k, k, rows, x.round); + matmul(h, w.v, v, rows, x.round); + matmul(h, w.g, g, rows, x.round); + const theta = if (step.kind == .full) c.theta_full else c.theta_sliding; + const cache = st.cacheOf(step); + for (0..rows) |r| { + const p = x.p0 + r; + for (0..c.heads) |hd| { + const head = q[r * qd + hd * dh ..][0..dh]; + rope(head, p, theta); + roundAll(x.round, head); + l2(head); + roundAll(x.round, head); + for (head, 0..) |*e, i| e.* = bfIf(x.round, x.sc.square(w.sqk[hd * dh + i], false) * e.*); + } + for (0..c.kv_heads) |hd| { + const head = k[r * kvd + hd * dh ..][0..dh]; + rope(head, p, theta); + roundAll(x.round, head); + l2(head); + roundAll(x.round, head); + } + @memcpy(st.k[cache][p * kvd ..][0..kvd], k[r * kvd ..][0..kvd]); + @memcpy(st.v[cache][p * kvd ..][0..kvd], v[r * kvd ..][0..kvd]); + } + const scores = try a.alloc(f64, x.p0 + rows); + const group = c.heads / c.kv_heads; + for (0..rows) |r| { + const p = x.p0 + r; + const lo = if (step.kind == .sliding and p + 1 > c.window) p + 1 - c.window else 0; + for (0..c.heads) |hd| { + const qh = q[r * qd + hd * dh ..][0..dh]; + const kh = hd / group; + var most = -std.math.inf(f64); + for (lo..p + 1) |j| { + var s: f64 = 0; + for (qh, st.k[cache][j * kvd + kh * dh ..][0..dh]) |e, f| s += e * f; + scores[j] = s * x.sc.attn(); + most = @max(most, scores[j]); + } + var total: f64 = 0; + for (lo..p + 1) |j| { + scores[j] = @exp(scores[j] - most); + total += scores[j]; + } + const oh = o[r * qd + hd * dh ..][0..dh]; + @memset(oh, 0); + for (lo..p + 1) |j| for (oh, st.v[cache][j * kvd + kh * dh ..][0..dh]) |*e, f| { + e.* += scores[j] / total * f; + }; + roundAll(x.round, oh); + } + const orow = o[r * qd ..][0..qd]; + for (orow, g[r * qd ..][0..qd], 0..) |*e, gv, i| e.* = bfIf(x.round, e.* * sigmoid(gv * x.sc.square(w.s_gate[i], true))); + l2(orow); + roundAll(x.round, orow); + } + matmul(o, w.o, out, rows, x.round); +} + +fn sigmoid(v: f64) f64 { + return 1 / (1 + @exp(-v)); +} + +fn mlp(x: Ctx, w: Mlp, h: []const f64, out: []f64) !void { + const f = x.c.intermediate; + const gt = try x.a.alloc(f64, x.rows * f); + const up = try x.a.alloc(f64, x.rows * f); + matmul(h, w.gate, gt, x.rows, x.round); + matmul(h, w.up, up, x.rows, x.round); + for (0..x.rows) |r| { + const m = gt[r * f ..][0..f]; + for (m, up[r * f ..][0..f], w.suv_gate, w.suv_up) |*gv, uv, sg, su| { + const gs = bfIf(x.round, gv.* * x.sc.suv(sg)); + const us = bfIf(x.round, uv * x.sc.suv(su)); + gv.* = bfIf(x.round, gs * sigmoid(gs) * us); + } + l2(m); + roundAll(x.round, m); + } + matmul(gt, w.down, out, x.rows, x.round); +} + +/// Logits [tokens.len, V] of a chunk at positions p0.. after the stream's earlier chunks (st.len must be p0). +pub fn forward(gpa: std.mem.Allocator, w: *const Weights, st: *State, tokens: []const u32, p0: usize, opt: Options) ![]f64 { + const c = st.c; + if (p0 != st.len or p0 + tokens.len > st.positions or tokens.len == 0) return error.BadChunk; + if (opt.mode != st.mode) return error.BadChunk; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const rows = tokens.len; + const d = c.hidden; + const sc: Scales = .{ .round = opt.round, .f = c.factors(), .c = c }; + const x: Ctx = .{ .a = a, .c = c, .sc = sc, .rows = rows, .p0 = p0, .round = opt.round }; + var h = try a.alloc(f64, rows * d); + for (tokens, 0..) |t, r| @memcpy(h[r * d ..][0..d], w.embed[t * d ..][0..d]); + const input_h = try a.alloc(f64, rows * d); + const state = try a.alloc(f64, rows * d); + const out = try a.alloc(f64, rows * d); + for (st.sched.slice(), 0..) |step, i| { + if (step.enter) { + if (step.loop == 0) { + @memcpy(input_h, h); + @memcpy(state, h); + } + for (h, state, input_h) |*e, s, in| e.* = bfIf(opt.round, s + in); + for (0..rows) |r| l2(h[r * d ..][0..d]); + roundAll(opt.round, h); + } + const alpha = switch (w.subs[step.sub]) { + .attn => |aw| blk: { + try attention(x, st, step, aw, h, out); + break :blk aw.alpha; + }, + .mlp => |mw| blk: { + try mlp(x, mw, h, out); + break :blk mw.alpha; + }, + }; + slerp(h, out, alpha, sc, rows, d); + if (step.leave) { + for (0..rows) |r| l2(h[r * d ..][0..d]); + roundAll(opt.round, h); + @memcpy(state, h); + } + if (opt.taps) |t| @memcpy(t[i * rows * d ..][0 .. rows * d], h); + } + st.len = p0 + rows; + const logits = try gpa.alloc(f64, rows * c.vocab); + for (0..rows) |r| for (0..c.vocab) |n| { + var s: f64 = 0; + for (h[r * d ..][0..d], w.head[n * d ..][0..d]) |e, f| s += e * f; + logits[r * c.vocab + n] = bfIf(opt.round, s) * sc.logit(w.logit_scale[n]); + }; + return logits; +} + +test { + _ = @import("ref_test.zig"); +} diff --git a/zig/src/families/dory/ref_fixture.zig b/zig/src/families/dory/ref_fixture.zig new file mode 100644 index 0000000000..99ea56f3d0 --- /dev/null +++ b/zig/src/families/dory/ref_fixture.zig @@ -0,0 +1,83 @@ +//! The dev-time oracle's numbers: vLLM #60149's dory.py (ce9dd0ed) in float64 on ref.random(names.tiny, seed). +pub const seed: u64 = 0xd0e1; +pub const Case = struct { loops: usize, rows: []const usize, logits: []const f64 }; +pub const cases = [_]Case{ + .{ .loops = 1, .rows = &.{ 0, 3, 11 }, .logits = &.{ + 0.09999391438248895, 0.13154584970370847, 0.10293708025163188, 0.2924127495741501, + 0.186555220886937, -0.06122914115470976, -0.04578685432031266, 0.16611501082010816, + 0.02830426137354789, -0.01248726366890309, 0.0022074780770686667, 0.03405137760517988, + -0.04901280581499187, -0.018756908645607223, 0.04753841420291959, -0.05218028445571712, + 0.10482749986251386, -0.0023096851346042237, 0.2740058531742164, 0.17250464171127391, + 0.01879628660203299, 0.0658653895733172, -0.03483957627425773, -0.06537271332220236, + -0.07671930810664027, 0.00802772984711647, -0.0015512327419252522, -0.029463142911226307, + 0.13230624348390496, -0.12313809789527817, 0.06526989244795121, 0.00881507890132587, + -0.2521379862087677, -0.02730344949827025, -0.07531001427634387, 0.01623004244057416, + -0.049129646366406, 0.05310989100283782, -0.005846538119214357, -0.04633713661061329, + -0.04568309358590196, 0.1456690670975681, -0.06715203332943064, -0.08628636278500487, + -0.13545138362324205, -0.018914618160367862, -0.020264665932942134, 0.1468698549181028, + -0.038580489837277514, 0.1547226689880456, -0.015110359030007096, 0.19290744033892324, + 0.002954960465123267, -0.03711934201726611, 0.006143138340454329, -0.09642782191398926, + -0.05975688961608956, -0.040829845354240765, -0.11481275513012186, 0.012610531981103878, + 0.006864936562921483, -0.08877049173256914, -0.04098930630195836, 0.019074984842730198, + -0.007349935929316478, -0.06023414427824233, -0.26287727582615183, 0.07618587383574708, + 0.046714892954293374, -0.05168761152718291, -0.04790492045200638, -0.17418614940458532, + -0.04131429260657254, 0.024141206204652828, -0.03925507264073149, 0.01834316045751918, + -0.042110303582961, -0.05544381409397682, 0.061684008516608047, 0.11188651702105945, + -0.03453732046128847, 0.07380828357594696, -0.02484591881166964, 0.02523192115067423, + 0.005622466601503181, -0.06079896483723641, -0.16067383640885266, -0.0007829945197161271, + 0.024790916622280213, 0.06697756922581959, -0.12208238860499403, -0.03247203787505084, + -0.067139447952211, -0.03535001630097715, 0.06214627214559448, -0.02946094723555471, + } }, + .{ .loops = 2, .rows = &.{ 0, 3, 11 }, .logits = &.{ + 0.11367021833912412, 0.11904307496223783, 0.06052550512427518, 0.25683556956523684, + 0.18053297909168894, -0.002724552527554604, -0.016307219079224534, 0.14664208934210787, + -0.026551610198372856, -0.041260053320581774, 0.015440766628291245, 0.09856245714193626, + -0.11680997954712521, -0.0007485500074362977, 0.023790011799636003, -0.11514654243679048, + 0.11528949626848585, -0.006254425602250209, 0.21477704959704272, 0.10707467664750457, + -0.047225415786799445, 0.02944779705458566, -0.028373476027534515, -0.032134952125101104, + -0.08479732852878837, 0.04931125814098024, -0.015559930055597335, -0.12137204763694615, + 0.10277440959572678, 0.009340732321087624, 0.09377850624667129, -0.06856405439786145, + -0.10819250397871807, 0.00015706074595524226, -0.03848520288437991, 0.0797621602414385, + 0.005383818853117581, 0.00240608689027137, 0.01703215232693508, 0.052970479130456945, + -0.11260000794076287, 0.13232085554674697, -0.005670043200110942, -0.14796747182956146, + -0.16930036600404527, 0.05185176224884835, 0.030582298013851268, 0.10268649267397312, + -0.023131547987350336, 0.17239289003070368, 0.019695278802568602, 0.19574557916757085, + -0.0104453789393829, -0.043368167305135034, -0.025781349458852153, -0.09772993100541169, + -0.06955261514738363, 0.058670371897001444, -0.11021878624907236, -0.04460424671060183, + -0.03358477822913512, -0.12975573396501197, -0.034774128196751415, -0.14033731002609118, + 0.1389556271794824, 0.06617610638917593, -0.149199115704294, 0.1611139950551834, + 0.2941915932985296, -0.1369923452859219, -0.08792547015218337, -0.01050599391037401, + -0.10365506935710184, 0.09447513158371583, -0.11336743007175912, -0.08192089827056208, + -0.011574748134543716, -0.05671023333499292, 0.08270983128005542, -0.10451458812195485, + -0.136779595486605, 0.06639237325736112, -0.04441351182700274, -0.0015760865807046744, + 0.023128347237608465, -0.04276904125859905, -0.12189683760558995, -0.049605697242238844, + -0.04930468884477375, 0.08079265267032978, -0.054115124806284394, -0.01619773377119024, + 0.019912151526570412, 0.03530385842237603, 0.13215232956223433, -0.037408192336297705, + } }, + .{ .loops = 4, .rows = &.{ 0, 3, 11 }, .logits = &.{ + 0.11212194428962761, 0.10851499316104132, 0.01275399656136828, 0.2420964383242402, + 0.1864884730653102, -0.028986030775549603, -0.0076726278499005636, 0.1649978354194932, + -0.01755188987517725, -0.042110423697741546, 0.017807313786058335, 0.08831788365238571, + -0.1599036204322905, 0.011557085589693325, 0.013107520557010153, -0.16972236697446064, + 0.09497175333379831, 0.01340027229987714, 0.13700457857699125, 0.04715213183654023, + -0.05717525285249378, 0.012989165675151801, -0.016714747767984792, -0.018411757035936525, + -0.07792956969964664, 0.07360559717273536, 0.0004438953330003915, -0.16836586815044147, + 0.13325083119218994, 0.09940711938945243, 0.10580453713434644, -0.07851908264324971, + -0.07923877529288886, 0.01662814350334148, -0.03429590118929948, 0.040241938971393235, + -0.032935991501544935, 0.09299656554266947, -0.006083033146947442, 0.011899552439768648, + -0.08255337913842666, 0.10233218300937907, -0.05314431241833622, -0.12914771260525557, + -0.12429416071662822, 0.03678160602251373, 0.09239845194189893, 0.030167449558176632, + 0.019247933472121625, 0.12316228591833886, -0.04975849321085094, 0.11823824279031896, + 0.00046993763526103217, -0.11165032285052635, -0.04229819121988338, -0.13996996850170793, + -0.0714666402045578, 0.16004529933716163, -0.11392817140912552, 0.019029283916361918, + -0.08946683976307494, -0.22341243293982635, -0.06697495814792473, -0.13740463514311277, + 0.08006411615424971, 0.027742428849581943, -0.14444469290597625, 0.23919148886147734, + 0.28405738233836614, -0.16943044336738783, -0.011068246483869728, -0.04309936793397592, + -0.14443187751361555, 0.039329780062217344, -0.09414049779612017, -0.042880558560585966, + -0.03021044886862581, -0.07285212218261204, 0.05849430725854382, -0.09724338384258853, + -0.08140950105266777, 0.11053304123362388, 0.06749812472451273, 0.050568559415255794, + 0.0204263685169831, -0.0549598728265329, -0.10960556716039625, -0.0175119735840323, + -0.08755555132489978, 0.029743908467058794, -0.078874439459035, -0.03735475526609189, + 0.04376874076793922, 0.0668630032772935, -0.041097133582121186, -0.028637869394776465, + } }, +}; diff --git a/zig/src/families/dory/ref_test.zig b/zig/src/families/dory/ref_test.zig new file mode 100644 index 0000000000..2de11d8d90 --- /dev/null +++ b/zig/src/families/dory/ref_test.zig @@ -0,0 +1,118 @@ +//! The float64 forward's checks: per_loop is chunk-invariant to the bit, last_loop is not, dory.py's numbers match. +const std = @import("std"); +const ref = @import("ref.zig"); +const names = @import("names.zig"); +const fixture = @import("ref_fixture.zig"); + +const testing = std.testing; + +const prompt = [_]u32{ 1, 7, 30, 2, 19, 4, 4, 25, 11, 0, 31, 8 }; + +/// The prompt's logits, fed as chunks of the given sizes (they sum to the prompt's length). +fn chunked(w: *const ref.Weights, loops: usize, sizes: []const usize, opt: ref.Options) ![]f64 { + const c = try names.tiny.withLoops(loops); + var st = try ref.State.init(testing.allocator, c, prompt.len, opt.mode); + defer st.deinit(); + const all = try testing.allocator.alloc(f64, prompt.len * c.vocab); + errdefer testing.allocator.free(all); + var at: usize = 0; + for (sizes) |n| { + const l = try ref.forward(testing.allocator, w, &st, prompt[at..][0..n], at, opt); + defer testing.allocator.free(l); + @memcpy(all[at * c.vocab ..][0 .. n * c.vocab], l); + at += n; + } + try testing.expectEqual(prompt.len, at); + return all; +} + +const cuts = [_][]const usize{ &.{12}, &.{ 5, 4, 3 }, &.{ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 }, &.{ 3, 9 } }; + +fn sameBits(a: []const f64, b: []const f64) bool { + return std.mem.eql(u8, std.mem.sliceAsBytes(a), std.mem.sliceAsBytes(b)); +} + +test "per_loop: every chunking and one-token decode give the same logits to the bit, 1 to 4 loops, both precisions" { + var w = try ref.random(testing.allocator, names.tiny, fixture.seed); + defer w.deinit(); + for (1..5) |loops| for ([_]bool{ false, true }) |round| { + const whole = try chunked(&w, loops, cuts[0], .{ .round = round }); + defer testing.allocator.free(whole); + for (whole) |v| try testing.expect(std.math.isFinite(v)); + for (cuts[1..]) |cut| { + const other = try chunked(&w, loops, cut, .{ .round = round }); + defer testing.allocator.free(other); + try testing.expect(sameBits(whole, other)); + } + }; +} + +test "last_loop: equal to per_loop at one loop, and at two loops its logits depend on where the prompt is cut" { + var w = try ref.random(testing.allocator, names.tiny, fixture.seed); + defer w.deinit(); + const one_per = try chunked(&w, 1, cuts[1], .{}); + defer testing.allocator.free(one_per); + const one_last = try chunked(&w, 1, cuts[1], .{ .mode = .last_loop }); + defer testing.allocator.free(one_last); + try testing.expect(sameBits(one_per, one_last)); + const whole = try chunked(&w, 2, cuts[0], .{ .mode = .last_loop }); + defer testing.allocator.free(whole); + const split = try chunked(&w, 2, cuts[1], .{ .mode = .last_loop }); + defer testing.allocator.free(split); + const v = names.tiny.vocab; + // the first chunk's rows see only their own chunk, as in one chunk; the later chunks' rows read last-loop keys + try testing.expect(sameBits(whole[0 .. 5 * v], split[0 .. 5 * v])); + try testing.expect(!sameBits(whole[5 * v ..], split[5 * v ..])); + const per = try chunked(&w, 2, cuts[0], .{}); + defer testing.allocator.free(per); + try testing.expect(sameBits(whole, per)); // in one chunk every key is the current loop's in both modes +} + +test "float64 against dory.py run in float64 (the dev-time oracle): every logit within 1e-12" { + var w = try ref.random(testing.allocator, names.tiny, fixture.seed); + defer w.deinit(); + for (fixture.cases) |case| { + const l = try chunked(&w, case.loops, cuts[0], .{}); + defer testing.allocator.free(l); + for (case.rows, 0..) |row, j| { + const got = l[row * names.tiny.vocab ..][0..names.tiny.vocab]; + for (got, case.logits[j * names.tiny.vocab ..][0..names.tiny.vocab]) |g, want| { + if (@abs(g - want) > 1e-12 * @max(1, @abs(want))) { + std.debug.print("{d} loops, row {d}: {e} against dory.py's {e}\n", .{ case.loops, row, g, want }); + return error.TestUnexpectedResult; + } + } + } + } +} + +test "the bf16 rounding points stay near float64: the same argmax on most rows, logits within a few hundredths" { + var w = try ref.random(testing.allocator, names.tiny, fixture.seed); + defer w.deinit(); + for (1..5) |loops| { + const exact = try chunked(&w, loops, cuts[0], .{}); + defer testing.allocator.free(exact); + const rounded = try chunked(&w, loops, cuts[0], .{ .round = true }); + defer testing.allocator.free(rounded); + var most: f64 = 0; + for (exact, rounded) |e, r| most = @max(most, @abs(e - r)); + try testing.expect(most > 0 and most < 0.25); + } +} + +test "taps: every step leaves each row on the unit sphere" { + var w = try ref.random(testing.allocator, names.tiny, fixture.seed); + defer w.deinit(); + const c = names.tiny; + var st = try ref.State.init(testing.allocator, c, prompt.len, .per_loop); + defer st.deinit(); + const taps = try testing.allocator.alloc(f64, st.sched.len * prompt.len * c.hidden); + defer testing.allocator.free(taps); + const l = try ref.forward(testing.allocator, &w, &st, &prompt, 0, .{ .taps = taps }); + defer testing.allocator.free(l); + for (0..st.sched.len) |i| for (0..prompt.len) |r| { + var ss: f64 = 0; + for (taps[(i * prompt.len + r) * c.hidden ..][0..c.hidden]) |v| ss += v * v; + try testing.expectApproxEqAbs(@as(f64, 1), ss, 1e-5); // eps on the sum: short of 1 by about eps / ss + }; +} From 9855f35688e69f6d003580fb2548c4fa9da13950 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 17:01:41 +0200 Subject: [PATCH 8/9] dory: MXFP8 and BF16 checkpoint maps (host only) nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-MXFP8 and -BF16 next to the NVFP4 map. - config.zig: Quant gains mxfp8. quantOf takes compressed-tensors' mxfp8-quantized (8-bit float codes, group 32) besides ModelOpt's W4A16_NVFP4 and plain bf16, and refuses other compressed-tensors groups by name. - names.zig: an MXFP8 projection is its E4M3 codes [n, k] and UE8M0 scales u8 [n, k/32]; the tiny checkpoint test runs all three quants, and the pinned tests (TF_DORY_DIR, TF_DORY_MXFP8_DIR, TF_DORY_BF16_DIR) share one body. - memory.zig: the published MXFP8 and BF16 files' tensor data (4,750,663,680 and 7,950,458,880 bytes) equal weightBytes; parameters are the same for every quant. --- zig/src/families/dory/config.zig | 94 +++++++++++++++++++++++++++++--- zig/src/families/dory/memory.zig | 19 ++++++- zig/src/families/dory/names.zig | 41 +++++++++++--- 3 files changed, 134 insertions(+), 20 deletions(-) diff --git a/zig/src/families/dory/config.zig b/zig/src/families/dory/config.zig index a096a0fc1f..efbd2dd7f4 100644 --- a/zig/src/families/dory/config.zig +++ b/zig/src/families/dory/config.zig @@ -10,8 +10,8 @@ pub const max_width = 8192; // the rows kernels' widest row: 4 sixteen-byte load /// A sublayer: attention over a sliding window, attention over every key before, or a SwiGLU MLP. pub const Kind = enum { sliding, full, mlp }; -/// How the projections are stored: ModelOpt NVFP4 W4A16 (group 16) or plain bf16. -pub const Quant = enum { nvfp4, bf16 }; +/// How the projections are stored: NVFP4 W4A16 at group 16, compressed-tensors MXFP8 at group 32, or plain bf16. +pub const Quant = enum { nvfp4, mxfp8, bf16 }; /// Why a check refused the checkpoint, kept for the caller's log line (tests read it instead). pub const Why = struct { @@ -210,7 +210,7 @@ fn rope(o: std.json.ObjectMap, kind: []const u8, why: *Why) !f64 { return number(p.object, "rope_theta", why); } -/// The quantization config.json declares: none (bf16) or ModelOpt's W4A16_NVFP4 at group 16. +/// The quantization config.json declares: none (bf16), ModelOpt's W4A16_NVFP4, or compressed-tensors' mxfp8. fn quantOf(o: std.json.ObjectMap, why: *Why) !Quant { const q = o.get("quantization_config") orelse return .bf16; if (q != .object) return error.BadConfig; @@ -218,11 +218,16 @@ fn quantOf(o: std.json.ObjectMap, why: *Why) !Quant { const algo = q.object.get("quant_algo"); const modelopt = method != null and method.? == .string and std.mem.eql(u8, method.?.string, "modelopt"); const w4a16 = algo != null and algo.? == .string and std.mem.eql(u8, algo.?.string, "W4A16_NVFP4"); - if (!modelopt or !w4a16) { - why.set("Dory on CUDA reads ModelOpt's W4A16_NVFP4 checkpoint or plain bf16; this one is quantized otherwise", .{}); - return error.UnsupportedQuantization; - } - if (q.object.get("config_groups")) |groups| if (groups == .object) { + if (modelopt and w4a16) return nvfp4Groups(q.object, why); + const mxfp8 = method != null and method.? == .string and std.mem.eql(u8, method.?.string, "compressed-tensors"); + if (mxfp8) return mxfp8Groups(q.object, why); + why.set("Dory on CUDA reads ModelOpt's W4A16_NVFP4, compressed-tensors' mxfp8-quantized or plain bf16; this one is quantized otherwise", .{}); + return error.UnsupportedQuantization; +} + +/// ModelOpt's groups: 16 weights a scale. +fn nvfp4Groups(q: std.json.ObjectMap, why: *Why) !Quant { + if (q.get("config_groups")) |groups| if (groups == .object) { var it = groups.object.iterator(); while (it.next()) |g| { const w = if (g.value_ptr.* == .object) g.value_ptr.object.get("weights") else null; @@ -236,6 +241,38 @@ fn quantOf(o: std.json.ObjectMap, why: *Why) !Quant { return .nvfp4; } +/// compressed-tensors' groups: mxfp8-quantized at 8 bits, float, 32 weights a scale. +fn mxfp8Groups(q: std.json.ObjectMap, why: *Why) !Quant { + const groups = q.get("config_groups") orelse { + why.set("mxfp8's config.json has no config_groups", .{}); + return error.BadConfig; + }; + if (groups != .object) return error.BadConfig; + var it = groups.object.iterator(); + while (it.next()) |g| if (g.value_ptr.* == .object) { + const w = g.value_ptr.object.get("weights"); + const ok = w != null and w.? == .object and intField(w.?.object, "num_bits", 8) and intField(w.?.object, "group_size", 32) and + stringField(w.?.object, "type", "float") and stringField(g.value_ptr.object, "format", "mxfp8-quantized"); + if (!ok) { + why.set("mxfp8 group {s} is not mxfp8-quantized at 8-bit float, group 32", .{g.key_ptr.*}); + return error.UnsupportedQuantization; + } + }; + return .mxfp8; +} + +/// The object's `key` as the integer `want`. +fn intField(o: std.json.ObjectMap, key: []const u8, want: i64) bool { + const v = o.get(key) orelse return false; + return v == .integer and v.integer == want; +} + +/// The object's `key` as the string `want`. +fn stringField(o: std.json.ObjectMap, key: []const u8, want: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .string and std.mem.eql(u8, v.string, want); +} + /// config.json's text; a refusal's reason goes to `why`. pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config { var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); @@ -475,3 +512,44 @@ test "bf16 checkpoints parse without a quantization_config; shapes outside the k odd.kv_heads = 3; try testing.expectError(error.UnsupportedShapes, checkShapes(odd, &why)); } + +/// The published MXFP8 quantization_config (targets shortened to "Linear"), spliced in place of NVFP4's. +const mxfp8_config = + \\{"config_groups": {"dory_matmul_weights": {"format": "mxfp8-quantized", + \\ "input_activations": null, + \\ "output_activations": null, + \\ "targets": ["Linear"], + \\ "weights": {"actorder": null, "block_structure": null, "dynamic": false, "group_size": 32, "num_bits": 8, + \\ "observer": "mse", "scale_dtype": "torch.uint8", "strategy": "group", "symmetric": true, "type": "float", + \\ "zp_dtype": null}}}, + \\ "format": "mxfp8-quantized", "global_compression_ratio": null, + \\ "ignore": ["backbone.embeddings", "lm_head"], + \\ "kv_cache_scheme": null, + \\ "quant_method": "compressed-tensors", + \\ "quantization_status": "compressed"} +; + +fn mxfp8Text(a: std.mem.Allocator) ![]u8 { + const t = try pinnedText(a); + defer a.free(t); + const at = std.mem.indexOf(u8, t, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, t, "\"head_dim\"").?; + return std.mem.concat(a, u8, &.{ t[0..at], "\"quantization_config\": ", mxfp8_config, ", ", t[end..] }); +} + +test "the MXFP8 quantization_config parses to mxfp8; group 16, 4-bit and other compressed-tensors refuse" { + var why: Why = .{}; + const mx = try mxfp8Text(testing.allocator); + defer testing.allocator.free(mx); + const c = try parse(testing.allocator, mx, &why); + try testing.expectEqual(Quant.mxfp8, c.quant); + const gs16 = try std.mem.replaceOwned(u8, testing.allocator, mx, "\"group_size\": 32", "\"group_size\": 16"); + defer testing.allocator.free(gs16); + try refused(gs16, error.UnsupportedQuantization, "mxfp8"); + const bits4 = try std.mem.replaceOwned(u8, testing.allocator, mx, "\"num_bits\": 8", "\"num_bits\": 4"); + defer testing.allocator.free(bits4); + try refused(bits4, error.UnsupportedQuantization, "mxfp8"); + const int4 = try std.mem.replaceOwned(u8, testing.allocator, mx, "mxfp8-quantized", "int4-quantized"); + defer testing.allocator.free(int4); + try refused(int4, error.UnsupportedQuantization, "mxfp8"); +} diff --git a/zig/src/families/dory/memory.zig b/zig/src/families/dory/memory.zig index c038654e3e..c74b3c955c 100644 --- a/zig/src/families/dory/memory.zig +++ b/zig/src/families/dory/memory.zig @@ -5,10 +5,11 @@ const schedule = @import("schedule.zig"); const Config = config.Config; -/// A projection's bytes as stored: NVFP4 codes, e4m3 scales a 16 and the fp32 scale, or bf16. +/// A projection's bytes as stored: NVFP4 codes, scales and the fp32 scale, MXFP8 codes and u8 scales, or bf16. pub fn projBytes(c: Config, n: usize, k: usize) usize { return switch (c.quant) { .nvfp4 => n * k / 2 + n * k / 16 + 4, + .mxfp8 => n * k + n * k / 32, .bf16 => n * k * 2, }; } @@ -52,11 +53,23 @@ fn pinned() !Config { return config.parse(testing.allocator, text, &why); } -test "3.97B parameters; the NVFP4 tensors' bytes equal the pinned file's data (3,202,498,032 bytes less its header)" { +test "3.97B parameters; the tensors' bytes equal each pinned checkpoint's data (file size less its header)" { const c = try pinned(); - try testing.expectEqual(@as(usize, 3_975_229_440), parameters(c)); + for ([_]config.Quant{ .nvfp4, .mxfp8, .bf16 }) |quant| { + var q = c; + q.quant = quant; + try testing.expectEqual(@as(usize, 3_975_229_440), parameters(q)); + } // model.safetensors at 41f02335: 8 bytes of length, a 121,192-byte header, then the tensors try testing.expectEqual(@as(usize, 3_202_498_032 - 8 - 121_192), weightBytes(c)); + // the published MXFP8 model.safetensors: 8 bytes of length, a 91,616-byte header, then the tensors + var mx = c; + mx.quant = .mxfp8; + try testing.expectEqual(@as(usize, 4_750_755_304 - 8 - 91_616), weightBytes(mx)); + // the published BF16 model.safetensors: 8 bytes of length, a 56,584-byte header, then the tensors + var bf = c; + bf.quant = .bf16; + try testing.expectEqual(@as(usize, 7_950_515_472 - 8 - 56_584), weightBytes(bf)); } test "KV a token: 36, 56, 96 KiB at 1, 2, 4 loops; 7 GiB at 128K and two loops, plus fixed sliding rings" { diff --git a/zig/src/families/dory/names.zig b/zig/src/families/dory/names.zig index c3fca61e0d..57e375f6fa 100644 --- a/zig/src/families/dory/names.zig +++ b/zig/src/families/dory/names.zig @@ -7,11 +7,11 @@ const Config = config.Config; const Tensor = core.checkpoint.Tensor; const Checkpoint = core.checkpoint.Checkpoint; -/// A projection [n, k] as stored: NVFP4 codes [n, k/2] with e4m3 scales a 16 and an fp32 scale, or bf16 [n, k]. +/// A projection [n, k] as stored: NVFP4 codes a 2 and scales a 16, MXFP8 codes with u8 scales a 32, or bf16. pub const Proj = struct { w: Tensor, - blocks: ?Tensor = null, // NVFP4's weight_scale - scale: f32 = 1, // NVFP4's weight_scale_2 + blocks: ?Tensor = null, // NVFP4's or MXFP8's weight_scale + scale: f32 = 1, // NVFP4's weight_scale_2; 1 for MXFP8 and bf16 n: usize, k: usize, }; @@ -48,6 +48,11 @@ const Names = struct { fn proj(self: *Names, n: usize, k: usize, comptime fmt: []const u8, args: anytype) !Proj { if (self.quant == .bf16) return .{ .w = try self.bf16(&.{ n, k }, fmt ++ ".weight", args), .n = n, .k = k }; + if (self.quant == .mxfp8) { + const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .f8_e4m3, &.{ n, k }); + const b = try self.ck.expect(try self.name(fmt ++ ".weight_scale", args), .u8, &.{ n, k / 32 }); + return .{ .w = w, .blocks = b, .n = n, .k = k }; + } const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .u8, &.{ n, k / 2 }); const b = try self.ck.expect(try self.name(fmt ++ ".weight_scale", args), .f8_e4m3, &.{ n, k / 16 }); const s = try self.ck.expect(try self.name(fmt ++ ".weight_scale_2", args), .f32, &.{}); @@ -62,7 +67,8 @@ const Names = struct { /// Every tensor the loader reads, sublayer by sublayer; a missing, mistyped or leftover tensor refuses the checkpoint. pub fn map(gpa: std.mem.Allocator, ck: *Checkpoint, c: Config) !Map { - if (c.hidden % 16 != 0 or c.intermediate % 16 != 0 or c.qDim() % 16 != 0) return error.UnsupportedShapes; + const g = if (c.quant == .mxfp8) @as(usize, 32) else 16; + if (c.hidden % g != 0 or c.intermediate % g != 0 or c.qDim() % g != 0) return error.UnsupportedShapes; var nm: Names = .{ .ck = ck, .quant = c.quant }; const d = c.hidden; const f = c.intermediate; @@ -117,6 +123,10 @@ pub fn writeFile(a: std.mem.Allocator, dir: std.Io.Dir, c: Config) !void { } fn proj(list: *std.ArrayList(E), al: std.mem.Allocator, cc: Config, pre: []const u8, n: usize, k: usize) !void { if (cc.quant == .bf16) return bf(list, al, try std.fmt.allocPrint(al, "{s}.weight", .{pre}), &.{ n, k }); + if (cc.quant == .mxfp8) { + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k }) }); + return list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale", .{pre}), .dtype = "U8", .shape = try al.dupe(usize, &.{ n, k / 32 }) }); + } try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "U8", .shape = try al.dupe(usize, &.{ n, k / 2 }) }); try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k / 16 }) }); try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_2", .{pre}), .dtype = "F32", .shape = &.{} }); @@ -174,11 +184,11 @@ fn tmpPath(a: std.mem.Allocator, tmp: testing.TmpDir) ![]u8 { return std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); } -test "every tensor of a tiny NVFP4 and a tiny bf16 checkpoint maps" { +test "every tensor of a tiny NVFP4, MXFP8 and bf16 checkpoint maps" { var arena: std.heap.ArenaAllocator = .init(testing.allocator); defer arena.deinit(); const a = arena.allocator(); - for ([_]config.Quant{ .nvfp4, .bf16 }) |quant| { + for ([_]config.Quant{ .nvfp4, .mxfp8, .bf16 }) |quant| { var c = tiny; c.quant = quant; var tmp = testing.tmpDir(.{}); @@ -193,12 +203,13 @@ test "every tensor of a tiny NVFP4 and a tiny bf16 checkpoint maps" { const o = m.subs[4].attn.o; try testing.expectEqual(@as(usize, 64), o.n); try testing.expectEqual(@as(usize, 64), o.k); - try testing.expectEqual(quant == .nvfp4, o.blocks != null); + try testing.expectEqual(quant != .bf16, o.blocks != null); } } -test "the pinned NVFP4 checkpoint (TF_DORY_DIR): 1,083 tensors, none left over" { - const dir = @import("sys").getenv("TF_DORY_DIR") orelse return error.SkipZigTest; +/// The checkpoint at the `env` directory: 72 sublayers, an mlp last, its end tokens. +fn pinnedMaps(env: [*:0]const u8) !void { + const dir = @import("sys").getenv(env) orelse return error.SkipZigTest; const c = try Config.read(testing.allocator, io, dir); var ck = try Checkpoint.openModel(testing.allocator, io, dir); defer ck.close(); @@ -208,3 +219,15 @@ test "the pinned NVFP4 checkpoint (TF_DORY_DIR): 1,083 tensors, none left over" try testing.expect(m.subs[71] == .mlp and m.subs[70] == .attn); try testing.expectEqualSlices(u32, &.{ 2, 11 }, c.eos[0..c.eos_count]); } + +test "the pinned NVFP4 checkpoint (TF_DORY_DIR): 1,083 tensors, none left over" { + try pinnedMaps("TF_DORY_DIR"); +} + +test "the pinned MXFP8 checkpoint (TF_DORY_MXFP8_DIR): 795 tensors, none left over" { + try pinnedMaps("TF_DORY_MXFP8_DIR"); +} + +test "the pinned BF16 checkpoint (TF_DORY_BF16_DIR): 507 tensors, none left over" { + try pinnedMaps("TF_DORY_BF16_DIR"); +} From b1a7f2386ce78ea216867f195ae9054d7c416f5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:39:22 +0200 Subject: [PATCH 9/9] dory: the CUDA forward for one stream on NVFP4, greedy and verify windows bit-equal cuda_weights.zig loads the ModelOpt W4A16 NVFP4 checkpoint: every projection packed on the host for the lane matmul with its own weight_scale_2 (q, k, v and g launch separately), the nGPT scale vectors formed in fp32 as torch forms them ((sqk f)^2, (s_gate f)^2, suv sqrt(hidden), logit_scale / std), embeddings and head bf16. cuda_model.zig runs the sublayer schedule over per_loop KV caches (8 + 20 L + 8: full caches over the context, sliding ones over a ring of a chunk and the window): embed, the loop entry and exit L2 rows, the projections (prompt GEMM for prompt chunks, the lane for decode rows), tf_qk_rows_l2 into the cache, window attention at head dim 256 (window 1024, or none over the linear cache), the gate and SwiGLU L2 rows, the slerp residual, and the scaled head. Prompt chunks to 2,048 rows, decode windows to 16. tf-cuda-test dory MODEL LOOPS N TEXT, on nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-NVFP4 (41f02335), GB10: - greedy text is sane at 1, 2 and 4 loops ("The capital of France is" -> " Paris. The capital of the United States is Washington, D.C. ..."); - drafted == plain: verify windows of 2, 3, 5, 8 and 16 rows copied from the plain run give every row's logits bit-equal to the plain one-row decode (35 and 54 windows a run); - a prompt cut in two chunks gives the last row's logits bit-equal to one chunk's, and the prompt fed as decode rows gives the same first token. Eager decode, no graphs: 14.7, 20.8 and 33.7 ms a token at 1, 2 and 4 loops. --- zig/build/cuda.zig | 2 +- zig/src/families/dory/cuda_model.zig | 261 +++++++++++++++++++++++++ zig/src/families/dory/cuda_weights.zig | 145 ++++++++++++++ zig/src/families/dory/dory.zig | 4 + zig/tests/cuda/dory_tests.zig | 149 ++++++++++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 563 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/dory/cuda_model.zig create mode 100644 zig/src/families/dory/cuda_weights.zig create mode 100644 zig/tests/cuda/dory_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index c83d97f215..c6249e5b6c 100644 --- a/zig/build/cuda.zig +++ b/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); diff --git a/zig/src/families/dory/cuda_model.zig b/zig/src/families/dory/cuda_model.zig new file mode 100644 index 0000000000..c6bf2fe913 --- /dev/null +++ b/zig/src/families/dory/cuda_model.zig @@ -0,0 +1,261 @@ +//! Dory's forward on CUDA for one stream: per_loop KV caches, prompt chunks and decode windows on shared operators. +const std = @import("std"); +const cuda = @import("cuda"); +const config = @import("config.zig"); +const schedule = @import("schedule.zig"); +const wts = @import("cuda_weights.zig"); + +const Config = config.Config; +const DeviceBuffer = cuda.DeviceBuffer; +const WindowAttention = cuda.window_attention.WindowAttention; + +pub const eps: f32 = 1e-6; +/// Prompt rows a forward takes at most. +pub const prompt_rows = 2048; +/// Decode rows (a verify window) a forward takes at most. +pub const window_rows = 16; +const key_block = 64; + +/// Keys a sliding cache keeps: a chunk's rows written before its first row reads back its window. +pub fn ringSize(rows: usize, window: usize) usize { + return (rows + window + key_block - 1) / key_block * key_block; +} + +/// The shared operators the forward launches. +pub const Kernels = struct { + lane: cuda.qmmf.Lane, + prompt: cuda.nvfp4.Prompt, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + qk: cuda.qk_rows.QkRows, + act: cuda.act_rows.ActRows, + dense: cuda.dense_rows.DenseRows, + win: WindowAttention, + + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + var k: Kernels = undefined; + k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10); + errdefer k.lane.unload(); + k.prompt = try cuda.nvfp4.Prompt.load(d); + errdefer k.prompt.unload(); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.qk = try cuda.qk_rows.QkRows.load(d); + errdefer k.qk.unload(); + k.act = try cuda.act_rows.ActRows.load(d); + errdefer k.act.unload(); + k.dense = try cuda.dense_rows.DenseRows.load(d); + errdefer k.dense.unload(); + k.win = try WindowAttention.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.win.unload(); + k.dense.unload(); + k.act.unload(); + k.qk.unload(); + k.norms.unload(); + k.embed.unload(); + k.prompt.unload(); + k.lane.unload(); + } +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + c: Config, + sched: schedule.Schedule, + context: usize, + ring: usize, + k: []u64, // per cache: [len, HK, D] bf16 (len: context for full caches, the ring for sliding ones) + v: []u64, + inv: [2]u64, // fp32 [D / 2] inverse frequencies: sliding, full + s: Scratch = undefined, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + const Scratch = struct { ids: u64, pos: u64, h: u64, h2: u64, state: u64, input: u64, out: u64, q: u64, kk: u64, vv: u64, g: u64, qn: u64, kn: u64, vn: u64, att: u64, gated: u64, mg: u64, mu: u64, m: u64, part: u64, win: u64, sel: u64 }; + + fn alloc(m: *Model, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const Kernels, context: usize) !Model { + const c = w.config; + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .c = c, .sched = schedule.build(c), .context = context, .ring = ringSize(prompt_rows, c.window), .k = &.{}, .v = &.{}, .inv = .{ 0, 0 } }; + errdefer m.deinit(); + m.k = try gpa.alloc(u64, m.sched.caches); + m.v = try gpa.alloc(u64, m.sched.caches); + for (m.k, m.v, m.sched.cache_kind[0..m.sched.caches]) |*kc, *vc, kind| { + const len = if (kind == .full) context else m.ring; + const bytes = len * c.kvDim() * 2; + kc.* = try m.alloc(bytes); + vc.* = try m.alloc(bytes); + try d.check(d.api.cuMemsetD8Async(kc.*, 0, bytes, s.handle), "cuMemsetD8Async"); + try d.check(d.api.cuMemsetD8Async(vc.*, 0, bytes, s.handle), "cuMemsetD8Async"); + } + for ([_]f64{ c.theta_sliding, c.theta_full }, &m.inv) |theta, *p| { + var inv: [256]f32 = undefined; // vLLM's inv_freq in fp32: 1 / theta^(2j / D) + const half = c.head_dim / 2; + for (inv[0..half], 0..) |*v, j| v.* = 1 / std.math.pow(f32, @floatCast(theta), @as(f32, @floatFromInt(2 * j)) / @as(f32, @floatFromInt(c.head_dim))); + p.* = try m.alloc(half * 4); + try m.upload(s, p.*, std.mem.sliceAsBytes(inv[0..half])); + } + try m.scratch(k); + try s.synchronize(); + return m; + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.k); + m.gpa.free(m.v); + } + + fn upload(m: *Model, s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); + } + + fn copy(m: *Model, s: cuda.Stream, dst: u64, src: u64, bytes: usize) !void { + try m.d.check(m.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, s.handle), "cuMemcpyDtoDAsync"); + } + + fn window(m: *const Model, full: bool) cuda.window_attention.Window { + const c = m.c; + return .{ .h = c.heads, .hk = c.kv_heads, .d = c.head_dim, .win = if (full) 0 else c.window, .ring = if (full) m.context else m.ring, .scale = c.factors().attn }; + } + + fn scratch(m: *Model, k: *const Kernels) !void { + const c = m.c; + const r = prompt_rows; + const d = c.hidden; + var part: usize = 0; + for (m.w.subs) |sub| switch (sub) { + .attn => |a| for ([_]cuda.qlinear.Weight{ a.q, a.k, a.v, a.g, a.o }) |q| for (1..window_rows + 1) |rows| { + part = @max(part, k.lane.partBytes(rows, q.nvfp4)); + }, + .mlp => |p| for ([_]cuda.qlinear.Weight{ p.gate, p.up, p.down }) |q| for (1..window_rows + 1) |rows| { + part = @max(part, k.lane.partBytes(rows, q.nvfp4)); + }, + }; + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .h = try m.alloc(r * d * 2), + .h2 = try m.alloc(r * d * 2), + .state = try m.alloc(r * d * 2), + .input = try m.alloc(r * d * 2), + .out = try m.alloc(r * d * 2), + .q = try m.alloc(r * c.qDim() * 2), + .kk = try m.alloc(r * c.kvDim() * 2), + .vv = try m.alloc(r * c.kvDim() * 2), + .g = try m.alloc(r * c.qDim() * 2), + .qn = try m.alloc(r * c.qDim() * 2), + .kn = try m.alloc(r * c.kvDim() * 2), + .vn = try m.alloc(r * c.kvDim() * 2), + .att = try m.alloc(r * c.qDim() * 2), + .gated = try m.alloc(r * c.qDim() * 2), + .mg = try m.alloc(r * c.intermediate * 2), + .mu = try m.alloc(r * c.intermediate * 2), + .m = try m.alloc(r * c.intermediate * 2), + .part = try m.alloc(part), + .win = try m.alloc(@max(WindowAttention.scratchBytes(m.window(true), window_rows), WindowAttention.scratchBytes(m.window(false), window_rows))), + .sel = try m.alloc(r * d * 2), + }; + } + + /// fp32 logits [rows.len, V] at `out` of `tokens` at positions p0.. (a prompt chunk, or decode rows: a window). + pub fn forward(m: *Model, k: *const Kernels, s: cuda.Stream, tokens: []const u32, p0: usize, prompt: bool, rows: []const usize, out: u64) !void { + const c = m.c; + const n = tokens.len; + if (n == 0 or n > @as(usize, if (prompt) prompt_rows else window_rows) or p0 + n > m.context) return error.TooManyRows; + if (n > m.ring - c.window) return error.ChainOverrunsRing; + var host: [2 * prompt_rows]i32 = undefined; + for (tokens, 0..) |t, i| { + host[i] = @intCast(t); + host[n + i] = @intCast(p0 + i); + } + try m.upload(s, m.s.ids, std.mem.sliceAsBytes(host[0..n])); + try m.upload(s, m.s.pos, std.mem.sliceAsBytes(host[n .. 2 * n])); + const lin: cuda.qlinear.Linear = .{ .lane = &k.lane, .nvfp4_prompt = &k.prompt }; + const d = c.hidden; + const f = c.factors(); + const x = m.s; + try k.embed.bf16(s, x.ids, m.w.embed, x.h, n, d); + var h = x.h; + var h2 = x.h2; + for (m.sched.slice()) |step| { + if (step.enter) { + if (step.loop == 0) { + try m.copy(s, x.input, h, n * d * 2); + try m.copy(s, x.state, h, n * d * 2); + } + try k.norms.l2(s, x.state, x.input, h, n, d, eps); + } + const alpha = switch (m.w.subs[step.sub]) { + .attn => |a| blk: { + try m.attention(k, s, lin, a, step, h, n, p0, prompt); + break :blk a.alpha; + }, + .mlp => |p| blk: { + try m.proj(lin, s, p.gate, h, x.mg, n, prompt); + try m.proj(lin, s, p.up, h, x.mu, n, prompt); + try k.act.swigluL2(s, x.mg, c.intermediate, x.mu, c.intermediate, p.su_g, p.su_u, x.m, n, c.intermediate, eps); + try m.proj(lin, s, p.down, x.m, x.out, n, prompt); + break :blk p.alpha; + }, + }; + try k.norms.slerp(s, h, x.out, alpha, f.alpha, h2, n, d, eps); + std.mem.swap(u64, &h, &h2); + if (step.leave) { + try k.norms.l2(s, h, null, x.state, n, d, eps); + try m.copy(s, h, x.state, n * d * 2); + } + } + for (rows, 0..) |r, j| try m.copy(s, x.sel + j * d * 2, h + r * d * 2, d * 2); + try k.dense.rowsScaled(s, x.sel, d, m.w.head, m.w.logit, out, c.vocab, rows.len, c.vocab, d); + } + + fn attention(m: *Model, k: *const Kernels, s: cuda.Stream, lin: cuda.qlinear.Linear, a: wts.Attn, step: schedule.Step, h: u64, n: usize, p0: usize, prompt: bool) !void { + const c = m.c; + const x = m.s; + try m.proj(lin, s, a.q, h, x.q, n, prompt); + try m.proj(lin, s, a.k, h, x.kk, n, prompt); + try m.proj(lin, s, a.v, h, x.vv, n, prompt); + try m.proj(lin, s, a.g, h, x.g, n, prompt); + const full = step.kind == .full; + const cache = step.cache.?; + const dh: i64 = @intCast(c.head_dim); + const src: cuda.qk_rows.Src = .{ .q = x.q, .q_ld = @intCast(c.qDim()), .q_hs = dh, .k = x.kk, .k_ld = @intCast(c.kvDim()), .k_hs = dh, .v = x.vv, .v_ld = @intCast(c.kvDim()), .v_hs = dh }; + const shape: cuda.qk_rows.Shape = .{ .l2 = true, .q_scale = a.qs, .inv = m.inv[@intFromBool(full)], .rot = c.head_dim, .eps = eps }; + const len = if (full) m.context else m.ring; + try k.qk.rowsFrom(s, src, .{ .h = c.heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, null, .{ .q = x.qn, .k = x.kn, .v = x.vn, .kc = m.k[cache], .vc = m.v[cache], .len = len }, n); + const w = m.window(full); + if (prompt) try k.win.promptRows(s, x.qn, m.k[cache], m.v[cache], x.att, p0, n, w, null) else try k.win.decode(s, x.qn, m.k[cache], m.v[cache], x.att, x.pos, null, w, x.win, n); + try k.act.gateL2(s, x.att, x.g, c.qDim(), a.sg, x.gated, n, c.qDim(), eps); + try m.proj(lin, s, a.o, x.gated, x.out, n, prompt); + } + + fn proj(m: *const Model, lin: cuda.qlinear.Linear, s: cuda.Stream, w: cuda.qlinear.Weight, in: u64, out: u64, n: usize, prompt: bool) !void { + const rows: cuda.qlinear.Rows = .{ .x = in, .part = m.s.part }; + try (if (prompt) lin.prompt(s, w, rows, out, n) else lin.decode(s, w, rows, out, n)); + } +}; + +test "the forward's launches type-check against the shared operators" { + std.testing.refAllDecls(Model); + std.testing.refAllDecls(Kernels); + std.testing.refAllDecls(wts); +} diff --git a/zig/src/families/dory/cuda_weights.zig b/zig/src/families/dory/cuda_weights.zig new file mode 100644 index 0000000000..e429435208 --- /dev/null +++ b/zig/src/families/dory/cuda_weights.zig @@ -0,0 +1,145 @@ +//! Dory's weights on the GPU: NVFP4 projections packed for the lane matmul, scale vectors as torch forms them. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const config = @import("config.zig"); +const names = @import("names.zig"); + +const Config = config.Config; +const DeviceBuffer = cuda.DeviceBuffer; +const Weight = cuda.qlinear.Weight; +const bfr = core.bf16_round; + +pub const Attn = struct { + alpha: u64, // bf16 [D] + q: Weight, + k: Weight, + v: Weight, + g: Weight, + o: Weight, + qs: u64, // fp32 [Q]: (sqk f)^2 + sg: u64, // fp32 [Q]: (s_gate f)^2 +}; + +pub const Mlp = struct { + alpha: u64, + gate: Weight, + up: Weight, + down: Weight, + su_g: u64, // fp32 [F]: suv_gate sqrt(D) + su_u: u64, +}; + +pub const Sub = union(enum) { attn: Attn, mlp: Mlp }; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: Config, + embed: u64 = 0, // bf16 [V, D] + head: u64 = 0, // bf16 [V, D] + logit: u64 = 0, // fp32 [V]: logit_scale / std as torch multiplies + subs: []Sub = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + bytes: usize = 0, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.subs); + w.* = undefined; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + try w.buffers.append(w.gpa, b); + w.bytes += bytes.len; + return b.ptr; + } + + /// A projection in its lane layout: NVFP4 words and scale tiles (padded to 128 rows), or bf16 rows. + fn proj(w: *Weights, d: *const cuda.Driver, p: names.Proj) !Weight { + if (p.blocks == null) { + const n = cuda.modelopt.padded(p.n); + const out = try w.gpa.alloc(u16, n * p.k); + defer w.gpa.free(out); + const src = try w.gpa.alloc(u16, p.n * p.k); + defer w.gpa.free(src); + @memcpy(std.mem.sliceAsBytes(src), p.w.bytes); + cuda.qmmf.packBf16(out, src, p.n, p.k); + return .{ .bf16 = cuda.modelopt.bf16Weight(try w.upload(d, std.mem.sliceAsBytes(out)), p.n, p.k) }; + } + const words = try w.gpa.alloc(u32, cuda.nvfp4.padded(p.n) / 64 * (p.k / 64) * 512); + defer w.gpa.free(words); + const tiles = try w.gpa.alloc(u8, cuda.nvfp4.padded(p.n) * (p.k / 16)); + defer w.gpa.free(tiles); + cuda.nvfp4.packWords(words, p.w.bytes, p.n, p.k); + cuda.nvfp4.packScales(tiles, p.blocks.?.bytes, p.n, p.k); + const wp = try w.upload(d, std.mem.sliceAsBytes(words)); + const tp = try w.upload(d, tiles); + return .{ .nvfp4 = cuda.nvfp4.weight(wp, tp, p.scale, p.n, p.k) }; + } + + /// fp32 vector: each bf16 value through `f` (the torch steps D3 of the progress log). + fn vector(w: *Weights, d: *const cuda.Driver, t: core.checkpoint.Tensor, f: anytype) !u64 { + const n = t.bytes.len / 2; + const v = try w.gpa.alloc(f32, n); + defer w.gpa.free(v); + for (v, 0..) |*o, i| o.* = f.at(bfr.value(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little))); + return w.upload(d, std.mem.sliceAsBytes(v)); + } +}; + +/// (v * factor)^2 in fp32, as dory.py squares sqk and s_gate after their factor. +const Square = struct { + f: f32, + fn at(s: Square, v: f32) f32 { + const t = v * s.f; + return t * t; + } +}; + +/// v * factor in fp32: suv by sqrt(hidden), logit_scale by 1 / std. +const Times = struct { + f: f32, + fn at(s: Times, v: f32) f32 { + return v * s.f; + } +}; + +/// config.json, every tensor mapped and checked, then sublayer by sublayer onto the GPU; `loops` overrides the count. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, loops: ?usize) !Weights { + var c = try Config.read(gpa, io, dir); + if (loops) |l| c = try c.withLoops(l); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var m = try names.map(gpa, &ck, c); + defer m.deinit(); + var w: Weights = .{ .gpa = gpa, .config = c }; + errdefer w.deinit(); + const f = c.factors(); + w.subs = try gpa.alloc(Sub, c.sublayers); + w.embed = try w.upload(d, m.embed.bytes); + for (w.subs, m.subs) |*out, s| out.* = switch (s) { + .attn => |a| .{ .attn = .{ + .alpha = try w.upload(d, a.alpha.bytes), + .q = try w.proj(d, a.q), + .k = try w.proj(d, a.k), + .v = try w.proj(d, a.v), + .g = try w.proj(d, a.g), + .o = try w.proj(d, a.o), + .qs = try w.vector(d, a.sqk, Square{ .f = f.sqk }), + .sg = try w.vector(d, a.s_gate, Square{ .f = f.gate }), + } }, + .mlp => |p| .{ .mlp = .{ + .alpha = try w.upload(d, p.alpha.bytes), + .gate = try w.proj(d, p.gate), + .up = try w.proj(d, p.up), + .down = try w.proj(d, p.down), + .su_g = try w.vector(d, p.suv_gate, Times{ .f = f.suv }), + .su_u = try w.vector(d, p.suv_up, Times{ .f = f.suv }), + } }, + }; + w.head = try w.upload(d, m.head.bytes); + w.logit = try w.vector(d, m.logit_scale, Times{ .f = f.logit }); + return w; +} diff --git a/zig/src/families/dory/dory.zig b/zig/src/families/dory/dory.zig index 50fc70624c..40046daeb4 100644 --- a/zig/src/families/dory/dory.zig +++ b/zig/src/families/dory/dory.zig @@ -7,6 +7,8 @@ pub const schedule = @import("schedule.zig"); pub const memory = @import("memory.zig"); /// The float64 forward the kernels and the CUDA model are checked against. pub const ref = @import("ref.zig"); +pub const cuda_weights = @import("cuda_weights.zig"); +pub const cuda_model = @import("cuda_model.zig"); test { _ = @import("config.zig"); @@ -14,4 +16,6 @@ test { _ = @import("schedule.zig"); _ = @import("memory.zig"); _ = @import("ref.zig"); + _ = @import("cuda_weights.zig"); + _ = @import("cuda_model.zig"); } diff --git a/zig/tests/cuda/dory_tests.zig b/zig/tests/cuda/dory_tests.zig new file mode 100644 index 0000000000..f221902292 --- /dev/null +++ b/zig/tests/cuda/dory_tests.zig @@ -0,0 +1,149 @@ +//! Dory on the real NVFP4 checkpoint: greedy tokens from a prompt, then verify windows of 2 to 16 rows bit-equal. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const dory = @import("dory"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const fwd = dory.cuda_model; + +pub const Rig = struct { + gpu: Gpu, + w: dory.cuda_weights.Weights, + k: fwd.Kernels, + s: cuda.Stream, + m: fwd.Model, + logits: cuda.DeviceBuffer, + + pub fn init(r: *Rig, gpu: Gpu, model: []const u8, loops: usize, context: usize) !void { + r.gpu = gpu; + r.w = try dory.cuda_weights.load(gpu.gpa, gpu.io, gpu.d, model, loops); + errdefer r.w.deinit(); + r.k = try fwd.Kernels.load(gpu.ctx); + errdefer r.k.deinit(); + r.s = try cuda.Stream.init(gpu.d, false); + errdefer r.s.deinit(); + r.m = try fwd.Model.init(gpu.gpa, gpu.d, r.s, &r.w, &r.k, context); + errdefer r.m.deinit(); + r.logits = try cuda.DeviceBuffer.alloc(gpu.d, fwd.prompt_rows * r.w.config.vocab * 4); + } + + pub fn deinit(r: *Rig) void { + r.logits.free(); + r.m.deinit(); + r.s.deinit(); + r.k.deinit(); + r.w.deinit(); + } + + /// fp32 logits of `tokens` at p0.. (each row's), on the host. + pub fn rows(r: *Rig, tokens: []const u32, p0: usize, prompt: bool, pick: []const usize) ![]f32 { + try r.m.forward(&r.k, r.s, tokens, p0, prompt, pick, r.logits.ptr); + try r.s.synchronize(); + const v = r.w.config.vocab; + const out = try r.gpu.gpa.alloc(f32, pick.len * v); + try r.logits.download(0, std.mem.sliceAsBytes(out)); + return out; + } +}; + +fn argmax(x: []const f32) u32 { + var best: usize = 0; + for (x, 0..) |v, i| if (v > x[best]) { + best = i; + }; + return @intCast(best); +} + +/// tf-cuda-test dory MODEL LOOPS N TEXT: N greedy tokens one row at a time, then every window width against them. +pub fn greedy(gpu: Gpu, model: []const u8, loops_text: []const u8, count_text: []const u8, text: []const u8) !void { + const gpa = gpu.gpa; + const loops = try std.fmt.parseInt(usize, loops_text, 10); + const count = try std.fmt.parseInt(usize, count_text, 10); + var tok = try core.tokenizer.loadTokenizer(gpu.io, gpa, model); + defer tok.deinit(); + const ids = try tok.encode(gpa, text); + defer gpa.free(ids); + var prompt: std.ArrayList(u32) = .empty; + defer prompt.deinit(gpa); + if (ids.len == 0 or ids[0] != 1) try prompt.append(gpa, 1); // + try prompt.appendSlice(gpa, ids); + var rig: Rig = undefined; + try rig.init(gpu, model, loops, 4096); + defer rig.deinit(); + const v = rig.w.config.vocab; + std.debug.print("INFO dory: {d} loops, {d} prompt tokens, {d} weight bytes on the GPU\n", .{ loops, prompt.items.len, rig.w.bytes }); + const n0 = prompt.items.len; + const t0 = check.now(gpu.io); + const first = try rig.rows(prompt.items, 0, true, &.{n0 - 1}); + defer gpa.free(first); + const plain = try gpa.alloc(u32, count); + defer gpa.free(plain); + const kept = try gpa.alloc(f32, count * v); // each decode row's logits + defer gpa.free(kept); + plain[0] = argmax(first); + const t1 = check.now(gpu.io); + for (1..count) |i| { + const row = try rig.rows(plain[i - 1 .. i], n0 + i - 1, false, &.{0}); + defer gpa.free(row); + @memcpy(kept[i * v ..][0..v], row); + plain[i] = argmax(row); + } + const t2 = check.now(gpu.io); + const said = try tok.decode(gpa, plain, false); + defer gpa.free(said); + std.debug.print("RESULT dory plain: {d} tokens, prompt {d:.1} ms, decode {d:.2} ms a token: \"{s}\"\n", .{ count, @as(f64, @floatFromInt(t1 - t0)) / 1e6, @as(f64, @floatFromInt(t2 - t1)) / 1e6 / @as(f64, @floatFromInt(count - 1)), said }); + std.debug.print("RESULT dory ids:", .{}); + for (plain) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + // a verify window of w rows: the last kept token and w - 1 drafts copied from the plain run + var bad: usize = 0; + var windows: usize = 0; + for ([_]usize{ 2, 3, 5, 8, 16 }) |w| { + gpa.free(try rig.rows(prompt.items, 0, true, &.{n0 - 1})); // the caches as the plain run had them + var at: usize = 1; + while (at + w <= count) : (at += w) { + var pick: [fwd.window_rows]usize = undefined; + for (pick[0..w], 0..) |*p, j| p.* = j; + const got = try rig.rows(plain[at - 1 .. at - 1 + w], n0 + at - 1, false, pick[0..w]); + defer gpa.free(got); + windows += 1; + for (0..w) |j| if (!std.mem.eql(u8, std.mem.sliceAsBytes(got[j * v ..][0..v]), std.mem.sliceAsBytes(kept[(at + j) * v ..][0..v]))) { + bad += 1; + if (bad < 4) std.debug.print("FAIL dory window {d} at {d} row {d}: logits differ from the plain row's (argmax {d} vs {d})\n", .{ w, at, j, argmax(got[j * v ..][0..v]), plain[at + j] }); + }; + } + } + try check.expect(bad == 0, "dory drafted == plain: {d} rows differ in {d} windows", .{ bad, windows }); + check.pass("dory drafted == plain: {d} windows of 2 to 16 rows, every row's logits bit-equal to the plain row's", .{windows}); + try seams(&rig, prompt.items, first); +} + +/// The prompt cut in two chunks (bit-equal: per_loop rows are chunk-invariant), and fed as decode rows (lane GEMM). +fn seams(rig: *Rig, prompt: []const u32, whole: []const f32) !void { + const gpa = rig.gpu.gpa; + const n = prompt.len; + if (n < 2) return; + const cut = n / 2; + gpa.free(try rig.rows(prompt[0..cut], 0, true, &.{cut - 1})); + const two = try rig.rows(prompt[cut..], cut, true, &.{n - cut - 1}); + defer gpa.free(two); + const same = std.mem.eql(u8, std.mem.sliceAsBytes(two), std.mem.sliceAsBytes(whole)); + try check.expect(same, "dory prompt cut at {d} of {d}: last-row logits differ from one chunk's", .{ cut, n }); + check.pass("dory prompt cut at {d} of {d}: last-row logits bit-equal to one chunk's", .{ cut, n }); + var last: []f32 = &.{}; + defer gpa.free(last); + for (prompt, 0..) |t, i| { + gpa.free(last); + last = try rig.rows(&.{t}, i, false, &.{0}); + } + var most: f32 = 0; + var top: f32 = 0; + for (last, whole) |a, b| { + most = @max(most, @abs(a - b)); + top = @max(top, @abs(b)); + } + try check.expect(argmax(last) == argmax(whole), "dory prompt as decode rows: first token {d}, the prompt GEMM's {d}", .{ argmax(last), argmax(whole) }); + check.pass("dory prompt as decode rows: the same first token, logits within {e} (largest {e})", .{ most, top }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index ccb288ab54..d13913ee83 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -99,6 +99,7 @@ const attn_back_tests = @import("attn_back_tests.zig"); const fn_hc_back_tests = @import("fn_hc_back_tests.zig"); const fn_pocket_tests = @import("fn_pocket_tests.zig"); const kolibri_tests = @import("kolibri_tests.zig"); +const dory_tests = @import("dory_tests.zig"); const kolibri_layer_tests = @import("kolibri_layer_tests.zig"); const aot_fallback_tests = @import("aot_fallback_tests.zig"); const qwen_draft_frames = @import("qwen_draft_frames.zig"); @@ -220,6 +221,7 @@ const usage = \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ dory MODEL LOOPS N TEXT Dory (NVFP4): N greedy tokens from TEXT, then verify windows of 2 to 16 rows bit-equal \\ kolibri-reference MODEL IDS Kolibri-1's layer 0 and head against float64 from the checkpoint (prompt and decode rows) \\ kolibri-layers MODEL IDS Kolibri-1's window edge (513 keys) and its fifth layer's full reach, on the checkpoint \\ kolibri-full MODEL IDS N Kolibri-1: eager rows, steps, drafted, sampled, shared rounds, cuts and kept states alike @@ -418,6 +420,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "dory")) return dory_tests.greedy(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "kolibri-reference")) return kolibri_layer_tests.reference(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "kolibri-layers")) return kolibri_layer_tests.layers(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "kolibri-full")) return kolibri_tests.full(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2));