diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 75bd900cda..32597e615b 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -149,6 +149,14 @@ CUDA also checks `window_attention`'s rows at head dims 256 and 512 (`window_row prompt tiles with P as bf16 or in two halves, a row's last key, windows over a ring or none over a linear cache), each output within float64's range over the spread its fp32 steps and tensor-core sums allow (`core/attention_range.zig`; seven controls, each one planted fault). +CUDA also checks Dory's nGPT rows: the L2 norm rows with and without their bf16 add, and the normalized lerp residual +(`norms_l2.cuh`), each output within float64's range at its bf16 points (three controls, each dropping one point). +CUDA also checks Dory's attention inputs (`qk_l2.cuh`): RoPE first, each head's L2 norm, q times its scale vector, v as is, at +head dim 256 and positions past 199,800 (two controls, each dropping one rounding point). +CUDA also checks Dory's activation rows (`act_l2.cuh`): the scaled sigmoid output gate over all heads and the scaled SwiGLU, +each row then L2-normed, against float64's ranges at their bf16 points and expf's ulps (four controls, one point each). +CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) times a per-column scale), bit-equal to its +unscaled tile's outputs rounded and scaled on the host. HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then 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/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/cuda_row_contracts.zig b/zig/build/cuda_row_contracts.zig index 2bf71d7940..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_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/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/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 91b3171bbf..9c43d2281a 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -76,6 +76,8 @@ extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_ctl_qk_round( } // Gate 3: act_rows.cu's tf_softcap_rows keeping only the rounding points in `keep`; the exact check must go red. +#include "act_round.cuh" + extern "C" __global__ void __launch_bounds__(256) tf_ctl_softcap(const __nv_bfloat16* x, __nv_bfloat16* out, int n, float cap, int keep) { tf_act_round::softcap_rows(x, out, n, cap, keep); @@ -93,8 +95,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_embed_affine(const int* // Gate 3: act_rows.cu's MLX activations keeping only the rounding points in `keep`; the exact check must go red. -#include "act_round.cuh" - extern "C" __global__ void __launch_bounds__(256) tf_ctl_act_round(const __nv_bfloat16* gate, long long g_ld, const __nv_bfloat16* up, long long u_ld, __nv_bfloat16* out, float* xs, int n, int geglu, @@ -235,3 +235,55 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_vit_g4_pool(const __nv_ float root, int keep) { tf_vit_g4::pool_std(x, bias, scale, out, gw, C, root, keep); } + +// Gate 3: norms.cu's nGPT rows with one rounding point dropped each; their range check must go red. +#include "norms_l2.cuh" + +extern "C" __global__ void __launch_bounds__(tf_l2::NT) tf_ctl_l2_drop_sum(const __nv_bfloat16* x, const __nv_bfloat16* r, + __nv_bfloat16* y, int width, float eps) { + tf_l2::l2_rows<0>(x, r, y, width, eps); +} + +#define TF_CTL_SLERP(NAME, KEEP) \ + extern "C" __global__ void __launch_bounds__(tf_l2::NT) NAME(const __nv_bfloat16* res, const __nv_bfloat16* out, \ + const __nv_bfloat16* alpha, float scale, \ + __nv_bfloat16* h, int width, float eps) { \ + tf_l2::slerp_rows(res, out, alpha, scale, h, width, eps); \ + } +TF_CTL_SLERP(tf_ctl_slerp_drop_a, tf_l2::B) +TF_CTL_SLERP(tf_ctl_slerp_drop_b, tf_l2::A) +#undef TF_CTL_SLERP + +// Gate 3: qk_rows.cu's tf_qk_rows_l2 with one rounding point dropped each; its range check must go red. +#include "qk_l2.cuh" + +#define TF_CTL_QK_L2(NAME, KEEP) \ + extern "C" __global__ void __launch_bounds__(tf_qk_l2::WIDE) NAME( \ + const tf_qk::QkSrc in, const float* qs, const float* cosv, const float* sinv, const float* inv, int rot, \ + const int* pos, const int* slot, const tf_qk::QkDst dst, __nv_bfloat16* kc, __nv_bfloat16* vc, \ + long long len, int h, int hk, float eps) { \ + tf_qk_l2::rows(in, qs, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps); \ + } +TF_CTL_QK_L2(tf_ctl_qk_l2_drop_rope, tf_qk_l2::NORM) +TF_CTL_QK_L2(tf_ctl_qk_l2_drop_norm, tf_qk_l2::ROPE) +#undef TF_CTL_QK_L2 + +// Gate 3: act_rows.cu's nGPT activations with one rounding point dropped each; their range checks must go red. +#include "act_l2.cuh" + +extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) tf_ctl_gate_l2_drop_gated(const __nv_bfloat16* o, const __nv_bfloat16* g, + long long g_ld, const float* sg, + __nv_bfloat16* out, int n, float eps) { + tf_act_l2::gate_rows<0>(o, g, g_ld, sg, out, n, eps); +} + +#define TF_CTL_SWIGLU_L2(NAME, KEEP) \ + extern "C" __global__ void __launch_bounds__(tf_act_l2::NT) NAME( \ + const __nv_bfloat16* g, long long g_ld, const __nv_bfloat16* u, long long u_ld, const float* sug, \ + const float* suu, __nv_bfloat16* out, int n, float eps) { \ + tf_act_l2::swiglu_rows(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/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/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/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 new file mode 100644 index 0000000000..a56d1d8782 --- /dev/null +++ b/zig/src/core/contract/dory_cuda_records.zig @@ -0,0 +1,14 @@ +//! 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 } } }, + .{ .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/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..3e267e3bbb --- /dev/null +++ b/zig/src/core/l2_ref.zig @@ -0,0 +1,245 @@ +//! 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) }; + }; + 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| { + 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); +} + +/// 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); + } +} + +/// 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/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/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/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/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/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/src/families/dory/config.zig b/zig/src/families/dory/config.zig new file mode 100644 index 0000000000..efbd2dd7f4 --- /dev/null +++ b/zig/src/families/dory/config.zig @@ -0,0 +1,555 @@ +//! 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: 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 { + 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), 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; + 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 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; + 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; +} + +/// 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, .{}); + 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)); +} + +/// 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/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 new file mode 100644 index 0000000000..40046daeb4 --- /dev/null +++ b/zig/src/families/dory/dory.zig @@ -0,0 +1,21 @@ +//! 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"); +pub const cuda_weights = @import("cuda_weights.zig"); +pub const cuda_model = @import("cuda_model.zig"); + +test { + _ = @import("config.zig"); + _ = @import("names.zig"); + _ = @import("schedule.zig"); + _ = @import("memory.zig"); + _ = @import("ref.zig"); + _ = @import("cuda_weights.zig"); + _ = @import("cuda_model.zig"); +} diff --git a/zig/src/families/dory/memory.zig b/zig/src/families/dory/memory.zig new file mode 100644 index 0000000000..c74b3c955c --- /dev/null +++ b/zig/src/families/dory/memory.zig @@ -0,0 +1,85 @@ +//! 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, 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, + }; +} + +/// 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 tensors' bytes equal each pinned checkpoint's data (file size less its header)" { + const c = try pinned(); + 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" { + 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..57e375f6fa --- /dev/null +++ b/zig/src/families/dory/names.zig @@ -0,0 +1,233 @@ +//! 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 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 or MXFP8's weight_scale + scale: f32 = 1, // NVFP4's weight_scale_2; 1 for MXFP8 and bf16 + 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 }; + 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, &.{}); + 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 { + 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; + 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 }); + 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 = &.{} }); + } + }; + 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, MXFP8 and bf16 checkpoint maps" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for ([_]config.Quant{ .nvfp4, .mxfp8, .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 != .bf16, o.blocks != null); + } +} + +/// 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(); + 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]); +} + +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"); +} 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 + }; +} 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 })); +} 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_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 new file mode 100644 index 0000000000..6735c08847 --- /dev/null +++ b/zig/tests/cuda/contract_dory.zig @@ -0,0 +1,64 @@ +//! 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 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; + +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, + 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 { + 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_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()); + 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()); + 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 { + 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(); + for (self.dense[0..self.denses]) |*t| t.deinit(); + if (self.dense_rows) |*q| q.unload(); + } +}; 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_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)))); +} 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") }; 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));