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/2] 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/2] 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") };