diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index bc1e3f6768..2a76a526a3 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -40,7 +40,9 @@ The test of a kernel fills a `contract.spec.Spec` and hands the harness a few ho 3. Shapes. Every served shape, then each tiled dimension at its least size, a tile minus a step, a tile, a tile plus a step and the largest served size. Each output must be within `|want| * outputs' half ulp + gamma(depth) * magnitude` of float64, or for a replacement kernel bit-equal to the kernel it replaces. Control: a row sum whose partial tail - loses its last input; it passes at whole tiles and fails at the edges. + loses its last input; it passes at whole tiles and fails at the edges. A kernel with bf16 rounding points between + its steps (MLX's op-by-op rows) checks each output against float64's value at those points, as a range over the + spread its fp32 steps allow, and its controls each drop one rounding point, which that check must catch. 4. fp32 sums. The inputs are integers whose partial sums reach about 2^20 in any order and stay below 2^24, so fp32 sums are exact whatever the kernel's order; the second half of each row mirrors the first, so the result is small. Before the kernel runs, the harness replays the sums on the CPU in fp16 and bf16 (one chain, contiguous halves, and @@ -104,15 +106,24 @@ Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), a formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10. -CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router logits (and their 128-row tiles, +CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and +the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles, bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and shared add over bf16 share sums (bit-equal to `tf_super_combine` and `tf_super_moe_out` over the fp32 shares), the Mamba scan of a verify window (its host twin's bytes), qmm_group's cluster tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel`), the checkpoint math's FP8 128-row prompt tiles (both slice meetings, bit-equal to the 64-row tiles) and NVFP4 experts: the decode-sized forms and the prompt tiles (up to down's NVFP4 -rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), and the -reduce-scatter's `tf_rank_sum`. +rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), the +reduce-scatter's `tf_rank_sum`, and `window_attention`'s rows at head dims 256 and 512 (`window_rows.cu`: decode chunks +and their merge, 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), and the activation rows in MLX's bf16 steps (GeGLU, and +SwiGLU as DFlash drafters round it; twelve controls, each dropping one rounding point; their 64-value sums bit-equal to +the fp32 tree), and the MLX affine token rows (4 bits, and 2 to 6 or 8 bits times a bf16 scale, each +output bit-equal to float64 at its rounding points; two controls, one dropping the rounding point, one folding the scale +into the group's scale and bias), and qk_rows' attention inputs at head dims to 512 with a unit v norm and MLX's bf16 +rows (three controls, each dropping one option). ## CUDA records per nvcc release diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6dfec448ca..c571d8f5ad 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -115,6 +115,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu 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)); + 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); // tf-tp-watch-check: the ranks' lines with no GPU or NCCL; a failing or dying rank stops every rank @@ -204,6 +205,7 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, .{ .name = "qwen35", .module = qwen35 }, .{ .name = "kolibri1", .module = kolibri1Module(b, target, optimize, cuda, core, lanes) }, + .{ .name = "gemma4_cuda", .module = @import("gemma4_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, }, }); return .{ .api = api, .engines = mod }; @@ -292,6 +294,14 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, }) })); step.dependOn(&writeback_contract.step); + const tail_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_norms_tail.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&tail_contract.step); const utility_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_utility_ref.zig"), .target = host, @@ -321,6 +331,16 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st }) })); step.dependOn(&qsa_contract.step); flash_contract.dependOn(&qsa_contract.step); + // the Gemma 4 operators' contract tests (#1048): window rows, MLX activations and soft cap, embed rows, qk rows + for ([_][]const u8{ "contract_window_rows.zig", "contract_act_round.zig", "contract_softcap.zig", "contract_embed.zig", "contract_qk_round.zig" }) |f| { + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path(b.fmt("zig/tests/cuda/{s}", .{f})), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })).step); + } const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)"); const projection_step = b.step("test-cuda-projection-fixture", "Projection fixture dispatch controls with CPU-only CUDA calls"); for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| { @@ -471,6 +491,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); @import("glm_family_host.zig").steps(b, cuda, glm, mods, step); @import("kolibri.zig").host(b, kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), native_modules.api, step); + @import("gemma4_cuda.zig").host(b, @import("gemma4_cuda.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); @import("glm_contracts.zig").tests(b, glm, mods.lanes, cuda, mods.core); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 25315ca223..13647715a8 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -70,7 +70,7 @@ pub const kernels = [_]Kernel{ .{ .name = "ck_lane", .flags = &.{"-O3"} }, // nvfp4/lane4.cu's A8: FP8 rows times FP8 weights (checkpoint math) .{ .name = "ck_experts", .flags = &.{"-O3"}, .arch_specific = true }, // ours: NVFP4 x NVFP4 experts (sm_121a) .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) - .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) + .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: MLX affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig) @@ -82,6 +82,7 @@ pub const kernels = [_]Kernel{ .{ .name = "tree_attention", .flags = &.{"-O3"} }, // ours: committed keys and paths (cuda/tree_attention.zig) .{ .name = "block_attention", .flags = &.{"-O3"} }, // ours: a drafter's block over its context (cuda/block_attention.zig) .{ .name = "window_attention", .flags = &.{"-O3"} }, // ours: sliding-window attention (cuda/window_attention.zig) + .{ .name = "window_rows", .flags = &.{"-O3"} }, // ours: window attention at head dims 256 and 512 (cuda/window_attention.zig) .{ .name = "moe_glue", .flags = &.{"-O3"} }, // ours: sigmoid top-k and the weighted combine (cuda/moe_glue.zig) .{ .name = "dense_rows", .flags = &.{"-O3"} }, // ours: bf16 rows times weight rows into fp32 (cuda/dense_rows.zig) .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts diff --git a/zig/build/gemma4_cuda.zig b/zig/build/gemma4_cuda.zig new file mode 100644 index 0000000000..32190bd895 --- /dev/null +++ b/zig/build/gemma4_cuda.zig @@ -0,0 +1,12 @@ +//! Gemma 4's CUDA family (#1048 step 8) over the maintained runtime, the core and the lane core; its host tests. +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, lanes: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); +} + +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-gemma4-cuda-host", "Gemma 4 CUDA config copies, RoPE tables, snapshot runs and refusals without a GPU").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/kernels/cuda/act_round.cuh b/zig/kernels/cuda/act_round.cuh new file mode 100644 index 0000000000..13bb103805 --- /dev/null +++ b/zig/kernels/cuda/act_round.cuh @@ -0,0 +1,116 @@ +// MLX's bf16-op-by-op activations (act_rows.cu's tf_geglu_rows, tf_swiglu_rows_round); `keep`: the rounding points. + +#pragma once + +#include +#include + +namespace tf_act_round { + +// mlx_lm's gelu_approx(x) * up as MLX's bf16 ops (Metal gemma4/glue.metal's gemma_geglu): its rounding points. +enum : int { CUBE = 1, CUBE_SCALED = 2, INNER = 4, INNER_SCALED = 8, TANH = 16, ONE_PLUS = 32, GELU = 64 }; +constexpr int GEGLU_ALL = 127; +// tf_draft_swiglu's bf16 steps (Metal core/draft_ops.metal, gemma4's DFlash drafter): its rounding points. +enum : int { EXP = 1, DENOM = 2, RECIP = 4, SIGMOID = 8, SILU = 16 }; +constexpr int SWIGLU_ALL = 31; +// MLX's logit soft cap cap * tanh(x / cap) as bf16 ops (Metal gemma4/glue.metal's gemma_softcap): its rounding points. +enum : int { CAP_DIV = 1, CAP_TANH = 2 }; +constexpr int SOFTCAP_ALL = 3; + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// v rounded to bf16 where `keep` holds the point p (a constant in the served kernels, so the test folds away). +__device__ __forceinline__ float point(float v, int keep, int p) { return (keep & p) != 0 ? bf(v) : v; } + +// 0.5 x (1 + tanh(c0 (x + c1 x^3))) up, each step rounded to bf16 (x^3 exact in fp32 for a bf16 x; c0, c1 as bf16). +__device__ __forceinline__ float geglu(float x, float up, int keep) { + const float c0 = bf(0.7978845608028654f), c1 = bf(0.044715f); + const float cube = point(__fmul_rn(__fmul_rn(x, x), x), keep, CUBE); + const float inner = point(__fadd_rn(x, point(__fmul_rn(c1, cube), keep, CUBE_SCALED)), keep, INNER); + const float th = point(tanhf(point(__fmul_rn(c0, inner), keep, INNER_SCALED)), keep, TANH); + const float half_x = bf(__fmul_rn(0.5f, x)); // exact for a bf16 x: no point to drop + const float gl = point(__fmul_rn(half_x, point(__fadd_rn(1.f, th), keep, ONE_PLUS)), keep, GELU); + return __fmul_rn(gl, up); +} + +// silu(g) up as sigmoid from exp(|g|): 1 / (1 + e) for g < 0, else 1 - that; each step rounded to bf16. +__device__ __forceinline__ float swiglu(float g, float up, int keep) { + const float e = point(expf(fabsf(g)), keep, EXP); + const float inv = point(__fdiv_rn(1.f, point(__fadd_rn(1.f, e), keep, DENOM)), keep, RECIP); + const float sig = g < 0.f ? inv : point(__fsub_rn(1.f, inv), keep, SIGMOID); + return __fmul_rn(point(__fmul_rn(g, sig), keep, SILU), up); +} + +// bf16(bf16(tanh(bf16(x / cap))) * cap): the division and product correctly rounded, tanhf within its 2 ulps. +__device__ __forceinline__ float softcap(float x, float cap, int keep) { + const float t = point(__fdiv_rn(x, cap), keep, CAP_DIV); + return __fmul_rn(point(tanhf(t), keep, CAP_TANH), cap); +} + +__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; +} + +// act_rows.cu's group_sum8: a thread's 8 stored values in order, then the group's 8 lanes in a fixed tree. +__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) { + float f[8]; + unpack8(packed, f); + float s = 0.f; + if (live) +#pragma unroll + for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4)); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2)); + return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1)); +} + +// out[r, i] = bf16(act(gate[r * g_ld + i], up[r * u_ld + i])) for i < n (n % 64 == 0), GEGLU or not; xs its sums. +template +__device__ __forceinline__ void rows(const __nv_bfloat16* __restrict__ gate, long long g_ld, + const __nv_bfloat16* __restrict__ up, long long u_ld, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, int n, int keep) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values + const bool live = 8 * c < n; + float g[8], u[8], o[8]; + if (live) { + unpack8(*reinterpret_cast(gate + r * g_ld + 8 * c), g); + unpack8(*reinterpret_cast(up + r * u_ld + 8 * c), u); + } +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = live ? (GEGLU ? geglu(g[i], u[i], keep) : swiglu(g[i], u[i], keep)) : 0.f; + const uint4 op = pack8(o); + if (live) *reinterpret_cast(out + r * n + 8 * c) = op; + if (xs != nullptr) { + const float s = group_sum8(op, live); + if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s; + } +} + +// out[r, i] = softcap(x[r, i]) for i < n (n % 8 == 0), rows packed n wide (out may be x); 8 values a thread. +__device__ __forceinline__ void softcap_rows(const __nv_bfloat16* x, __nv_bfloat16* out, int n, float cap, int keep) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; + if (8 * c >= n) return; + float v[8]; + unpack8(*reinterpret_cast(x + r * n + 8 * c), v); +#pragma unroll + for (int i = 0; i < 8; ++i) v[i] = softcap(v[i], cap, keep); + *reinterpret_cast(out + r * n + 8 * c) = pack8(v); +} + +} // namespace tf_act_round diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index a10b0ed2c9..3563f1f021 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/zig/kernels/cuda/act_rows.cu @@ -141,3 +141,26 @@ extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat } *reinterpret_cast(out + r * d + 8 * c) = pack8(y); } + +// MLX's bf16-op-by-op activations (act_round.cuh), the same rows and sums as tf_swiglu_rows. +#include "act_round.cuh" + +// out = bf16(gelu_tanh(gate) * up) as mlx_lm's geglu in MLX's bf16 ops (Gemma 4's MLP); xs (or null) out's sums. +extern "C" __global__ void __launch_bounds__(NT) tf_geglu_rows(const __nv_bfloat16* gate, long long g_ld, + const __nv_bfloat16* up, long long u_ld, + __nv_bfloat16* out, float* xs, int n) { + tf_act_round::rows(gate, g_ld, up, u_ld, out, xs, n, tf_act_round::GEGLU_ALL); +} + +// out = bf16(silu(gate) * up) in MLX's bf16 steps (tf_draft_swiglu's, Gemma 4's DFlash drafter); xs (or null) its sums. +extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows_round(const __nv_bfloat16* gate, long long g_ld, + const __nv_bfloat16* up, long long u_ld, + __nv_bfloat16* out, float* xs, int n) { + tf_act_round::rows(gate, g_ld, up, u_ld, out, xs, n, tf_act_round::SWIGLU_ALL); +} + +// out = bf16(cap * tanh(x / cap)) in MLX's bf16 steps: a soft-capped head's logits (final_logit_softcapping). +extern "C" __global__ void __launch_bounds__(NT) tf_softcap_rows(const __nv_bfloat16* x, __nv_bfloat16* out, int n, + float cap) { + tf_act_round::softcap_rows(x, out, n, cap, tf_act_round::SOFTCAP_ALL); +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index e0d3bb2d51..de86c047b5 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -63,6 +63,45 @@ extern "C" __global__ void tf_ctl_half_dot(const __nv_bfloat16* X, int K, float* if (lane == 0) Y[row] = __half2float(s); } +// 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, + int keep) { + if (geglu) tf_act_round::rows(gate, g_ld, up, u_ld, out, xs, n, keep); + else tf_act_round::rows(gate, g_ld, up, u_ld, out, xs, n, keep); +} + +// Gate 3: act_rows.cu's tf_softcap_rows keeping only the rounding points in `keep`; the exact check must go red. +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); +} + +// Gate 3: embed_rows.cu's scaled MLX rows keeping the rounding points in `keep`; the exact check must go red. +#include "embed_affine.cuh" + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_embed_affine(const int* ids, const uint32_t* w, + const __nv_bfloat16* s, const __nv_bfloat16* b, + __nv_bfloat16* out, int d, int bits, int group, + float scale, int keep) { + tf_embed_affine::rows(ids, w, s, b, out, d, bits, group, scale, keep); +} + +// Gate 3: qk_rows.cu's tf_qk_rows_round keeping only the options in `keep`; the exact check must go red. +#include "qk_round.cuh" + +extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_ctl_qk_round(const tf_qk::QkSrc in, const tf_qk::QkW nw, + 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, int keep) { + tf_qk_round::rows(in, nw, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps, keep); +} + // Gate 5: bounded to 256 threads a block, which the registry says its launcher gives 512. extern "C" __global__ void __launch_bounds__(256) tf_ctl_bounded(const float* X, float* Y) { const int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -85,3 +124,66 @@ extern "C" __global__ void tf_ctl_heavy(const float* X, float* Y) { for (int j = 0; j < 160; ++j) s += v[j] * v[159 - j]; Y[t] = s; } + +// Gate 3: window_rows.cu's kernels at head dims 256 and 512, each with one planted fault (window_rows.cuh's Fault). +#include "window_rows.cuh" + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_part_edge(const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, const int* slot, + float* po, float* pm, float* pl, int h, int hk, int win, + long long ring, float scale) { + tf_wrows::part<256, tf_wrows::EDGE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_rescale(const float* po, const float* pm, const float* pl, + __nv_bfloat16* out, int h, int hk, int chunks, + const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, const int* slot, + int win, long long ring, float scale) { + tf_wrows::merge<256, tf_wrows::RESCALE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_trunc(const float* po, const float* pm, const float* pl, + __nv_bfloat16* out, int h, int hk, int chunks, + const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, const int* slot, + int win, long long ring, float scale) { + tf_wrows::merge<256, tf_wrows::TRUNC>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT, 1) tf_ctl_wrows_prompt_edge(const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, __nv_bfloat16* out, + const int* hi, int p0, int W, int H, int HK, int G, + float scale, int win, int ring) { + tf_wrows::prompt::kernel<256, true, tf_wrows::EDGE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT, 1) tf_ctl_wrows_prompt_last(const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, __nv_bfloat16* out, + const int* hi, int p0, int W, int H, int HK, int G, + float scale, int win, int ring) { + tf_wrows::prompt::kernel<256, true, tf_wrows::LAST>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +extern "C" __global__ void __launch_bounds__(tf_wrows::NT) tf_ctl_wrows_merge_rescale_512(const float* po, const float* pm, const float* pl, + __nv_bfloat16* out, int h, int hk, int chunks, + const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, const int* slot, + int win, long long ring, float scale) { + tf_wrows::merge<512, tf_wrows::RESCALE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale); +} + +// Gate 3: norms.cu's MLX tail with one rounding point dropped each; its test's exact rounding-point check must go red. +#include "norms_tail.cuh" + +#define TF_CTL_TAIL(NAME, DROP) \ + extern "C" __global__ void __launch_bounds__(tf_rms_tail::NT) NAME( \ + const __nv_bfloat16* x, const __nv_bfloat16* r, const __nv_bfloat16* wb, const __nv_bfloat16* sc, \ + 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); \ + } +TF_CTL_TAIL(tf_ctl_tail_drop_branch_row, tf_rms_tail::BRANCH_ROW) +TF_CTL_TAIL(tf_ctl_tail_drop_branch_out, tf_rms_tail::BRANCH_OUT) +TF_CTL_TAIL(tf_ctl_tail_drop_sum, tf_rms_tail::SUM) +TF_CTL_TAIL(tf_ctl_tail_drop_row, tf_rms_tail::ROW) +#undef TF_CTL_TAIL diff --git a/zig/kernels/cuda/embed_affine.cuh b/zig/kernels/cuda/embed_affine.cuh new file mode 100644 index 0000000000..23646fd4e8 --- /dev/null +++ b/zig/kernels/cuda/embed_affine.cuh @@ -0,0 +1,44 @@ +// MLX affine token rows with a bf16 scale (embed_rows.cu's tf_embed_rows_affine); `keep`: its rounding point. +#pragma once + +#include +#include + +#include "affine_codes.cuh" + +namespace tf_embed_affine { + +// keep's bit: the dequantized value rounded to bf16 before the scale (MLX's dequantize, then its multiply). +constexpr int ROUND = 1; +// A control's only: the scale folded into the group's scale and bias, q (s scale) + b scale, rounded once. +constexpr int FOLD = 2; + +// Row ids[blockIdx.x], 8 values a thread: bf16(bf16(fma(q, s, b)) * scale), q `bits` wide in MLX's words. +__device__ __forceinline__ void rows(const int* __restrict__ ids, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ s, const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ out, int d, int bits, int group, float scale, int keep) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; + if (8 * c >= d) return; + const int64_t tok = ids[r]; + const uint32_t* row = w + tok * (d / 32) * bits; + const int64_t g = tok * (d / group) + 8 * c / group; + float sc = __bfloat162float(s[g]), bi = __bfloat162float(b[g]); + if (keep & FOLD) sc = __fmul_rn(sc, scale), bi = __fmul_rn(bi, scale); + uint4 o; + __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&o); +#pragma unroll + for (int i = 0; i < 4; ++i) { + float v[2]; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const float e = __fmaf_rn(static_cast(tf_affine8::code_at(row, 8 * c + 2 * i + j, bits)), sc, bi); + if (keep & FOLD) v[j] = e; + else v[j] = __fmul_rn((keep & ROUND) ? __bfloat162float(__float2bfloat16_rn(e)) : e, scale); + } + h[i] = __floats2bfloat162_rn(v[0], v[1]); + } + *reinterpret_cast(out + r * d + 8 * c) = o; +} + +} // namespace tf_embed_affine diff --git a/zig/kernels/cuda/embed_rows.cu b/zig/kernels/cuda/embed_rows.cu index 664e6a41ec..cd317656d9 100644 --- a/zig/kernels/cuda/embed_rows.cu +++ b/zig/kernels/cuda/embed_rows.cu @@ -1,4 +1,4 @@ -// Token rows for every CUDA family (float64 reference: core/act_ref.zig's affine4): 4-bit affine or bf16 tables. +// Token rows for every CUDA family (float64 references: core/embed_ref.zig): MLX affine or bf16 tables. #include #include @@ -56,3 +56,16 @@ extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16_copies(const for (int k = 0; k < copies; ++k) *reinterpret_cast(out + (r * copies + k) * d + 8 * c) = value; } + +// MLX's affine rows at 2 to 6 or 8 bits with a bf16 scale (embed_affine.cuh). +#include "embed_affine.cuh" + +// out[r] = bf16(bf16(q * s + b) * scale) of token ids[r]: q `bits` wide, s and b a group's; scale 1: MLX's dequantize. +extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine(const int* __restrict__ ids, + const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ s, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ out, int d, int bits, + int group, float scale) { + tf_embed_affine::rows(ids, w, s, b, out, d, bits, group, scale, tf_embed_affine::ROUND); +} diff --git a/zig/kernels/cuda/norms.cu b/zig/kernels/cuda/norms.cu index 73902985bd..fa214406c6 100644 --- a/zig/kernels/cuda/norms.cu +++ b/zig/kernels/cuda/norms.cu @@ -503,3 +503,20 @@ extern "C" __global__ void __launch_bounds__(256) tf_gated_rms_rows_256(const __ float* xs, int heads, float eps) { gated_rms<32>(y, y_ld, z, z_ld, w, out, xs, heads, eps); } + +// The bf16 residual's tails (norms_tail.cuh): a norm on the branch before the add, a bf16 scalar on the sum. +#include "norms_tail.cuh" + +// h = bf16(r + bf16(rms(x) wb)) [* sc], y = bf16(rms(h) w), xs y's 64-value sums: each norm rounded once. +extern "C" __global__ void __launch_bounds__(tf_rms_tail::NT) tf_add_rms_rows_bf16h_tail( + const __nv_bfloat16* x, const __nv_bfloat16* r, const __nv_bfloat16* wb, const __nv_bfloat16* sc, + 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); +} + +// The same rows with MLX's bf16 steps: each norm bf16(w * bf16(x * inv)), the sum rounded before the scalar. +extern "C" __global__ void __launch_bounds__(tf_rms_tail::NT) tf_add_rms_rows_bf16h_tail_round( + const __nv_bfloat16* x, const __nv_bfloat16* r, const __nv_bfloat16* wb, const __nv_bfloat16* sc, + 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); +} diff --git a/zig/kernels/cuda/norms_tail.cuh b/zig/kernels/cuda/norms_tail.cuh new file mode 100644 index 0000000000..e27c642d5e --- /dev/null +++ b/zig/kernels/cuda/norms_tail.cuh @@ -0,0 +1,164 @@ +// bf16 residual tails (norms.cu's tf_add_rms_rows_bf16h_tail*): branch norm, add, bf16 scalar, output norm and sums. + +#pragma once + +#include +#include + +namespace tf_rms_tail { + +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 branch's normalized row, the branch's output, the sum before the scalar, the output's row. +enum : int { BRANCH_ROW = 1, BRANCH_OUT = 2, SUM = 4, ROW = 8 }; +// One rounding a norm, bf16 rows between the steps. +constexpr int ONCE = BRANCH_OUT | SUM; +// MLX's bf16 op by op (Metal gemma4/glue.metal's tails): bf16(w * bf16(x * inv)) for each norm. +constexpr int MLX = BRANCH_ROW | BRANCH_OUT | SUM | ROW; + +__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; +} + +// A normalized value times its weight: w * bf16(v * inv) where R holds the row's point P, else (v * inv) * w. +template +__device__ __forceinline__ float normed(float v, float inv, float w) { + if constexpr ((R & P) != 0) return __fmul_rn(w, bf(__fmul_rn(v, inv))); + else return __fmul_rn(__fmul_rn(v, inv), w); +} + +__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 / width + eps), IEEE division and square root: norms.cu's inv_rms_v tree over the 8 warps. +template +__device__ __forceinline__ float inv_rms(float part, int width, float eps) { + __shared__ float warps[NT / 32]; + 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); + const float r = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(s, static_cast(width)), eps))); + if constexpr (!ONCE_ONLY) __syncthreads(); // `warps` is read again by the output's norm + return r; +} + +// A thread's 8 stored outputs summed in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). +__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) { + float f[8]; + unpack8(packed, f); + float s = 0.f; + if (live) +#pragma unroll + for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4)); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2)); + return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1)); +} + +// b = wb ? bf16(rms(x) wb) : x, h = bf16(sc ? (r + b) sc : r + b) (if r, wb or sc), y = bf16(rms(h) w), xs its sums. +template +__device__ __forceinline__ void tail(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ r, + const __nv_bfloat16* __restrict__ wb, const __nv_bfloat16* __restrict__ sc, + const __nv_bfloat16* __restrict__ w, __nv_bfloat16* __restrict__ h, + __nv_bfloat16* __restrict__ y, float* __restrict__ xs, int width, float eps) { + const int64_t at = static_cast(blockIdx.x) * width; + const int chunks = width / 8; + float v[MOST][8]; + uint4 xp[MOST], rp[MOST], bp[MOST], wp[MOST]; +#pragma unroll + for (int j = 0; j < MOST; ++j) { // every load issued before any is used + const int c = threadIdx.x + j * NT; + if (c < chunks) { + xp[j] = reinterpret_cast(x + at)[c]; + if (r != nullptr) rp[j] = reinterpret_cast(r + at)[c]; + if (wb != nullptr) bp[j] = reinterpret_cast(wb)[c]; + wp[j] = reinterpret_cast(w)[c]; + } else xp[j] = rp[j] = bp[j] = wp[j] = make_uint4(0, 0, 0, 0); // past the row: zeros that add nothing + } +#pragma unroll + for (int j = 0; j < MOST; ++j) unpack8(xp[j], v[j]); + if (wb != nullptr) { + 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); // zeros past the row add nothing + const float inv = inv_rms(part, width, eps); +#pragma unroll + for (int j = 0; j < MOST; ++j) { + if (threadIdx.x + j * NT >= chunks) continue; // past the row v stays 0 + float bw[8]; + unpack8(bp[j], bw); +#pragma unroll + for (int i = 0; i < 8; ++i) v[j][i] = point(normed(v[j][i], inv, bw[i])); + } + } + const float scale = sc != nullptr ? __bfloat162float(*sc) : 1.f; + if (r != nullptr || wb != nullptr || sc != nullptr) { +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int c = threadIdx.x + j * NT; + if (c < chunks) { + if (r != nullptr) { + float rv[8]; + unpack8(rp[j], rv); +#pragma unroll + for (int i = 0; i < 8; ++i) v[j][i] = __fadd_rn(rv[i], v[j][i]); + } + if (sc != nullptr) +#pragma unroll + for (int i = 0; i < 8; ++i) v[j][i] = __fmul_rn(point(v[j][i]), scale); + const uint4 hp = pack8(v[j]); + reinterpret_cast(h + at)[c] = hp; + unpack8(hp, v[j]); // the output's norm reads h as stored + } + } + } + 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); + const float inv = inv_rms(part, width, eps); +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int c = threadIdx.x + j * NT; + const bool live = c < chunks; + float wv[8], o[8]; + unpack8(wp[j], wv); +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = normed(v[j][i], inv, wv[i]); + const uint4 yp = pack8(o); + if (live) reinterpret_cast(y + at)[c] = yp; + if (xs != nullptr && j * NT < chunks) { // whole warps of whole groups: width % 64 == 0 + const float s = group_sum8(yp, live); + if (live && (c & 7) == 0) xs[static_cast(blockIdx.x) * (width / 64) + c / 8] = s; + } + } +} + +} // namespace tf_rms_tail diff --git a/zig/kernels/cuda/qk_round.cuh b/zig/kernels/cuda/qk_round.cuh new file mode 100644 index 0000000000..299fdf56f4 --- /dev/null +++ b/zig/kernels/cuda/qk_round.cuh @@ -0,0 +1,84 @@ +// qk_rows.cu's tf_qk_rows_once and tf_qk_rows_round: head dims to 512, a unit v norm and MLX's rounding points. + +#pragma once + +#include "qk_rows.cuh" + +namespace tf_qk_round { + +using namespace tf_qk; + +constexpr int WIDE = 512; // head_dim at most + +// The options: the normalized row rounded before its weight, the weighted row before RoPE, v by a unit RMS norm. +enum : int { ROW = 1, OUT = 2, VNORM = 4 }; +// MLX's bf16 op by op (Metal gemma4/glue.metal's gemma_norm_rope): bf16(w * bf16(x * inv)), then RoPE. +constexpr int MLX = ROW | OUT; + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// 1 / sqrt(sum / dim + eps) of a head's squares: each warp's xor tree, then the warps in order (tf_qk_rows' order). +__device__ __forceinline__ float inv_rms(float x, int dim, float eps) { + __shared__ float warps[WIDE / 32]; + float part = __fmul_rn(x, x); + 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(__fdiv_rn(s, static_cast(dim)), eps))); +} + +// A normed value: w * bf16(x * inv) with ROW, else (x * inv) * w; rounded to bf16 with OUT. +__device__ __forceinline__ float normed(float x, float inv, float w, int pts) { + const float y = (pts & ROW) ? __fmul_rn(w, bf(__fmul_rn(x, inv))) : __fmul_rn(__fmul_rn(x, inv), w); + return (pts & OUT) ? bf(y) : y; +} + +// tf_qk_rows with head dims to 512 and the options in `pts`: a block a (row, head), a thread an element. +__device__ __forceinline__ void rows(const QkSrc& in, const QkW& nw, 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, int pts) { + 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 float x = __bfloat162float(src[d]); + 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, or by a unit RMS norm + const int k = hd - h - hk; + const __nv_bfloat16 out = (pts & VNORM) ? __float2bfloat16_rn(__fmul_rn(x, inv_rms(x, dim, eps))) : src[d]; + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + d] = out; + if (vc != nullptr) vc[(at * hk + k) * dim + d] = out; + return; + } + float y = x; + float yp = 0.f; + const int half = rot / 2; + const int partner = d < half ? d + half : d - half; + if (nw.q != nullptr) { + const float rstd = inv_rms(x, dim, eps); + y = normed(x, rstd, wv(nw, hd >= h, d), pts); + if (d < rot) yp = normed(__bfloat162float(src[partner]), rstd, wv(nw, hd >= h, partner), pts); + } else if (d < rot) { + yp = __bfloat162float(src[partner]); + } + if ((cosv != nullptr || inv != nullptr) && d < rot) { + 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)); + } + const __nv_bfloat16 out = __float2bfloat16_rn(y); + if (hd < h) { + dst.q[r * dst.q_ld + hd * dst.q_hs + d] = out; + } else { + const int k = hd - h; + 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_round diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu index 395b1ca677..b56c7ff642 100644 --- a/zig/kernels/cuda/qk_rows.cu +++ b/zig/kernels/cuda/qk_rows.cu @@ -2,65 +2,7 @@ // A block a (row, head), a thread an element; a head's norm sums in one fixed order, never depending on the row count. -#include -#include - -namespace tf_qk { - -constexpr int MOST = 256; // head_dim at most - -// Where a row's heads come from: q, k and v (null: none) at row stride *_ld and head stride *_hs (elements). -struct QkSrc { - const __nv_bfloat16* q; - long long q_ld, q_hs; - const __nv_bfloat16* k; - long long k_ld, k_hs; - const __nv_bfloat16* v; - long long v_ld, v_hs; -}; - -// The q and k norm weights: fp32 or bf16 (null q: no norm). -struct QkW { - const void* q; - const void* k; - int bf16; -}; - -__device__ __forceinline__ float wv(const QkW& w, bool is_k, int i) { - const void* p = is_k ? w.k : w.q; - return w.bf16 ? __bfloat162float(static_cast(p)[i]) : static_cast(p)[i]; -} - -// cos and sin of p's angle on rotated dim d: from tables [positions, rot], or computed from inv [rot / 2]. -__device__ __forceinline__ void rope_at(const float* cosv, const float* sinv, const float* inv, int p, int rot, int d, - float& c, float& sn) { - if (inv != nullptr) { - const float a = __fmul_rn(static_cast(p), inv[d < rot / 2 ? d : d - rot / 2]); - c = cosf(a); - sn = sinf(a); - } else { - c = cosv[static_cast(p) * rot + d]; - sn = sinv[static_cast(p) * rot + d]; - } -} - -// Where a row's q, k and v (null: none) go, at row stride *_ld and head stride *_hs (elements). -struct QkDst { - __nv_bfloat16* q; - long long q_ld, q_hs; - __nv_bfloat16* k; - long long k_ld, k_hs; - __nv_bfloat16* v; - long long v_ld, v_hs; -}; - -__device__ __forceinline__ const __nv_bfloat16* head_src(const QkSrc& s, int64_t r, int hd, int h, int hk) { - if (hd < h) return s.q + r * s.q_ld + hd * s.q_hs; - if (hd < h + hk) return s.k + r * s.k_ld + (hd - h) * s.k_hs; - return s.v + r * s.v_ld + (hd - h - hk) * s.v_hs; -} - -} // namespace tf_qk +#include "qk_rows.cuh" using namespace tf_qk; @@ -192,3 +134,25 @@ extern "C" __global__ void __launch_bounds__(128) tf_qk_rows_w(const QkSrc in, c } } } + +// tf_qk_rows to head dim 512 (qk_round.cuh), and the options it alone has: a unit v norm and MLX's rounding points. +#include "qk_round.cuh" + +// tf_qk_rows' arithmetic at head dims to 512; vnorm: v by a unit RMS norm, bf16(x * inv), else as is. +extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_qk_rows_once(const QkSrc in, const QkW nw, 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, int vnorm) { + tf_qk_round::rows(in, nw, cosv, sinv, inv, rot, pos, slot, dst, kc, vc, len, h, hk, eps, vnorm ? tf_qk_round::VNORM : 0); +} + +// As tf_qk_rows_once with MLX's rounding points: q, k = bf16(w * bf16(x * inv)) before RoPE (Metal gemma_norm_rope). +extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_qk_rows_round(const QkSrc in, const QkW nw, 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, int vnorm) { + 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)); +} diff --git a/zig/kernels/cuda/qk_rows.cuh b/zig/kernels/cuda/qk_rows.cuh new file mode 100644 index 0000000000..f655175fb8 --- /dev/null +++ b/zig/kernels/cuda/qk_rows.cuh @@ -0,0 +1,63 @@ +// qk_rows.cu's shared parts: where heads come from and go, the norm weights, cos and sin. + +#pragma once + +#include +#include + +namespace tf_qk { + +constexpr int MOST = 256; // tf_qk_rows' head_dim at most + +// Where a row's heads come from: q, k and v (null: none) at row stride *_ld and head stride *_hs (elements). +struct QkSrc { + const __nv_bfloat16* q; + long long q_ld, q_hs; + const __nv_bfloat16* k; + long long k_ld, k_hs; + const __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +// The q and k norm weights: fp32 or bf16 (null q: no norm). +struct QkW { + const void* q; + const void* k; + int bf16; +}; + +__device__ __forceinline__ float wv(const QkW& w, bool is_k, int i) { + const void* p = is_k ? w.k : w.q; + return w.bf16 ? __bfloat162float(static_cast(p)[i]) : static_cast(p)[i]; +} + +// cos and sin of p's angle on rotated dim d: from tables [positions, rot], or computed from inv [rot / 2]. +__device__ __forceinline__ void rope_at(const float* cosv, const float* sinv, const float* inv, int p, int rot, int d, + float& c, float& sn) { + if (inv != nullptr) { + const float a = __fmul_rn(static_cast(p), inv[d < rot / 2 ? d : d - rot / 2]); + c = cosf(a); + sn = sinf(a); + } else { + c = cosv[static_cast(p) * rot + d]; + sn = sinv[static_cast(p) * rot + d]; + } +} + +// Where a row's q, k and v (null: none) go, at row stride *_ld and head stride *_hs (elements). +struct QkDst { + __nv_bfloat16* q; + long long q_ld, q_hs; + __nv_bfloat16* k; + long long k_ld, k_hs; + __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +__device__ __forceinline__ const __nv_bfloat16* head_src(const QkSrc& s, int64_t r, int hd, int h, int hk) { + if (hd < h) return s.q + r * s.q_ld + hd * s.q_hs; + if (hd < h + hk) return s.k + r * s.k_ld + (hd - h) * s.k_hs; + return s.v + r * s.v_ld + (hd - h - hk) * s.v_hs; +} + +} // namespace tf_qk diff --git a/zig/kernels/cuda/window_rows.cu b/zig/kernels/cuda/window_rows.cu new file mode 100644 index 0000000000..46887ad334 --- /dev/null +++ b/zig/kernels/cuda/window_rows.cu @@ -0,0 +1,72 @@ +// window_attention's rows at head dims 256 and 512 (window_rows.cuh): decode partials and merge, prompt tiles. + +#include "window_rows.cuh" +#include "window_rows512.cuh" + +using namespace tf_wrows; + +extern "C" __global__ void __launch_bounds__(NT) tf_wrows_part_256(const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, + const int* slot, float* po, float* pm, float* pl, + int h, int hk, int win, long long ring, float scale) { + part<256, NONE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_wrows_merge_256(const float* po, const float* pm, const float* pl, + __nv_bfloat16* out, int h, int hk, int chunks, + const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, + const int* slot, int win, long long ring, + float scale) { + merge<256, NONE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256(const __nv_bfloat16* q, + const __nv_bfloat16* kc, + const __nv_bfloat16* vc, __nv_bfloat16* out, + const int* hi, int p0, int W, int H, int HK, + int G, float scale, int win, int ring) { + prompt::kernel<256, false, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256_p32(const __nv_bfloat16* q, + const __nv_bfloat16* kc, + const __nv_bfloat16* vc, + __nv_bfloat16* out, const int* hi, int p0, + int W, int H, int HK, int G, float scale, + int win, int ring) { + prompt::kernel<256, true, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_wrows_part_512(const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, + const int* slot, float* po, float* pm, float* pl, + int h, int hk, int win, long long ring, float scale) { + part_staged<512, NONE>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_wrows_merge_512(const float* po, const float* pm, const float* pl, + __nv_bfloat16* out, int h, int hk, int chunks, + const __nv_bfloat16* q, const __nv_bfloat16* kc, + const __nv_bfloat16* vc, const int* pos, + const int* slot, int win, long long ring, + float scale) { + merge<512, NONE>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_512(const __nv_bfloat16* q, + const __nv_bfloat16* kc, + const __nv_bfloat16* vc, __nv_bfloat16* out, + const int* hi, int p0, int W, int H, int HK, + int G, float scale, int win, int ring) { + prompt::kernel<512, false, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_512_p32(const __nv_bfloat16* q, + const __nv_bfloat16* kc, + const __nv_bfloat16* vc, + __nv_bfloat16* out, const int* hi, int p0, + int W, int H, int HK, int G, float scale, + int win, int ring) { + prompt::kernel<512, true, NONE>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} diff --git a/zig/kernels/cuda/window_rows.cuh b/zig/kernels/cuda/window_rows.cuh new file mode 100644 index 0000000000..947668c4ca --- /dev/null +++ b/zig/kernels/cuda/window_rows.cuh @@ -0,0 +1,560 @@ +// window_attention at head dims 256 and 512: decode by 256-key chunks and a merge, prompt tiles (attention_range.zig). +#pragma once + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_wrows { + +constexpr int NT = 256; // threads a block +constexpr int NW = NT / 32; // warps a block +constexpr int CK = 256; // keys a decode chunk: chunk c holds keys [c * CK, c * CK + CK) +constexpr int GMAX = 16; // query heads a key head +constexpr int SP = GMAX + 1; // the decode scores' key stride (a bank further each key) +constexpr int BN = 32; // keys a prompt tile +constexpr unsigned FULL = 0xffffffffu; + +// Planted faults, compiled only into contract_controls.cu's controls; every served kernel has NONE. +enum Fault : int { + NONE = 0, + EDGE = 1, // the window's first key one too early + LAST = 2, // a prompt row's last key: its own position, the given one ignored + RESCALE = 3, // the merge adds the chunks' partials without exp(m_c - M) + TRUNC = 4, // the merge's output rounded toward zero +}; + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +__device__ __forceinline__ float lo16(uint32_t w) { return __uint_as_float(w << 16); } +__device__ __forceinline__ float hi16(uint32_t w) { return __uint_as_float(w & 0xffff0000u); } + +// A float64 value to bf16 in one rounding, ties to even: toward zero to fp32 with a sticky last bit, then nearest. +__device__ __forceinline__ __nv_bfloat16 d2bf(double x) { + if (x != x) return __float2bfloat16_rn(__int_as_float(0x7fc00000)); + float f = __double2float_rz(x); + if ((double)f != x && isfinite(f)) f = __int_as_float(__float_as_int(f) | 1); + return __float2bfloat16_rn(f); +} + +// q . k in float64, lanes over the head's dims, then a xor tree; every lane gets the sum. +template +__device__ __forceinline__ double dot64(const __nv_bfloat16* q, const __nv_bfloat16* k, int lane) { + double s = 0.0; +#pragma unroll 4 + for (int e = lane; e < D; e += 32) s += (double)__bfloat162float(q[e]) * (double)__bfloat162float(k[e]); + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(FULL, s, o); + return s; +} + +// A head over keys [lo, hi] in float64 for columns c0 + threadIdx.x + NT i < c1 (kv(key, 0): K row, 1: V); a block. +template +__device__ __noinline__ void slow_head(const __nv_bfloat16* qh, KV kv, int lo, int hi, double scale, + __nv_bfloat16* out, int c0, int c1) { + __shared__ double red[NW]; + __shared__ double pw[NW]; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + double m = -INFINITY; + for (int j = lo + warp; j <= hi; j += NW) m = fmax(m, dot64(qh, kv(j, 0), lane) * scale); + __syncthreads(); + if (lane == 0) red[warp] = m; + __syncthreads(); + m = -INFINITY; + for (int w = 0; w < NW; ++w) m = fmax(m, red[w]); + constexpr int CPT = D / NT; + double acc[CPT], l = 0.0; +#pragma unroll + for (int i = 0; i < CPT; ++i) acc[i] = 0.0; + for (int jb = lo; jb <= hi; jb += NW) { + const int j = jb + warp; + if (j <= hi) { + const double s = dot64(qh, kv(j, 0), lane) * scale; + if (lane == 0) pw[warp] = exp(s - m); + } + __syncthreads(); + const int nb = min(NW, hi - jb + 1); + for (int w = 0; w < nb; ++w) { + const double p = pw[w]; + l += p; + const __nv_bfloat16* v = kv(jb + w, 1); +#pragma unroll + for (int i = 0; i < CPT; ++i) { + const int c = c0 + threadIdx.x + NT * i; + if (c < c1) acc[i] += p * (double)__bfloat162float(v[c]); + } + } + __syncthreads(); + } +#pragma unroll + for (int i = 0; i < CPT; ++i) { + const int c = c0 + threadIdx.x + NT * i; + if (c < c1) out[c] = d2bf(acc[i] / l); + } +} + +// Decode partials, a block a (row, key head, chunk): the chunk's max, weights and o; empty chunks write m = -inf. +template +__device__ __forceinline__ void part(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, const int* __restrict__ pos, + const int* __restrict__ slot, float* __restrict__ po, float* __restrict__ pm, + float* __restrict__ pl, int h, int hk, int win, long long ring, float scale) { + constexpr int DPL = D / 32; // a lane's dims of q and k, and its columns of v + constexpr int HW = 32 / DPL; // query heads a warp takes: 32 q values, then 32 accumulators, a lane + constexpr int U = 4; // keys a warp loads at once + __shared__ float sp[CK * SP]; + __shared__ float mx[GMAX], ls[GMAX]; + __shared__ float ox[7 * 4 * 256]; // key runs 1.. partial o, (runs - 1) G D floats at most (G HW: runs 8) + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, c = blockIdx.z, chunks = gridDim.z, G = h / hk; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int p = pos[r]; + const int64_t base = (slot != nullptr ? static_cast(slot[r]) : 0) * ring; + const int lo = win > 0 ? max(p - win + (FAULT == EDGE ? 0 : 1), 0) : 0; + const int k0 = max((lo / CK + c) * CK, lo), k1 = min((lo / CK + c) * CK + CK - 1, p); + const int64_t part = (r * hk + kh) * chunks + c; // this block's partial: G heads of D values + if (k0 > k1) { + if (threadIdx.x < G) pm[part * G + threadIdx.x] = -INFINITY, pl[part * G + threadIdx.x] = 0.0f; + return; + } + const int nk = k1 - k0 + 1; + const int64_t r0 = k0 % ring; + // warp (hg, kg): heads hg HW.., keys [j0, j1), the kg-th of WG runs in key order + const int HG = (G + HW - 1) / HW, WG = max(1, NW / HG); + const int hg = warp % HG, kg = warp / HG, run = (nk + WG - 1) / WG; + const bool busy = warp < HG * WG; + const int j0 = min(kg * run, nk), j1 = min(j0 + run, nk); + const auto rowp = [&](const __nv_bfloat16* cache, int jj) { + int64_t row = r0 + jj; + if (row >= ring) row -= ring; + return cache + ((base + row) * hk + kh) * D + lane * DPL; + }; + const auto load = [&](const __nv_bfloat16* src, float (&x)[DPL]) { +#pragma unroll + for (int i = 0; i < DPL; i += 8) { + const uint4 u = *reinterpret_cast(src + i); + const uint32_t w[4] = {u.x, u.y, u.z, u.w}; +#pragma unroll + for (int t = 0; t < 4; ++t) x[i + 2 * t] = lo16(w[t]), x[i + 2 * t + 1] = hi16(w[t]); + } + }; + if (busy) { // scores: a lane's DPL dims in an fma chain, then a xor tree over the lanes + float qv[HW][DPL]; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = hg * HW + hh; + load(q + (r * h + kh * G + min(hd, G - 1)) * D + lane * DPL, qv[hh]); + } + for (int jj = j0; jj < j1; jj += U) { + float kf[U][DPL]; +#pragma unroll + for (int u = 0; u < U; ++u) + if (jj + u < j1) load(rowp(kc, jj + u), kf[u]); +#pragma unroll + for (int u = 0; u < U; ++u) { + if (jj + u >= j1) break; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + if (hg * HW + hh >= G) break; // the same for the whole warp + float s = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; ++i) s = __fmaf_rn(qv[hh][i], kf[u][i], s); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) s = __fadd_rn(s, __shfl_xor_sync(FULL, s, o)); + if (lane == 0) sp[(jj + u) * SP + hg * HW + hh] = __fmul_rn(s, scale); + } + } + } + } + __syncthreads(); + for (int hd = warp; hd < G; hd += NW) { // a warp a head: max, then weights and their sum (lanes, then a xor tree) + float m = -INFINITY; + for (int jj = lane; jj < nk; jj += 32) m = fmaxf(m, sp[jj * SP + hd]); + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(FULL, m, o)); + float l = 0.0f; + for (int jj = lane; jj < nk; jj += 32) { + const float s = sp[jj * SP + hd]; + const float e = s == -INFINITY ? 0.0f : expf(__fsub_rn(s, m)); + sp[jj * SP + hd] = e; + l = __fadd_rn(l, e); + } + for (int o = 16; o > 0; o >>= 1) l = __fadd_rn(l, __shfl_xor_sync(FULL, l, o)); + if (lane == 0) mx[hd] = m, ls[hd] = l; + } + __syncthreads(); + // o: a lane's DPL columns of heads hg HW.., an fma chain over its run's keys in order; then the runs in order + float acc[HW][DPL]; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) +#pragma unroll + for (int i = 0; i < DPL; ++i) acc[hh][i] = 0.0f; + if (busy) { + for (int jj = j0; jj < j1; jj += U) { + float vf[U][DPL]; +#pragma unroll + for (int u = 0; u < U; ++u) + if (jj + u < j1) load(rowp(vc, jj + u), vf[u]); +#pragma unroll + for (int u = 0; u < U; ++u) { + if (jj + u >= j1) break; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + if (hg * HW + hh >= G) break; + const float w = sp[(jj + u) * SP + hg * HW + hh]; +#pragma unroll + for (int i = 0; i < DPL; ++i) acc[hh][i] = __fmaf_rn(w, vf[u][i], acc[hh][i]); + } + } + } + if (kg > 0) +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = hg * HW + hh; + if (hd >= G) break; +#pragma unroll + for (int i = 0; i < DPL; ++i) ox[((kg - 1) * G + hd) * D + lane * DPL + i] = acc[hh][i]; + } + } + __syncthreads(); + if (busy && kg == 0) { +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = hg * HW + hh; + if (hd >= G) break; +#pragma unroll + for (int i = 0; i < DPL; ++i) { + float a = acc[hh][i]; + for (int k = 1; k < WG; ++k) a = __fadd_rn(a, ox[((k - 1) * G + hd) * D + lane * DPL + i]); + po[(part * G + hd) * D + lane * DPL + i] = a; + } + } + } + if (threadIdx.x < G) pm[part * G + threadIdx.x] = mx[threadIdx.x], pl[part * G + threadIdx.x] = ls[threadIdx.x]; +} + +// out = sum_c a_c o_c / sum_c a_c l_c, a_c = expf(m_c - M), chunks in order; NaN or inf: the head again in float64. +template +__device__ __forceinline__ void merge(const float* __restrict__ po, const float* __restrict__ pm, + const float* __restrict__ pl, __nv_bfloat16* __restrict__ out, int h, int hk, + int chunks, const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, + const int* __restrict__ pos, const int* __restrict__ slot, int win, + long long ring, float scale) { + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, G = h / hk; + const int o = blockIdx.z * NT + threadIdx.x; // NT divides D: the block's outputs are one head's + const int hd = o / D, e = o % D; + const int64_t part0 = (r * hk + kh) * chunks; + float m = -INFINITY; + for (int c = 0; c < chunks; ++c) m = fmaxf(m, pm[(part0 + c) * G + hd]); + float acc = 0.0f, l = 0.0f; + for (int c = 0; c < chunks; ++c) { + const float mc = pm[(part0 + c) * G + hd]; + if (mc == -INFINITY) continue; + const float a = FAULT == RESCALE ? 1.0f : expf(__fsub_rn(mc, m)); + acc = __fmaf_rn(a, po[((part0 + c) * G + hd) * D + e], acc); + l = __fmaf_rn(a, pl[(part0 + c) * G + hd], l); + } + const float x = __fdiv_rn(acc, l); + __nv_bfloat16* oh = out + (r * h + kh * G + hd) * D; + oh[e] = FAULT == TRUNC ? __float2bfloat16_rz(x) : __float2bfloat16_rn(x); + if (__syncthreads_or(!isfinite(x))) { + const int p = pos[r]; + const int64_t base = (slot != nullptr ? static_cast(slot[r]) : 0) * ring; + const auto kv = [&](int key, int which) { + return (which == 0 ? kc : vc) + ((base + key % ring) * hk + kh) * D; + }; + const int c0 = blockIdx.z * NT % D; + slow_head(q + (r * h + kh * G + hd) * D, kv, win > 0 ? max(p - win + 1, 0) : 0, p, (double)scale, oh, c0, + c0 + NT); + } +} + +namespace prompt { + +__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(qmm_frag::smem(p))); +} + +template +__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +// The bf16 pair's fp32 values. +__device__ __forceinline__ float2 unpack(uint32_t w) { return make_float2(lo16(w), hi16(w)); } + +template +__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } + +template +struct Cfg { + static constexpr int SL = 128; // dims a warp's slice: its share of q . k, its columns of o + static constexpr int CS = D / SL; // warps sharing 16 rows + static constexpr int K16 = SL / 16, NB = SL / 8; // a slice's mma k-steps and n8 output tiles + static constexpr int NG = NW / CS; // 16-row groups a block + static constexpr int C = D / 8; // 16-byte chunks a key row + static constexpr int NS = D == 256 ? 4 : 2; // staging slots of BN keys (K or V) + static constexpr int SLOT = BN * C; // uint4 a slot + static constexpr int XS = NW * 16 * 32; // floats: the warps' partial scores, [warp][element][lane] + static constexpr int SMEM = NS * SLOT * 16 + XS * 4; // dynamic shared bytes +}; + +// Prompt rows p0..: row i sees keys [p0 + i - win + 1, hi[i]]; CS warps a 16-row band, slices summed; P32: P halves. +template +__device__ __forceinline__ bool rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, + const int* __restrict__ hi, int p0, int W, int H, int HK, int G, float scale, + int win, int ring) { + using K = Cfg; + constexpr int CS = K::CS, NG = K::NG, C = K::C, NS = K::NS, SL = K::SL, K16 = K::K16, NB = K::NB; + extern __shared__ uint4 wsmem[]; + uint4* slots = wsmem; + float* xs = reinterpret_cast(wsmem + NS * K::SLOT); + __shared__ int kmax; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g = lane / 4, qd = lane % 4, mt = lane / 8; + const int grp = warp / CS, cs = warp % CS; + const int HB = G % NG == 0 ? NG : 1, RB = NG / HB; // heads and 16-row bands a block + const int hk = blockIdx.y / (G / HB), hg = blockIdx.y % (G / HB); + const int head = hk * G + hg * HB + grp % HB; + const int row0 = blockIdx.x * 16 * RB, wrow = row0 + (grp / HB) * 16; + const int r0 = wrow + g, r1 = r0 + 8; + const int pos0 = p0 + r0, pos1 = p0 + r1; + const auto first = [&](int p) { return win > 0 ? max(p - win + (FAULT == EDGE ? 0 : 1), 0) : 0; }; + const auto last = [&](int r) { return hi != nullptr && FAULT != LAST && r < W ? hi[r] : p0 + r; }; + const int lo0 = first(pos0), lo1 = first(pos1), hi0 = last(r0), hi1 = last(r1); + if (threadIdx.x == 0) kmax = 0; + __syncthreads(); + if (threadIdx.x < 16 * RB && row0 + (int)threadIdx.x < W) atomicMax(&kmax, last(row0 + threadIdx.x)); + __syncthreads(); + const int kfirst = first(p0 + row0), klast = kmax, limit = p0 + W; + const int t0 = kfirst / BN, tiles = klast / BN + 1 - t0, items = 2 * tiles; + auto stage = [&](int item) { // item 2t: tile t's keys, 2t + 1: its values + const __nv_bfloat16* src = item % 2 == 0 ? kc : vc; + const int key0 = (t0 + item / 2) * BN; + uint4* dst = slots + (item % NS) * K::SLOT; + for (int i = threadIdx.x; i < BN * C; i += NT) { + const int key = key0 + i / C; + const bool in = key >= kfirst && key <= klast && key < limit; // outside every row's keys: zeros + const int row = in ? (ring > 0 ? key % ring : key) : 0; + cp16z(dst + sw(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in); + } + }; + uint32_t qa[K16][4]; // this warp's SL dims of rows r0, r1 + { + const uint32_t* q0 = reinterpret_cast(q + ((int64_t)min(r0, W - 1) * H + head) * D + cs * SL); + const uint32_t* q1 = reinterpret_cast(q + ((int64_t)min(r1, W - 1) * H + head) * D + cs * SL); +#pragma unroll + for (int k = 0; k < K16; ++k) { + qa[k][0] = r0 < W ? q0[8 * k + qd] : 0u; + qa[k][1] = r1 < W ? q1[8 * k + qd] : 0u; + qa[k][2] = r0 < W ? q0[8 * k + qd + 4] : 0u; + qa[k][3] = r1 < W ? q1[8 * k + qd + 4] : 0u; + } + } +#pragma unroll + for (int i = 0; i < NS; ++i) { + if (i < items) stage(i); + commit(); + } + auto take = [&](int it) -> const uint4* { // item it's slot, once every thread's copy has landed + wait_group(); + __syncthreads(); + return slots + (it % NS) * K::SLOT; + }; + auto done = [&](int it) { // the slot is free: stage the item NS ahead into it + __syncthreads(); + if (it + NS < items) stage(it + NS); + commit(); + }; + float o[NB][4]; +#pragma unroll + for (int i = 0; i < NB; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.0f; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + for (int t = 0; t < tiles; ++t) { + float s[4][4]; + { + const uint4* sl = take(2 * t); +#pragma unroll + for (int k = 0; k < K16; ++k) { +#pragma unroll + for (int n = 0; n < 4; n += 2) { + uint32_t b[4]; + ldmatrix4(b, sl + sw(8 * (n + mt / 2) + lane % 8, cs * (SL / 8) + 2 * k + mt % 2)); + if (k == 0) { + mma0(s[n], qa[k], b[0], b[1]); + mma0(s[n + 1], qa[k], b[2], b[3]); + } else { + mma(s[n], qa[k], b[0], b[1]); + mma(s[n + 1], qa[k], b[2], b[3]); + } + } + } +#pragma unroll + for (int n = 0; n < 4; ++n) +#pragma unroll + for (int j = 0; j < 4; ++j) xs[(warp * 16 + 4 * n + j) * 32 + lane] = s[n][j]; + done(2 * t); + } + // the group's scores: its slices' partials in slice order +#pragma unroll + for (int n = 0; n < 4; ++n) +#pragma unroll + for (int j = 0; j < 4; ++j) { + float a = xs[(grp * CS * 16 + 4 * n + j) * 32 + lane]; +#pragma unroll + for (int c = 1; c < CS; ++c) a = __fadd_rn(a, xs[((grp * CS + c) * 16 + 4 * n + j) * 32 + lane]); + s[n][j] = a; + } + const int kb = (t0 + t) * BN + 2 * qd; + float tm0 = -INFINITY, tm1 = -INFINITY; + bool in[4][4]; +#pragma unroll + for (int n = 0; n < 4; ++n) { + const int c = kb + 8 * n; + in[n][0] = c >= lo0 && c <= hi0; + in[n][1] = c + 1 >= lo0 && c + 1 <= hi0; + in[n][2] = c >= lo1 && c <= hi1; + in[n][3] = c + 1 >= lo1 && c + 1 <= hi1; +#pragma unroll + for (int j = 0; j < 4; ++j) s[n][j] = in[n][j] ? __fmul_rn(s[n][j], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(s[n][0], s[n][1])); + tm1 = fmaxf(tm1, fmaxf(s[n][2], s[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(FULL, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(FULL, tm1, x)); + } + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : __expf(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : __expf(__fsub_rn(m1, n1))) : 1.0f; + float w0[4], w1[4]; +#pragma unroll + for (int n = 0; n < 4; ++n) { + s[n][0] = in[n][0] && act0 ? __expf(__fsub_rn(s[n][0], n0)) : 0.0f; + s[n][1] = in[n][1] && act0 ? __expf(__fsub_rn(s[n][1], n0)) : 0.0f; + s[n][2] = in[n][2] && act1 ? __expf(__fsub_rn(s[n][2], n1)) : 0.0f; + s[n][3] = in[n][3] && act1 ? __expf(__fsub_rn(s[n][3], n1)) : 0.0f; + const float x0 = __fadd_rn(s[n][0], s[n][1]), x1 = __fadd_rn(s[n][2], s[n][3]); + const float y0 = __fadd_rn(x0, __shfl_xor_sync(FULL, x0, 2)), y1 = __fadd_rn(x1, __shfl_xor_sync(FULL, x1, 2)); + w0[n] = __fadd_rn(y0, __shfl_xor_sync(FULL, y0, 1)); + w1[n] = __fadd_rn(y1, __shfl_xor_sync(FULL, y1, 1)); + } + l0 = __fmaf_rn(l0, a0, __fadd_rn(__fadd_rn(w0[0], w0[2]), __fadd_rn(w0[1], w0[3]))); + l1 = __fmaf_rn(l1, a1, __fadd_rn(__fadd_rn(w1[0], w1[2]), __fadd_rn(w1[1], w1[3]))); + m0 = n0; + m1 = n1; + uint32_t ph[2][4], pr[2][4]; // P's bf16 fragments, and (P32) the rest P - bf16(P) +#pragma unroll + for (int j = 0; j < 2; ++j) { + ph[j][0] = pack(s[2 * j][0], s[2 * j][1]); + ph[j][1] = pack(s[2 * j][2], s[2 * j][3]); + ph[j][2] = pack(s[2 * j + 1][0], s[2 * j + 1][1]); + ph[j][3] = pack(s[2 * j + 1][2], s[2 * j + 1][3]); + if constexpr (P32) { + const float2 h0 = unpack(ph[j][0]), h1 = unpack(ph[j][1]), h2 = unpack(ph[j][2]), h3 = unpack(ph[j][3]); + pr[j][0] = pack(__fsub_rn(s[2 * j][0], h0.x), __fsub_rn(s[2 * j][1], h0.y)); + pr[j][1] = pack(__fsub_rn(s[2 * j][2], h1.x), __fsub_rn(s[2 * j][3], h1.y)); + pr[j][2] = pack(__fsub_rn(s[2 * j + 1][0], h2.x), __fsub_rn(s[2 * j + 1][1], h2.y)); + pr[j][3] = pack(__fsub_rn(s[2 * j + 1][2], h3.x), __fsub_rn(s[2 * j + 1][3], h3.y)); + } + } + { + const uint4* sl = take(2 * t + 1); +#pragma unroll + for (int hf = 0; hf < NB / 8; ++hf) { // 64 columns at a time: the tile's P x V, then folded into o + float tv[8][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) { +#pragma unroll + for (int i = 0; i < 8; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, sl + sw(16 * j + 8 * (mt % 2) + lane % 8, cs * (SL / 8) + 8 * hf + i + mt / 2)); + if (j == 0) { + mma0(tv[i], ph[j], b[0], b[1]); + mma0(tv[i + 1], ph[j], b[2], b[3]); + } else { + mma(tv[i], ph[j], b[0], b[1]); + mma(tv[i + 1], ph[j], b[2], b[3]); + } + if constexpr (P32) { + mma(tv[i], pr[j], b[0], b[1]); + mma(tv[i + 1], pr[j], b[2], b[3]); + } + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + float* x = o[8 * hf + i]; + x[0] = __fmaf_rn(x[0], a0, tv[i][0]); + x[1] = __fmaf_rn(x[1], a0, tv[i][1]); + x[2] = __fmaf_rn(x[2], a1, tv[i][2]); + x[3] = __fmaf_rn(x[3], a1, tv[i][3]); + } + } + done(2 * t + 1); + } + } + wait_group<0>(); + bool odd = false; +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = cs * SL + 8 * i + 2 * qd; + const float x0 = __fdiv_rn(o[i][0], l0), x1 = __fdiv_rn(o[i][1], l0); + const float x2 = __fdiv_rn(o[i][2], l1), x3 = __fdiv_rn(o[i][3], l1); + if (r0 < W) { + *reinterpret_cast(out + ((int64_t)r0 * H + head) * D + d) = pack(x0, x1); + odd |= !isfinite(x0) || !isfinite(x1); + } + if (r1 < W) { + *reinterpret_cast(out + ((int64_t)r1 * H + head) * D + d) = pack(x2, x3); + odd |= !isfinite(x2) || !isfinite(x3); + } + } + return __syncthreads_or(odd); +} + +// A block with a NaN or infinite output: each of its rows' heads again in float64 (off the common path). +template +__device__ __noinline__ void odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, + __nv_bfloat16* out, const int* hi, int p0, int W, int H, int HK, int G, float scale, + int win, int ring) { + using K = Cfg; + const int HB = G % K::NG == 0 ? K::NG : 1, RB = K::NG / HB; + const int hk = blockIdx.y / (G / HB), hg = blockIdx.y % (G / HB); + const int row0 = blockIdx.x * 16 * RB; + const auto kv = [&](int key, int which) { + return (which == 0 ? kc : vc) + ((int64_t)(ring > 0 ? key % ring : key) * HK + hk) * D; + }; + for (int i = 0; i < 16 * RB * HB; ++i) { + const int r = row0 + i % (16 * RB), head = hk * G + hg * HB + i / (16 * RB); + if (r >= W) continue; + const int p = p0 + r; + slow_head(q + ((int64_t)r * H + head) * D, kv, win > 0 ? max(p - win + 1, 0) : 0, + hi != nullptr ? hi[r] : p, (double)scale, out + ((int64_t)r * H + head) * D, 0, D); + } +} + +template +__device__ __forceinline__ void kernel(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, + __nv_bfloat16* out, const int* hi, int p0, int W, int H, int HK, int G, + float scale, int win, int ring) { + if (rows(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring)) + odd(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); +} + +} // namespace prompt + +} // namespace tf_wrows diff --git a/zig/kernels/cuda/window_rows512.cuh b/zig/kernels/cuda/window_rows512.cuh new file mode 100644 index 0000000000..e974967e3c --- /dev/null +++ b/zig/kernels/cuda/window_rows512.cuh @@ -0,0 +1,181 @@ +// window_rows.cuh's decode partials at head dim 512: K and V tiles in shared memory, a block 4 query heads. +#pragma once + +#include "window_rows.cuh" + +namespace tf_wrows { + +// part's arithmetic for head dim 512: a block 4 heads, K and V tiles in shared memory, a warp a key run. +constexpr int TK = 32, NSG = 2; + +template +__device__ __forceinline__ void part_staged(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, const int* __restrict__ pos, + const int* __restrict__ slot, float* __restrict__ po, float* __restrict__ pm, + float* __restrict__ pl, int h, int hk, int win, long long ring, float scale) { + constexpr int DPL = D / 32, HW = 4, C = D / 8; // a lane's dims and columns, heads a block (and warp), uint4 a key + extern __shared__ uint4 tiles[]; // [NSG][TK][C] K or V tiles; then the runs' partial o + __shared__ float sp[CK * SP]; + __shared__ float mx[GMAX], ls[GMAX]; + float* ox = reinterpret_cast(tiles); + const int64_t r = blockIdx.x; + const int G = h / hk, HG = (G + HW - 1) / HW, c = blockIdx.z, chunks = gridDim.z; + const int kh = blockIdx.y / HG, hg = blockIdx.y % HG; // the block's heads: hg HW.. of key head kh + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int p = pos[r]; + const int64_t base = (slot != nullptr ? static_cast(slot[r]) : 0) * ring; + const int lo = win > 0 ? max(p - win + (FAULT == EDGE ? 0 : 1), 0) : 0; + const int k0 = max((lo / CK + c) * CK, lo), k1 = min((lo / CK + c) * CK + CK - 1, p); + const int64_t part = (r * hk + kh) * chunks + c; + const int own = min(HW, G - hg * HW); // heads of this block + if (k0 > k1) { + if (threadIdx.x < own) pm[part * G + hg * HW + threadIdx.x] = -INFINITY, pl[part * G + hg * HW + threadIdx.x] = 0.0f; + return; + } + const int nk = k1 - k0 + 1, nt = (nk + TK - 1) / TK; + const int64_t r0 = k0 % ring; + // warp kg: a tile's keys kg, kg + NW, ...; its keys in tile order are its run + constexpr int WG = NW; + const int kg = warp; + constexpr bool busy = true; + const auto stage = [&](const __nv_bfloat16* cache, int t) { + uint4* dst = tiles + (t % NSG) * TK * C; + for (int i = threadIdx.x; i < TK * C; i += NT) { + const int jj = t * TK + i / C; + int64_t row = r0 + min(jj, nk - 1); + if (row >= ring) row -= ring; + cp16z(dst + i, cache + ((base + row) * hk + kh) * D + (i % C) * 8, jj < nk); + } + commit(); + }; + const auto load = [&](const uint4* src, float (&x)[DPL]) { +#pragma unroll + for (int i = 0; i < DPL / 8; ++i) { + const uint4 u = src[i]; + const uint32_t w[4] = {u.x, u.y, u.z, u.w}; +#pragma unroll + for (int t = 0; t < 4; ++t) x[8 * i + 2 * t] = lo16(w[t]), x[8 * i + 2 * t + 1] = hi16(w[t]); + } + }; + float qv[HW][DPL]; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = min(hg * HW + hh, G - 1); + load(reinterpret_cast(q + (r * h + kh * G + hd) * D + lane * DPL), qv[hh]); + } + for (int t = 0; t < NSG - 1; ++t) + if (t < nt) stage(kc, t); + else commit(); + for (int t = 0; t < nt; ++t) { // scores: a lane's DPL dims in an fma chain, then a xor tree over the lanes + if (t + NSG - 1 < nt) stage(kc, t + NSG - 1); + else commit(); + asm volatile("cp.async.wait_group %0;\n" ::"n"(NSG - 1)); + __syncthreads(); + if (busy) { + for (int u0 = kg; u0 < TK; u0 += 4 * WG) { // four keys' chains side by side + float kf[4][DPL], sc[4][HW]; +#pragma unroll + for (int x = 0; x < 4; ++x) { + const int u = min(u0 + x * WG, TK - 1); + load(tiles + (t % NSG) * TK * C + u * C + lane * (DPL / 8), kf[x]); + } +#pragma unroll + for (int x = 0; x < 4; ++x) +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + sc[x][hh] = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; ++i) sc[x][hh] = __fmaf_rn(qv[hh][i], kf[x][i], sc[x][hh]); + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) +#pragma unroll + for (int x = 0; x < 4; ++x) +#pragma unroll + for (int hh = 0; hh < HW; ++hh) + sc[x][hh] = __fadd_rn(sc[x][hh], __shfl_xor_sync(FULL, sc[x][hh], o)); + if (lane == 0) +#pragma unroll + for (int x = 0; x < 4; ++x) { + const int u = u0 + x * WG; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) + if (u < TK && t * TK + u < nk && hg * HW + hh < G) + sp[(t * TK + u) * SP + hg * HW + hh] = __fmul_rn(sc[x][hh], scale); + } + } + } + __syncthreads(); + } + for (int t = 0; t < NSG - 1; ++t) // the values' first tiles load while the weights are formed + if (t < nt) stage(vc, t); + else commit(); + if (warp < own) { // a warp a head + const int hd = hg * HW + warp; + float m = -INFINITY; + for (int jj = lane; jj < nk; jj += 32) m = fmaxf(m, sp[jj * SP + hd]); + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(FULL, m, o)); + float l = 0.0f; + for (int jj = lane; jj < nk; jj += 32) { + const float sv = sp[jj * SP + hd]; + const float e = sv == -INFINITY ? 0.0f : expf(__fsub_rn(sv, m)); + sp[jj * SP + hd] = e; + l = __fadd_rn(l, e); + } + for (int o = 16; o > 0; o >>= 1) l = __fadd_rn(l, __shfl_xor_sync(FULL, l, o)); + if (lane == 0) mx[hd] = m, ls[hd] = l; + } + float acc[HW][DPL]; +#pragma unroll + for (int hh = 0; hh < HW; ++hh) +#pragma unroll + for (int i = 0; i < DPL; ++i) acc[hh][i] = 0.0f; + for (int t = 0; t < nt; ++t) { // o: a lane's DPL columns, an fma chain over its run's keys in order + if (t + NSG - 1 < nt) stage(vc, t + NSG - 1); + else commit(); + asm volatile("cp.async.wait_group %0;\n" ::"n"(NSG - 1)); + __syncthreads(); + if (busy) { + for (int u = kg; u < TK && t * TK + u < nk; u += WG) { + float vf[DPL]; + load(tiles + (t % NSG) * TK * C + u * C + lane * (DPL / 8), vf); +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + if (hg * HW + hh >= G) break; + const float w = sp[(t * TK + u) * SP + hg * HW + hh]; +#pragma unroll + for (int i = 0; i < DPL; ++i) acc[hh][i] = __fmaf_rn(w, vf[i], acc[hh][i]); + } + } + } + __syncthreads(); + } + if (busy && kg > 0) +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = hg * HW + hh; + if (hd >= G) break; +#pragma unroll + for (int i = 0; i < DPL; ++i) ox[((kg - 1) * HW + hh) * D + lane * DPL + i] = acc[hh][i]; + } + __syncthreads(); + if (busy && kg == 0) { +#pragma unroll + for (int hh = 0; hh < HW; ++hh) { + const int hd = hg * HW + hh; + if (hd >= G) break; +#pragma unroll + for (int i = 0; i < DPL; ++i) { + float a = acc[hh][i]; + for (int k = 1; k < WG; ++k) a = __fadd_rn(a, ox[((k - 1) * HW + hh) * D + lane * DPL + i]); + po[(part * G + hd) * D + lane * DPL + i] = a; + } + } + } + if (threadIdx.x < own) { + const int hd = hg * HW + threadIdx.x; + pm[part * G + hd] = mx[hd], pl[part * G + hd] = ls[hd]; + } +} + +} // namespace tf_wrows diff --git a/zig/src/core/act_ref.zig b/zig/src/core/act_ref.zig index f56f282175..77d696b461 100644 --- a/zig/src/core/act_ref.zig +++ b/zig/src/core/act_ref.zig @@ -1,5 +1,6 @@ //! Activation rows in float64, the reference every backend's activation kernels are checked against. const std = @import("std"); +const bf = @import("bf16_round.zig"); pub fn sigmoid(x: f64) f64 { return 1 / (1 + @exp(-x)); @@ -20,6 +21,54 @@ pub fn affine4(q: u4, scale: f64, bias: f64) f64 { return @as(f64, @floatFromInt(q)) * scale + bias; } +/// mlx_lm's gelu_approx constants, sqrt(2 / pi) and 0.044715, as the fp32 values the kernels round to bf16. +pub const gelu_c0: f32 = 0.7978845608028654; +pub const gelu_c1: f32 = 0.044715; + +/// MLX's geglu at its bf16 rounding points (Metal gemma4/glue.metal's gemma_geglu), tanh's value times (1 + t). +pub fn geglu(x: f64, up: f64, t: f64) f64 { + const cube = bf.bf16(x * x * x); // exact before its rounding for a bf16 x + const inner = bf.bf16(x + bf.bf16(bf.bf16(gelu_c1) * cube)); + const th = bf.bf16(std.math.tanh(bf.bf16(bf.bf16(gelu_c0) * inner)) * (1 + t)); + return bf.bf16(bf.bf16(bf.bf16(0.5 * x) * bf.bf16(1 + th)) * up); +} + +/// tf_draft_swiglu's bf16 steps (Metal core/draft_ops.metal), exp's value times (1 + te), the reciprocal's (1 + td). +pub fn swigluRound(g: f64, up: f64, te: f64, td: f64) f64 { + const e = bf.bf16(@exp(@abs(g)) * (1 + te)); + const inv = bf.bf16(1 / bf.bf16(1 + e) * (1 + td)); + const sig = if (g < 0) inv else bf.bf16(1 - inv); + return bf.bf16(bf.bf16(g * sig) * up); +} + +/// MLX's soft cap at its points: bf16(fp32(bf16(tanh(bf16(fp32(x / cap))) (1 + t)) cap)), fp32 first as the kernel. +pub fn softcap(x: f64, cap: f64, t: f64) f64 { + const q = bf.bf16(@as(f32, @floatCast(x / cap))); + const th = bf.bf16(std.math.tanh(q) * (1 + t)); + return bf.bf16(@as(f32, @floatCast(th * cap))); +} + +/// softcap's float64 ends with tanhf within `spread` either way (monotonic in t). +pub fn softcapRange(x: f64, cap: f64, spread: f64) [2]f64 { + const a = softcap(x, cap, -spread); + const b = softcap(x, cap, spread); + if (std.math.isNan(a) or std.math.isNan(b)) return @splat(std.math.nan(f64)); + return .{ @min(a, b), @max(a, b) }; +} + +/// One of the two at its float64 ends: every fp32 libm value (tanh, exp) and division within `spread` either way. +pub fn mlxRange(geglu_: bool, x: f64, up: f64, spread: f64) [2]f64 { + var lo = std.math.inf(f64); + var hi = -std.math.inf(f64); + for ([_]f64{ -spread, spread }) |a| for ([_]f64{ -spread, spread }) |b| { + const v = if (geglu_) geglu(x, up, a) else swigluRound(x, up, a, b); + if (std.math.isNan(v)) return @splat(std.math.nan(f64)); + lo = @min(lo, v); + hi = @max(hi, v); + }; + return .{ lo, hi }; +} + /// A bf16 output of fp32 arithmetic through expf: half a bf16 ulp plus 2^-18 of |want|. pub fn tolerance(want: f64) f64 { return @abs(want) * (std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -18)) + 1e-30; @@ -31,3 +80,34 @@ test "activations against hand values" { try std.testing.expectApproxEqRel(@as(f64, 2 * sigmoid(2) * 3), swiglu(2, 3), 1e-15); try std.testing.expect(tolerance(1) > tolerance(0.5)); } + +test "MLX's steps: rounding points against hand values, the ranges ordered, and both near the once-rounded forms" { + try std.testing.expectEqual(@as(f64, 0), geglu(0, 3, 0)); + try std.testing.expectEqual(@as(f64, 0), swigluRound(0, 1, 0, 0)); + try std.testing.expectEqual(bf.bf16(-2 * bf.bf16(1 / bf.bf16(1 + bf.bf16(@exp(2.0))))), swigluRound(-2, 1, 0, 0)); // g < 0: 1 / (1 + e) + try std.testing.expectEqual(bf.bf16(2 * bf.bf16(1 - bf.bf16(1 / bf.bf16(1 + bf.bf16(@exp(2.0)))))), swigluRound(2, 1, 0, 0)); + var prng = std.Random.DefaultPrng.init(9); + for (0..2000) |_| { + const x = bf.bf16(prng.random().float(f64) * 8 - 4); + const up = bf.bf16(prng.random().float(f64) * 4 - 2); + for ([_]bool{ true, false }) |g| { + const r = mlxRange(g, x, up, 0x1p-21); + try std.testing.expect(r[0] <= r[1]); + const once = if (g) 0.5 * x * (1 + std.math.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) * up else swiglu(x, up); + try std.testing.expect(@abs(r[0] - once) <= 0x1p-5 * @abs(x * up) + 0x1p-10); // a step of 1 + tanh: 2^-9 of x up + } + } +} + +test "the soft cap: rounding points against hand values, ends ordered, within a step of cap tanh(x / cap)" { + try std.testing.expectEqual(@as(f64, 0), softcap(0, 30, 0)); + try std.testing.expectEqual(@as(f64, 30), softcap(1e6, 30, 0)); // tanh rounds to 1 + try std.testing.expectEqual(-softcap(7.25, 30, 0), softcap(-7.25, 30, 0)); + var prng = std.Random.DefaultPrng.init(3); + for (0..4000) |_| { + const x = bf.bf16(prng.random().float(f64) * 120 - 60); + const r = softcapRange(x, 30, 0x1p-21); + try std.testing.expect(r[0] <= r[1]); + try std.testing.expect(@abs(r[0] - 30 * std.math.tanh(x / 30)) <= 0x1p-6 * 30); + } +} diff --git a/zig/src/core/attention_range.zig b/zig/src/core/attention_range.zig new file mode 100644 index 0000000000..196d9046d6 --- /dev/null +++ b/zig/src/core/attention_range.zig @@ -0,0 +1,170 @@ +//! A head's attention in float64 at Metal's rounding points (the output's), and the range window_rows.cu's fp32 allows. +const std = @import("std"); +const bfr = @import("bf16_round.zig"); + +const u: f64 = 0x1p-24; + +fn gamma(n: f64) f64 { + return n * u / (1 - n * u); +} + +/// A form's fp32 arithmetic, as the range needs it; each count is in roundings of 2^-24. +pub const Form = struct { + /// A score's spread over sum_e |q_e k_e|: lanes' fma chains and a xor tree, or mma steps and slice sums. + score: f64, + /// __expf (2 + 1.173 |x| ulps), else expf (2 ulps). + fast_exp: bool, + /// The merge's expf(m_c - M) of each chunk's weights (decode), on top of the weights' own. + chunk_exp: bool, + /// P's error into P x V, relative: bf16 2^-8, two bf16 halves 2^-16, fp32 0. + p_rel: f64, + /// The chains of o's terms p_j v_j and of l's terms p_j, as roundings of their terms' magnitude. + num: f64, + den: f64, +}; + +/// An mma k16 step's error in 2^-24 of |acc| + sum |a b|: 17 terms cut at the largest's ulp, then rounded down. +pub const mma_step = 36; + +/// window_rows.cu's decode at head dim d over n keys: lanes' chains of d / 32 then a 5-level tree, 256-key chunks. +pub fn decode(d: usize, n: usize) Form { + const chunks: f64 = @floatFromInt(n / 256 + 2); + return .{ .score = @floatFromInt(d / 32 + 5), .fast_exp = false, .chunk_exp = true, .p_rel = 0, .num = 264 + chunks + 1, .den = 13 + chunks + 1 }; +} + +/// Prompt rows over n keys: 8 mma steps a 128-dim slice, slices in order, an fma a 32-key tile; P32: 4 P x V steps. +pub fn prompt(d: usize, n: usize, p32: bool) Form { + const tiles: f64 = @floatFromInt(n / 32 + 2); + const steps: f64 = if (p32) 4 else 2; + return .{ .score = 8 * mma_step + @as(f64, @floatFromInt(d / 128)), .fast_exp = true, .chunk_exp = false, .p_rel = if (p32) 0x1p-16 else 0x1p-8, .num = mma_step * steps + tiles + 2, .den = tiles + 6 }; +} + +/// A head over n keys (k, v [n][d]): `point` float64's output, `half` how far fp32 may land (w, tau: n scratch). +pub fn head(q: []const f64, k: []const f64, v: []const f64, n: usize, d: usize, scale: f64, f: Form, w: []f64, tau: []f64, point: []f64, half: []f64) void { + var m: f64 = -std.math.inf(f64); + var dmax: f64 = 0; + for (0..n) |j| { + var s: f64 = 0; + var a: f64 = 0; + for (0..d) |e| { + s += q[e] * k[j * d + e]; + a += @abs(q[e] * k[j * d + e]); + } + w[j] = s * scale; + tau[j] = f.score * u * a * @abs(scale) + u * @abs(w[j]); // the score's spread; below, the weight's + m = @max(m, w[j]); + if (std.math.isFinite(tau[j])) dmax = @max(dmax, tau[j]); + } + var l: f64 = 0; + var tmax: f64 = 0; + var vmax: f64 = 0; + for (w[0..n], tau[0..n]) |*p, *t| { + const x = (m - p.*) + 2 * dmax; // |the computed s_j - max| at most + var e = t.* + u * x + if (f.fast_exp) (3 + 1.2 * x) * 0x1p-23 else 0x1p-22; + if (f.chunk_exp) e += 0x1p-22 + u * x; + p.* = @exp(p.* - m); + l += p.*; + t.* = std.math.expm1(e); + if (std.math.isFinite(t.*)) tmax = @max(tmax, t.*); + } + for (v[0 .. n * d]) |x| if (std.math.isFinite(x)) { + vmax = @max(vmax, @abs(x)); + }; + const big = tmax >= 0.5; // scores too far apart for the spread's first order: any finite value + const kappa = (1 + tmax) / (1 - tmax); + const tiny = @as(f64, @floatFromInt(n + 8)) * 0x1p-120 * (1 + vmax); + for (0..d) |e| { + var o: f64 = 0; + for (0..n) |j| o += w[j] * v[j * d + e]; + o /= l; + point[e] = o; + if (!std.math.isFinite(o) or big) { + half[e] = if (big) std.math.inf(f64) else 0; + continue; + } + var r: f64 = 0; + var dv: f64 = 0; + for (0..n) |j| { + const x = v[j * d + e]; + r += w[j] * @abs(x); + dv += w[j] * tau[j] * @abs(x - o); + } + r /= l; + dv /= l; + const e1 = dv / (1 - tmax); + const en = (gamma(f.num) + f.p_rel) * r * kappa; + const base = @abs(o) + e1 + en; + half[e] = (e1 + en + (gamma(f.den) + u) * base) * (1 + 0x1p-20) + tiny; + } +} + +/// What a check saw: outputs whose range holds one bf16 value, several, and outputs non-finite in float64. +pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0 }; + +/// bf16 bits in value order, -0 just below +0 (NaN bits fall outside every finite range). +fn key(b: u16) i32 { + return if (b & 0x8000 != 0) -@as(i32, b & 0x7fff) - 1 else b; +} + +fn bitsOf(v: f64) u16 { + return bfr.bits(@floatCast(bfr.bf16(v))); +} + +/// bf16 bits `got` in point +- half's ends rounded to bf16, in value order (-0 < +0); NaN for NaN; infinities by bits. +pub fn within(point: f64, half: f64, got: u16, stats: *Stats) bool { + const nan = got & 0x7f80 == 0x7f80 and got & 0x7f != 0; + if (std.math.isNan(point)) { + stats.nonfinite += 1; + return nan; + } + if (std.math.isInf(point)) { + stats.nonfinite += 1; + return got == bitsOf(point); + } + const lo = bitsOf(point - half); + const hi = bitsOf(point + half); + if (lo == hi) stats.decided += 1 else stats.ranged += 1; + return !nan and key(lo) <= key(got) and key(got) <= key(hi); +} + +test "one key gives its value; a range's ends, -0 below +0, NaN only for NaN, infinities by their bits" { + var w: [2]f64 = undefined; + var t: [2]f64 = undefined; + var pt: [2]f64 = undefined; + var hf: [2]f64 = undefined; + head(&.{ 1, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, 4, 6, 8 }, 1, 2, 1, decode(256, 1), &w, &t, &pt, &hf); + try std.testing.expectEqualSlices(f64, &.{ 2, 4 }, &pt); + var st: Stats = .{}; + try std.testing.expect(within(2, hf[0], 0x4000, &st)); // 2.0 + try std.testing.expect(!within(2, hf[0], 0x4001, &st)); + try std.testing.expect(!within(2, hf[0], 0x0000, &st)); + try std.testing.expectEqual(@as(usize, 3), st.decided); + // a range across zero holds both zeros; one below zero holds -0 and not +0 + try std.testing.expect(within(0, 0x1p-140, 0x8000, &st)); + try std.testing.expect(within(0, 0x1p-140, 0x0000, &st)); + try std.testing.expect(within(-0x1p-140, 0x1p-142, 0x8000, &st)); + try std.testing.expect(!within(-0x1p-140, 0x1p-142, 0x0000, &st)); + // NaN only for NaN (either sign), infinities by their bits + try std.testing.expect(within(std.math.nan(f64), 0, 0xffc1, &st)); + try std.testing.expect(!within(std.math.nan(f64), 0, 0x7f80, &st)); + try std.testing.expect(!within(1e300, 1e300, 0x7fc0, &st)); + try std.testing.expect(within(-std.math.inf(f64), 0, 0xff80, &st)); + try std.testing.expect(!within(-std.math.inf(f64), 0, 0x7f80, &st)); + // a range whose upper end rounds to +Inf holds +Inf, and the largest finite, but not zero + try std.testing.expect(within(3.3e38, 1e37, 0x7f80, &st)); + try std.testing.expect(within(3.3e38, 1e37, 0x7f7f, &st)); + try std.testing.expect(!within(3.3e38, 1e37, 0x0000, &st)); +} + +test "a weight's spread moves an output by about its share of |v - o|" { + var w: [2]f64 = undefined; + var t: [2]f64 = undefined; + var pt: [1]f64 = undefined; + var hf: [1]f64 = undefined; + // two keys of equal score: o is the mean; the range grows with the scores' magnitude + head(&.{1}, &.{ 3, 3 }, &.{ 0, 1 }, 2, 1, 1, prompt(256, 2, true), &w, &t, &pt, &hf); + try std.testing.expectEqual(@as(f64, 0.5), pt[0]); + const small = hf[0]; + head(&.{1}, &.{ 30, 30 }, &.{ 0, 1 }, 2, 1, 1, prompt(256, 2, true), &w, &t, &pt, &hf); + try std.testing.expect(hf[0] > small and hf[0] < 0x1p-9); +} diff --git a/zig/src/core/bf16_round.zig b/zig/src/core/bf16_round.zig new file mode 100644 index 0000000000..15fb39ebbf --- /dev/null +++ b/zig/src/core/bf16_round.zig @@ -0,0 +1,79 @@ +//! bf16 rounding of float64 values, for float64 references of kernels that round to bf16 between their steps. +const std = @import("std"); + +/// v to the nearest bf16, ties to even, in one rounding (not through fp32); NaN and infinities pass, overflow is Inf. +pub fn bf16(v: f64) f64 { + if (!std.math.isFinite(v) or v == 0) return v; + const e = @max(std.math.frexp(v).exponent, -125); // |v| in [2^(e-1), 2^e); below 2^-126 the subnormal step + const step = std.math.ldexp(@as(f64, 1), e - 8); + const q = v / step; // a power of two apart: exact + var t = @floor(q); + const d = q - t; + if (d > 0.5 or (d == 0.5 and @mod(t, 2) != 0)) t += 1; + const out = std.math.copysign(t * step, v); // an underflow to zero keeps v's sign, as bits() does + return if (@abs(out) >= 0x1p128) std.math.copysign(std.math.inf(f64), v) else out; +} + +/// fp32 to bf16 bits, ties to even, as __float2bfloat16_rn rounds (NaN stays NaN). +pub fn bits(v: f32) u16 { + const u: u32 = @bitCast(v); + if (std.math.isNan(v)) return @truncate((u >> 16) | 0x40); + return @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +/// The fp32 value of bf16 bits. +pub fn value(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 rounded to bf16 and back: a kernel's `__bfloat162float(__float2bfloat16_rn(v))`. +pub fn round32(v: f32) f32 { + return value(bits(v)); +} + +test "ties to even, the subnormal step, overflow, and fp32 values rounded once as the fp32 path rounds them" { + try std.testing.expectEqual(@as(f64, 1), bf16(1 + 0x1p-8)); // a tie, to the even 1 + try std.testing.expectEqual(@as(f64, 1 + 0x1p-6), bf16(1 + 3 * 0x1p-8)); // a tie, to the even 1 + 2^-6 + try std.testing.expectEqual(@as(f64, 1 + 0x1p-7), bf16(1 + 0x1p-8 + 0x1p-40)); // just past the tie + try std.testing.expectEqual(@as(f64, 0x1p-133), bf16(0x1p-133 * 1.4)); + try std.testing.expectEqual(@as(f64, 0), bf16(0x1p-135)); + try std.testing.expect(std.math.isInf(bf16(0x1p128 * 0.9999))); + // the largest, under half a step past it + try std.testing.expectEqual(@as(f64, -255 * 0x1p120), bf16(-255 * 0x1p120 - 0x1p111)); + try std.testing.expect(std.math.isNan(bf16(std.math.nan(f64)))); + var prng = std.Random.DefaultPrng.init(5); + for (0..100000) |_| { + const u = prng.random().int(u32); + const v: f32 = @bitCast(u); + if (!std.math.isFinite(v)) continue; + try std.testing.expectEqual(bits(v), bitsOf(bf16(v))); // bits: a zero's sign counts + } +} + +/// bf16(v)'s bits, for tests that must tell -0 from +0 (float equality does not). +fn bitsOf(v: f64) u16 { + return bits(@floatCast(v)); +} + +test "raw bits at the edges: signed zeros from underflow, subnormal ties, the normal boundary, overflow" { + const cases = [_]struct { f64, u16 }{ + .{ -0x1p-135, 0x8000 }, // underflow to zero keeps the sign + .{ 0x1p-135, 0x0000 }, + .{ -0.0, 0x8000 }, + .{ -0x1p-134, 0x8000 }, // half the least subnormal: a tie, to the even zero + .{ -0x1p-134 - 0x1p-160, 0x8001 }, // just past the tie + .{ 0x1p-133, 0x0001 }, // the least subnormal + .{ -3 * 0x1p-134, 0x8002 }, // a subnormal tie, to the even 2 + .{ 0x1p-126 - 0x1p-134, 0x0080 }, // the largest subnormal's tie, up to the least normal + .{ -(1 + 0x1p-8), 0xbf80 }, // a tie, to the even -1 + .{ -(1 + 3 * 0x1p-8), 0xbf82 }, + .{ 255 * 0x1p120, 0x7f7f }, // the largest finite + .{ -(255 * 0x1p120 + 0x1p119), 0xff80 }, // its tie with 2^128: to the even -Inf + .{ -(255 * 0x1p120 + 0x1p119 - 0x1p90), 0xff7f }, + .{ 0x1p300, 0x7f80 }, + .{ -std.math.inf(f64), 0xff80 }, + }; + for (cases) |cs| try std.testing.expectEqual(cs[1], bitsOf(bf16(cs[0]))); + try std.testing.expect(std.math.isNan(value(bitsOf(bf16(-std.math.nan(f64)))))); + try std.testing.expectEqual(@as(u16, 0x8000), bits(-0x1p-135)); // fp32's path agrees +} diff --git a/zig/src/core/contract/gemma_cuda_records.zig b/zig/src/core/contract/gemma_cuda_records.zig new file mode 100644 index 0000000000..9fe70593d0 --- /dev/null +++ b/zig/src/core/contract/gemma_cuda_records.zig @@ -0,0 +1,14 @@ +//! Gemma 4's CUDA launch records (act_rows' MLX activations and soft cap, embed_rows' and qk_rows' MLX rows). +const CudaRecord = @import("registry.zig").CudaRecord; + +pub const cuda = [_]CudaRecord{ + .{ .launch = .{ .name = "act_rows tf_geglu_rows 256 threads", .symbol = "tf_geglu_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_geglu_rows 128 threads", .symbol = "tf_geglu_rows", .image = "act_rows", .threads = 128 }, .usage = .{ .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows_round 256 threads", .symbol = "tf_swiglu_rows_round", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows_round 128 threads", .symbol = "tf_swiglu_rows_round", .image = "act_rows", .threads = 128 }, .usage = .{ .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_softcap_rows 256 threads", .symbol = "tf_softcap_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_softcap_rows 128 threads", .symbol = "tf_softcap_rows", .image = "act_rows", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "embed_rows tf_embed_rows_affine", .symbol = "tf_embed_rows_affine", .image = "embed_rows", .threads = 256 }, .usage = .{ .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows_once 512 threads", .symbol = "tf_qk_rows_once", .image = "qk_rows", .threads = 512 }, .usage = .{ .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 36, .static = 64, .bound = 512 }, .{ .regs = 36, .static = 64, .bound = 512 }, .{ .regs = 34, .static = 64, .bound = 512 }, .{ .regs = 34, .static = 64, .bound = 512 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows_round 512 threads", .symbol = "tf_qk_rows_round", .image = "qk_rows", .threads = 512 }, .usage = .{ .{ .regs = 30, .static = 64, .bound = 512 }, .{ .regs = 36, .static = 64, .bound = 512 }, .{ .regs = 36, .static = 64, .bound = 512 }, .{ .regs = 34, .static = 64, .bound = 512 }, .{ .regs = 34, .static = 64, .bound = 512 } } }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..c93a6d8da4 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -54,6 +54,8 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = .{ .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "norms tf_add_rms_rows_bf16h_tail", .symbol = "tf_add_rms_rows_bf16h_tail", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 94, .static = 64, .bound = 256 }, .{ .regs = 95, .static = 64, .bound = 256 }, .{ .regs = 95, .static = 64, .bound = 256 }, .{ .regs = 147, .static = 64, .bound = 256 }, .{ .regs = 147, .static = 64, .bound = 256 } } }, + .{ .launch = .{ .name = "norms tf_add_rms_rows_bf16h_tail_round", .symbol = "tf_add_rms_rows_bf16h_tail_round", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 94, .static = 64, .bound = 256 }, .{ .regs = 95, .static = 64, .bound = 256 }, .{ .regs = 95, .static = 64, .bound = 256 }, .{ .regs = 147, .static = 64, .bound = 256 }, .{ .regs = 147, .static = 64, .bound = 256 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "router tf_nemo_topk_wide", .symbol = "tf_nemo_topk_wide", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, @@ -221,7 +223,7 @@ pub const 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_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("gemma_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 }; /// One launch's usage under a later nvcc, by the record's name. diff --git a/zig/src/core/contract/window_rows_records.zig b/zig/src/core/contract/window_rows_records.zig new file mode 100644 index 0000000000..c4c9e0db06 --- /dev/null +++ b/zig/src/core/contract/window_rows_records.zig @@ -0,0 +1,14 @@ +//! window_rows.cu's launches on CUDA (head dims 256 and 512), read by kernel-contract-limits from every SM. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "window_rows tf_wrows_part_256", .symbol = "tf_wrows_part_256", .image = "window_rows", .threads = 256 }, .usage = .{ .{ .regs = 96, .static = 46208, .bound = 256 }, .{ .regs = 98, .static = 46208, .bound = 256 }, .{ .regs = 98, .static = 46208, .bound = 256 }, .{ .regs = 104, .static = 46208, .bound = 256 }, .{ .regs = 104, .static = 46208, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_merge_256", .symbol = "tf_wrows_merge_256", .image = "window_rows", .threads = 256 }, .usage = .{ .{ .regs = 48, .static = 128, .bound = 256 }, .{ .regs = 48, .static = 128, .bound = 256 }, .{ .regs = 48, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_prompt_256", .symbol = "tf_wrows_prompt_256", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 240, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_prompt_256_p32", .symbol = "tf_wrows_prompt_256_p32", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 252, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_part_512", .symbol = "tf_wrows_part_512", .image = "window_rows", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 160, .static = 17536, .bound = 256 }, .{ .regs = 163, .static = 17536, .bound = 256 }, .{ .regs = 163, .static = 17536, .bound = 256 }, .{ .regs = 125, .static = 17536, .bound = 256 }, .{ .regs = 125, .static = 17536, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_merge_512", .symbol = "tf_wrows_merge_512", .image = "window_rows", .threads = 256 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 62, .static = 128, .bound = 256 }, .{ .regs = 62, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_prompt_512", .symbol = "tf_wrows_prompt_512", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 241, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 } } }, + .{ .launch = .{ .name = "window_rows tf_wrows_prompt_512_p32", .symbol = "tf_wrows_prompt_512_p32", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 249, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 254, .static = 144, .bound = 256 }, .{ .regs = 254, .static = 144, .bound = 256 } } }, +}; diff --git a/zig/src/core/embed_ref.zig b/zig/src/core/embed_ref.zig new file mode 100644 index 0000000000..48a8d55157 --- /dev/null +++ b/zig/src/core/embed_ref.zig @@ -0,0 +1,47 @@ +//! Token rows in float64 at their rounding points: MLX's affine dequantize and a bf16 scale (cuda/embed_rows.zig). +const std = @import("std"); +const bf = @import("bf16_round.zig"); + +/// Code i of an MLX row packed `bits` wide in little-endian u32 words, low code first (a code may cross a word). +pub fn code(row: []const u32, i: usize, bits: u4) u32 { + const bit = i * bits; + const sh: u6 = @intCast(bit % 32); + var v: u64 = row[bit / 32]; + if (sh + bits > 32) v |= @as(u64, row[bit / 32 + 1]) << 32; + return @intCast((v >> sh) & ((@as(u64, 1) << bits) - 1)); +} + +/// MLX's dequantize, q s + b in fp32 then bf16, times the bf16 `scale` (null: none) in fp32 then bf16 again. +pub fn affine(q: u32, s: f64, b: f64, scale: ?f64) f64 { + const e: f32 = @mulAdd(f32, @floatFromInt(q), @floatCast(s), @floatCast(b)); // q s is exact in fp32 + const v = bf.bf16(e); + const x = scale orelse return v; + const p: f32 = @floatCast(v * x); // exact in float64: two bf16 values + return bf.bf16(p); +} + +test "MLX's packing: codes in order across word edges at every width" { + var prng = std.Random.DefaultPrng.init(11); + for ([_]u4{ 2, 3, 4, 5, 6, 8 }) |bits| { + var want: [128]u32 = undefined; + var row: [33]u32 = @splat(0); + for (&want, 0..) |*q, i| { + q.* = prng.random().int(u32) & ((@as(u32, 1) << bits) - 1); + const bit = i * bits; + row[bit / 32] |= q.* << @intCast(bit % 32); + if (bit % 32 + bits > 32) row[bit / 32 + 1] |= q.* >> @intCast(32 - bit % 32); + } + for (want, 0..) |q, i| try std.testing.expectEqual(q, code(&row, i, bits)); + } + try std.testing.expectEqual(@as(u32, 8), code(&.{0x8765_4321}, 7, 4)); // the top nibble is code 7 +} + +test "the rounding points: the bf16 value before the scale, and the scale's product rounded once" { + try std.testing.expectEqual(@as(f64, 3 * 0.5 - 1), affine(3, 0.5, -1, null)); + try std.testing.expectEqual(@as(f64, 31), affine(15, 2, 1, 1)); // exact throughout + // 1 + 2^-8 rounds to 1 (a tie to even) before the scale: times 3 gives 3, where the unrounded value gives 3 + 2^-6 + try std.testing.expectEqual(@as(f64, 3), affine(1, 1, 0x1p-8, 3)); + try std.testing.expectEqual(@as(f64, 3 + 0x1p-6), bf.bf16((1 + 0x1p-8) * 3)); + try std.testing.expect(std.math.isNan(affine(1, std.math.nan(f64), 0, 62))); + try std.testing.expect(std.math.isInf(affine(255, 0x1p120, 0, 0x1p10))); // past fp32's range: Inf, as the fp32 product +} diff --git a/zig/src/core/norms_ref.zig b/zig/src/core/norms_ref.zig index ef2a2d14e2..0ac09eaa51 100644 --- a/zig/src/core/norms_ref.zig +++ b/zig/src/core/norms_ref.zig @@ -1,5 +1,6 @@ //! RMS norm rows in float64, the reference every backend's norm kernels are checked against (CUDA: cuda/norms.zig). const std = @import("std"); +const bf = @import("bf16_round.zig"); /// y = x / sqrt(mean(x^2) + eps) * w for one row (w null: no weight). pub fn rms(x: []const f64, w: ?[]const f64, eps: f64, y: []f64) void { @@ -63,6 +64,67 @@ pub fn gatedRms(y: []const f64, z: []const f64, w: []const f64, eps: f64, out: [ for (out, z) |*o, zz| o.* *= zz / (1 + @exp(-zz)); } +/// A bf16 residual's tail (cuda/norms.zig's addRmsTail) at its rounding points; `round`: MLX's bf16(w * bf16(v * inv)). +pub const Tail = struct { + round: bool, + residual: ?[]const f64 = null, + branch: ?[]const f64 = null, + scalar: ?f64 = null, + eps: f64 = 1e-6, + /// The relative spread of an fp32 inverse RMS and the products by it: the outputs are ranges over it (0: points). + spread: f64 = 0, +}; + +/// Sums of squares in float64, the tail's exact sums; a test may pass its own `squares` (another accumulator). +pub const Exact = struct { + pub fn squares(_: Exact, v: []const f64) f64 { + var s: f64 = 0; + for (v) |x| s += x * x; + return s; + } +}; + +/// norm(v) w at one value: w bf16(v inv) rounded (`round`), else v inv w rounded once. +fn normed(v: f64, inv: f64, w: f64, round: bool) f64 { + return bf.bf16(if (round) w * bf.bf16(v * inv) else v * inv * w); +} + +/// a and b in order; NaN in either gives NaN for both. +fn order(a: f64, b: f64) [2]f64 { + if (std.math.isNan(a) or std.math.isNan(b)) return @splat(std.math.nan(f64)); + return .{ @min(a, b), @max(a, b) }; +} + +/// One row's h = bf16(bf16(r + norm(x) wb) sc) and y = norm(h) w, each a range over inverse RMS values within spread. +pub fn tail(t: Tail, x: []const f64, w: []const f64, sums: anytype, h_lo: []f64, h_hi: []f64, y_lo: []f64, y_hi: []f64) void { + const n: f64 = @floatFromInt(x.len); + const inv1 = if (t.branch != null) 1 / @sqrt(sums.squares(x) / n + t.eps) else 0; + for (x, h_lo, h_hi, 0..) |v, *lo, *hi, i| { + var ends: [2]f64 = undefined; + for (&ends, [_]f64{ 1 - t.spread, 1 + t.spread }) |*e, f| { + var s = if (t.branch) |wb| normed(v, inv1 * f, wb[i], t.round) else v; + if (t.residual) |r| s = r[i] + s; + e.* = if (t.scalar) |sc| bf.bf16(bf.bf16(s) * sc) else bf.bf16(s); + } + lo.*, hi.* = order(ends[0], ends[1]); + } + for (h_lo, h_hi, y_lo, y_hi) |lo, hi, *small, *big| { // the squares' least and largest terms, as scratch + const a, const b = order(@abs(lo), @abs(hi)); + small.* = if (lo <= 0 and hi >= 0) 0 else a; + big.* = b; + } + const inv_lo = (1 - t.spread) / @sqrt(sums.squares(y_hi) / n + t.eps); + const inv_hi = (1 + t.spread) / @sqrt(sums.squares(y_lo) / n + t.eps); + for (h_lo, h_hi, w, y_lo, y_hi) |lo, hi, ww, *ylo, *yhi| { + var most = [2]f64{ std.math.inf(f64), -std.math.inf(f64) }; + for ([_]f64{ lo * inv_lo, lo * inv_hi, hi * inv_lo, hi * inv_hi }) |p| { + const o = normed(p, 1, ww, t.round); // monotone in p: the corners bound it + if (std.math.isNan(o)) most = @splat(std.math.nan(f64)) else most = .{ @min(most[0], o), @max(most[1], o) }; + } + ylo.*, yhi.* = most; + } +} + /// The element kinds the kernels read and write, numbered as they are (hip/ops/common.hpp: 0 fp32, 1 fp16, 2 bf16). pub const Kind = enum(u8) { f32 = 0, f16 = 1, bf16 = 2 }; @@ -81,3 +143,32 @@ test "a row's norm and the residual step against hand values" { try std.testing.expectEqualSlices(f64, &.{ 3, -3, 3, -3 }, &res); try std.testing.expect(tolerance(1, .bf16) > tolerance(1, .f32)); } + +test "a tail's rounding points: none set passes x, the scalar rounds the sum first, MLX's rows differ from one rounding" { + const x = [_]f64{ 1, -1, 1, -1 }; + const w = [_]f64{ 1, 1, 1, 1 }; + var h: [2][4]f64 = undefined; + var y: [2][4]f64 = undefined; + tail(.{ .round = true, .eps = 0 }, &x, &w, Exact{}, &h[0], &h[1], &y[0], &y[1]); + try std.testing.expectEqualSlices(f64, &x, &h[0]); + try std.testing.expectEqualSlices(f64, &x, &y[0]); + const r = [_]f64{ 1, 1, 1, 1 }; + const tiny = [_]f64{ 0x1p-8, 0x1p-8, 0x1p-8, 0x1p-8 }; + tail(.{ .round = false, .residual = &r, .scalar = 1 + 0x1p-7, .eps = 0 }, &tiny, &w, Exact{}, &h[0], &h[1], &y[0], &y[1]); + try std.testing.expectEqual(@as(f64, 1 + 0x1p-7), h[0][0]); // bf16(1 + 2^-8) = 1 (a tie to even), then 1 + 2^-7 + var prng = std.Random.DefaultPrng.init(3); + var v: [64]f64 = undefined; + var wb: [64]f64 = undefined; + for (&v, &wb) |*a, *b| { + a.* = bf.bf16(prng.random().float(f64) * 4 - 2); + b.* = bf.bf16(0.5 + prng.random().float(f64)); + } + var rows: [4][64]f64 = undefined; + var ys: [2][64]f64 = undefined; + tail(.{ .round = true, .branch = &wb }, &v, &wb, Exact{}, &rows[0], &rows[1], &ys[0], &ys[1]); + tail(.{ .round = false, .branch = &wb }, &v, &wb, Exact{}, &rows[2], &rows[3], &ys[0], &ys[1]); + try std.testing.expectEqualSlices(f64, &rows[0], &rows[1]); // spread 0: points + try std.testing.expect(!std.mem.eql(f64, &rows[0], &rows[2])); + tail(.{ .round = true, .branch = &wb, .spread = 0x1p-12 }, &v, &wb, Exact{}, &rows[0], &rows[1], &ys[0], &ys[1]); + for (rows[0], rows[1], ys[0], ys[1]) |a, b, c, d| try std.testing.expect(a <= b and c <= d); +} diff --git a/zig/src/core/qk_ref.zig b/zig/src/core/qk_ref.zig index b1763d9f34..f5faef8c4c 100644 --- a/zig/src/core/qk_ref.zig +++ b/zig/src/core/qk_ref.zig @@ -1,5 +1,6 @@ //! QK norm and RoPE of one head in float64, the reference for every backend's attention inputs (cuda/qk_rows.zig). const std = @import("std"); +const bf = @import("bf16_round.zig"); /// One head: RMS-normed by w (none: as is), then rotate-half RoPE on its first `rot` dims (no cos/sin: none). pub fn head(x: []const f64, w: ?[]const f64, eps: f64, cos: ?[]const f64, sin: ?[]const f64, rot: usize, y: []f64) void { @@ -7,7 +8,7 @@ pub fn head(x: []const f64, w: ?[]const f64, eps: f64, cos: ?[]const f64, sin: ? var ss: f64 = 0; for (x) |v| ss += v * v; const inv = if (w != null) 1 / @sqrt(ss / @as(f64, @floatFromInt(n)) + eps) else 1; - var z: [256]f64 = undefined; + var z: [512]f64 = undefined; for (x, 0..) |v, i| z[i] = v * inv * (if (w) |ww| ww[i] else 1); const half = rot / 2; for (y, 0..) |*o, i| { @@ -21,6 +22,82 @@ pub fn head(x: []const f64, w: ?[]const f64, eps: f64, cos: ?[]const f64, sin: ? } } +/// The options of a head beyond `head`'s: MLX's bf16 rows (bf16(x inv) before the weight, w bf16(x inv) before RoPE). +pub const Points = struct { row: bool = false, out: bool = false }; + +/// An fp32 run's stray: the inverse RMS by `spread`, each product or sum by `u`, cos and sin by `trig` (0: float64). +pub const Stray = struct { spread: f64 = 0, u: f64 = 0, trig: f64 = 0 }; + +/// a and b in order; NaN in either gives NaN for both. +fn order(a: f64, b: f64) [2]f64 { + if (std.math.isNan(a) or std.math.isNan(b)) return @splat(std.math.nan(f64)); + return .{ @min(a, b), @max(a, b) }; +} + +/// The least and largest of the products of two intervals, widened by u of the largest magnitude. +fn mul(a: [2]f64, b: [2]f64, u: f64) [2]f64 { + var r = [2]f64{ std.math.inf(f64), -std.math.inf(f64) }; + for (a) |x| for (b) |y| { + const p = x * y; + if (std.math.isNan(p)) return @splat(std.math.nan(f64)); + r = .{ @min(r[0], p), @max(r[1], p) }; + }; + return widen(r, u); +} + +fn widen(r: [2]f64, u: f64) [2]f64 { + const m = @max(@abs(r[0]), @abs(r[1])) * u; + return .{ r[0] - m, r[1] + m }; +} + +/// A normed value's ends over the inverse's spread: (x inv) w in fp32, or w bf16(x inv) (`row`); bf16 with `out`. +fn normed(x: f64, inv: f64, w: f64, p: Points, st: Stray) [2]f64 { + var r = [2]f64{ std.math.inf(f64), -std.math.inf(f64) }; + const a = st.spread + st.u + st.spread * st.u; + for ([_]f64{ 1 - a, 1 + a }) |f| for ([_]f64{ 1 - st.u, 1 + st.u }) |f2| { + const v = x * inv * f; + var y = if (p.row) w * bf.bf16(v) else v * w * f2; + if (p.out) y = bf.bf16(y); + if (std.math.isNan(y)) return @splat(std.math.nan(f64)); + r = .{ @min(r[0], y), @max(r[1], y) }; + }; + return r; +} + +/// One head's bf16 outputs as ranges: normed by w (none: as is) from `ss` at points `p`, then RoPE on `rot` dims. +pub fn headRange(x: []const f64, w: ?[]const f64, ss: f64, eps: f64, cos: ?[]const f64, sin: ?[]const f64, rot: usize, p: Points, st: Stray, lo: []f64, hi: []f64) void { + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + var z: [512][2]f64 = undefined; + for (x, 0..) |v, i| z[i] = if (w) |ww| normed(v, inv, ww[i], p, st) else .{ v, v }; + const half = rot / 2; + for (lo, hi, 0..) |*l, *h, i| { + var t = z[i]; + if (cos != null and i < rot) { + const pi = if (i < half) i + half else i - half; + const zp = if (i < half) [2]f64{ -z[pi][1], -z[pi][0] } else z[pi]; + const c = cos.?[i]; + const sn = sin.?[i]; + const a = mul(t, order(c - st.trig * @abs(c), c + st.trig * @abs(c)), st.u); + const b = mul(zp, order(sn - st.trig * @abs(sn), sn + st.trig * @abs(sn)), st.u); + t = widen(.{ a[0] + b[0], a[1] + b[1] }, st.u); + } + l.*, h.* = order(bf.bf16(t[0]), bf.bf16(t[1])); + } +} + +/// A v head by a unit RMS norm from its sum of squares: bf16(x inv), as ranges over the inverse's spread. +pub fn unitRange(x: []const f64, ss: f64, eps: f64, st: Stray, lo: []f64, hi: []f64) void { + const inv = 1 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + for (x, lo, hi) |v, *l, *h| l.*, h.* = normed(v, inv, 1, .{ .out = true }, .{ .spread = st.spread, .u = st.u }); +} + +/// A head's sum of squares in float64. +pub fn squares(x: []const f64) f64 { + var s: f64 = 0; + for (x) |v| s += v * v; + return s; +} + test "rotate-half pairs (d, d + rot / 2) and leave the rest" { var y: [4]f64 = undefined; head(&.{ 1, 2, 3, 4 }, null, 0, &.{ 0, 0 }, &.{ 1, 1 }, 2, &y); // angle 90 degrees on the first pair @@ -28,3 +105,39 @@ test "rotate-half pairs (d, d + rot / 2) and leave the rest" { head(&.{ 2, 2, 2, 2 }, &.{ 1, 1, 1, 3 }, 0, null, null, 0, &y); try std.testing.expectEqualSlices(f64, &.{ 1, 1, 1, 3 }, &y); } + +test "rounding points: no options is head's value, MLX's rows round before the weight and before RoPE" { + var prng = std.Random.DefaultPrng.init(7); + var x: [64]f64 = undefined; + var w: [64]f64 = undefined; + var cs: [64]f64 = undefined; + var sn: [64]f64 = undefined; + for (&x, &w, &cs, &sn, 0..) |*a, *b, *c, *s, i| { + a.* = bf.bf16(prng.random().floatNorm(f64)); + b.* = bf.bf16(0.9 + 0.2 * prng.random().float(f64)); + const ang = 0.37 * @as(f64, @floatFromInt(i % 16 + 1)); + c.* = @cos(ang); + s.* = @sin(ang); + } + var want: [64]f64 = undefined; + var lo: [64]f64 = undefined; + var hi: [64]f64 = undefined; + head(&x, &w, 1e-6, &cs, &sn, 32, &want); + headRange(&x, &w, squares(&x), 1e-6, &cs, &sn, 32, .{}, .{}, &lo, &hi); + for (want, lo, hi) |v, l, h| { + try std.testing.expectEqual(l, h); + try std.testing.expectEqual(bf.bf16(v), l); + } + headRange(&x, &w, squares(&x), 1e-6, null, null, 0, .{ .row = true, .out = true }, .{}, &lo, &hi); + const inv = 1 / @sqrt(squares(&x) / 64 + 1e-6); + for (x, w, lo) |v, ww, l| try std.testing.expectEqual(bf.bf16(ww * bf.bf16(v * inv)), l); + headRange(&x, &w, squares(&x), 1e-6, &cs, &sn, 64, .{ .row = true, .out = true }, .{ .spread = 0x1p-20, .u = 0x1p-24, .trig = 0x1p-22 }, &lo, &hi); + var wide: usize = 0; + for (lo, hi) |l, h| { + try std.testing.expect(l <= h); + if (l != h) wide += 1; + } + try std.testing.expect(wide < 16); // most outputs stay decided + unitRange(&x, squares(&x), 1e-6, .{}, &lo, &hi); + for (x, lo) |v, l| try std.testing.expectEqual(bf.bf16(v * inv), l); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 721310a5eb..17d0804785 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -27,6 +27,8 @@ pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_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"); /// The kernel contract every GPU kernel's test registers with (zig/tests/contract, tf-cuda-test kernel-contract). pub const contract = @import("contract.zig"); pub const nvfp4 = @import("nvfp4.zig"); @@ -40,6 +42,7 @@ pub const drafter_release = @import("drafter_release.zig"); pub const drafter_layout = @import("drafter_layout.zig"); pub const tree_round = @import("tree_round.zig"); pub const attention_ref = @import("attention_ref.zig"); +pub const attention_range = @import("attention_range.zig"); pub const gdn_prep_ref = @import("gdn_prep_ref.zig"); pub const qk_ref = @import("qk_ref.zig"); pub const moe_ref = @import("moe_ref.zig"); @@ -64,9 +67,12 @@ test { _ = tuning; _ = norms_ref; _ = act_ref; + _ = bf16_round; + _ = embed_ref; _ = contract; _ = tree_round; _ = attention_ref; + _ = attention_range; _ = gdn_prep_ref; _ = qk_ref; _ = moe_ref; diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig index 1d343f85f3..e23223e7d5 100644 --- a/zig/src/cuda/act_rows.zig +++ b/zig/src/cuda/act_rows.zig @@ -1,4 +1,4 @@ -//! act_rows.cu's activation rows for every CUDA family: SwiGLU, sigmoid output gates and 64-value group sums. +//! act_rows.cu's activation rows for every CUDA family: SwiGLU, GeGLU, sigmoid output gates and 64-value group sums. const std = @import("std"); const kernels = @import("kernels.zig"); const module = @import("module.zig"); @@ -28,12 +28,15 @@ pub const ActRows = struct { gate_fn: Function, sums_fn: Function, dconv_fn: Function, + geglu_fn: Function, + swiglu_round_fn: Function, + softcap_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") }; + 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") }; } pub fn unload(g: *ActRows) void { @@ -42,6 +45,20 @@ pub const ActRows = struct { /// out [rows, n] = bf16(silu(gate) * up), rows at strides g_ld, u_ld; xs (or null): out's 64-value sums. pub fn swiglu(g: *const ActRows, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { + return g.glu(g.swiglu_fn, s, gate, g_ld, up, u_ld, out, xs, rows, n); + } + + /// swiglu in MLX's bf16 steps (tf_draft_swiglu's: exp, 1 + e, its reciprocal, the sigmoid, silu, each rounded). + pub fn swigluRound(g: *const ActRows, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { + return g.glu(g.swiglu_round_fn, s, gate, g_ld, up, u_ld, out, xs, rows, n); + } + + /// out = bf16(gelu_tanh(gate) * up) as mlx_lm's geglu in MLX's bf16 ops, rows as swiglu's; xs (or null): sums. + pub fn geglu(g: *const ActRows, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { + return g.glu(g.geglu_fn, s, gate, g_ld, up, u_ld, out, xs, rows, n); + } + + fn glu(_: *const ActRows, f: Function, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { if (n % 64 != 0 or !aligned(&.{ gate, up, out }, &.{ g_ld, u_ld })) return error.RowShape; if (rows == 0) return; var a: launch_.Args = .{}; @@ -53,7 +70,20 @@ pub const ActRows = struct { a.add(xs orelse 0); a.add(@as(i32, @intCast(n))); const t = launch_.rowThreads(rows, threads, 128); - try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n, t), .block = .{ .x = t } }, s, &a); + try launch_.launch(f, .{ .grid = grid(rows, n, t), .block = .{ .x = t } }, s, &a); + } + + /// out [rows, n] = bf16(cap * tanh(x / cap)) in MLX's bf16 steps (x / cap and tanh rounded); out may be x. + pub fn softcap(g: *const ActRows, s: Stream, x: u64, out: u64, rows: usize, n: usize, cap: f32) !void { + if (n % 8 != 0 or n > std.math.maxInt(i32) or !aligned(&.{ x, out }, &.{n})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(out); + a.add(@as(i32, @intCast(n))); + a.add(cap); + const t = launch_.rowThreads(rows, threads, 128); + try launch_.launch(g.softcap_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((n / 8 + t - 1) / t), .z = 1 }, .block = .{ .x = t } }, s, &a); } /// out = bf16(o * sigmoid(gate)), head h's gate at r * g_ld + h * g_hs ([q | gate]: q + d, 2 d); xs: sums. diff --git a/zig/src/cuda/embed_rows.zig b/zig/src/cuda/embed_rows.zig index 1fc5cda5b7..38771adc0c 100644 --- a/zig/src/cuda/embed_rows.zig +++ b/zig/src/cuda/embed_rows.zig @@ -1,4 +1,4 @@ -//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 32 or 64) or bf16 tables. +//! embed_rows.cu's token rows for every CUDA family: MLX affine (2 to 6 or 8 bits, a bf16 scale) or bf16 tables. const std = @import("std"); const kernels = @import("kernels.zig"); const module = @import("module.zig"); @@ -17,12 +17,13 @@ pub const EmbedRows = struct { affine4_fn: Function, bf16_fn: Function, bf16_copies_fn: Function, + affine_fn: Function, pub fn load(d: *const Driver) !EmbedRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.embed_rows); errdefer mod.unload(); - return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16"), .bf16_copies_fn = try mod.function("tf_embed_rows_bf16_copies") }; + return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16"), .bf16_copies_fn = try mod.function("tf_embed_rows_bf16_copies"), .affine_fn = try mod.function("tf_embed_rows_affine") }; } pub fn unload(g: *EmbedRows) void { @@ -44,6 +45,24 @@ pub const EmbedRows = struct { try launch_.launch(g.affine4_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); } + /// `affine`'s table: codes `bits` wide (2 to 6 or 8) in groups of `group`; `scale` (a bf16 value) on each value. + pub const Affine = struct { bits: u4 = 4, group: usize = 64, scale: ?f32 = null }; + + /// out [rows, d] = MLX's bf16(q s + b) of w u32 [vocab, d bits / 32] at ids, times o.scale and rounded again. + pub fn affine(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize, o: Affine) !void { + const bits_ok = o.bits >= 2 and o.bits <= 8 and o.bits != 7; + const scale_ok = if (o.scale) |x| std.math.isFinite(x) and @as(u32, @bitCast(x)) & 0xFFFF == 0 else true; + if (!bits_ok or !scale_ok or (o.group != 32 and o.group != 64 and o.group != 128) or d == 0 or d % o.group != 0 or + d > std.math.maxInt(i32) or rows > std.math.maxInt(u32) or out % 16 != 0) return error.RowShape; + if (o.bits == 4 and o.scale == null and o.group != 128) return g.affine4Groups(s, ids, w, scales, biases, out, rows, d, o.group, 1); + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ ids, w, scales, biases, out }) |v| a.add(v); + for ([_]usize{ d, o.bits, o.group }) |v| a.add(@as(i32, @intCast(v))); + a.add(o.scale orelse 1); + try launch_.launch(g.affine_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + /// out [rows, d] = w [vocab, d] bf16 rows at ids. pub fn bf16(g: *const EmbedRows, s: Stream, ids: u64, w: u64, out: u64, rows: usize, d: usize) !void { if (d % 8 != 0 or out % 16 != 0 or w % 16 != 0) return error.RowShape; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 80c5f50d6f..40d9846fe5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -50,6 +50,7 @@ pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes e pub const tree_attention: []const u8 = if (available) &Blob("fatbin_tree_attention").bytes else &.{}; pub const block_attention: []const u8 = if (available) &Blob("fatbin_block_attention").bytes else &.{}; pub const window_attention: []const u8 = if (available) &Blob("fatbin_window_attention").bytes else &.{}; +pub const window_rows: []const u8 = if (available) &Blob("fatbin_window_rows").bytes else &.{}; pub const moe_glue: []const u8 = if (available) &Blob("fatbin_moe_glue").bytes else &.{}; pub const dense_rows: []const u8 = if (available) &Blob("fatbin_dense_rows").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/cuda/norms.zig b/zig/src/cuda/norms.zig index 6f774d50d3..4814b42240 100644 --- a/zig/src/cuda/norms.zig +++ b/zig/src/cuda/norms.zig @@ -33,6 +33,16 @@ fn vectorOk(width: usize, ptrs: []const u64) bool { return vectorWidth(width) and aligned(ptrs); } +/// addRmsTail's options: b = norm(x) wb (or x), h = bf16(r + b) [* scalar], y = bf16(norm(h) w), xs y's 64-value sums. +pub const Tail = struct { + /// Each norm as MLX's bf16 steps, bf16(w * bf16(v * inv)), where false gives bf16(v * inv * w). + round: bool = false, + /// A norm on x before the add (a sandwich norm): its bf16 weight; null adds x as it is. + branch: ?u64 = null, + /// One bf16 the sum is multiplied by, after the sum is rounded, before the output's norm; null: none. + scalar: ?u64 = null, +}; + pub const Norms = struct { mod: Module, rms_fn: Function, @@ -43,6 +53,7 @@ pub const Norms = struct { add_rms_h_f32w_fn: Function, 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 pub fn load(d: *const Driver) !Norms { if (!kernels.available) return error.BuiltWithoutKernels; @@ -58,6 +69,7 @@ pub const Norms = struct { .add_rms_h_f32w_fn = try mod.function("tf_add_rms_rows_bf16h_f32w"), .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") }, }; } @@ -143,6 +155,19 @@ pub const Norms = struct { try launch_.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = vec_threads } }, s, &a); } + /// addRmsH with bf16 weights and a tail's options; with none set it is addRmsH itself, the same kernel and bits. + pub fn addRmsTail(n: *const Norms, s: Stream, x: u64, r: ?u64, w: u64, h: u64, y: u64, xs: ?u64, rows: usize, width: usize, eps: f32, t: Tail) !void { + if (!t.round and t.branch == null and t.scalar == null) return n.addRmsH(s, x, r, w, .bf16, h, y, xs, rows, width, eps); + if (width % 64 != 0 or !vectorOk(width, &.{ x, r orelse 0, t.branch orelse 0, w, h, y }) or (t.scalar orelse 0) % 2 != 0) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ x, r orelse 0, t.branch orelse 0, t.scalar orelse 0, w, h, y, xs orelse 0 }) |v| a.add(v); + a.add(@as(i32, @intCast(width))); + a.add(eps); + const f = n.tail_fns[@intFromBool(t.round)]; + try launch_.launch(f, .{ .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 d1ccb1a4ed..528baa233b 100644 --- a/zig/src/cuda/qk_rows.zig +++ b/zig/src/cuda/qk_rows.zig @@ -10,14 +10,27 @@ const Module = module.Module; const Function = module.Function; const Stream = stream_.Stream; -/// A packed row's heads: h query, hk key and hk value heads of `dim` (a multiple of 32, at most 256). +/// A packed row's heads: h query, hk key and hk value heads of `dim` (a multiple of 32, at most 512). pub const Heads = struct { h: usize, hk: usize, dim: usize }; /// Where a row's heads come from: q, k and v (0: none) at row stride *_ld and head stride *_hs (elements). pub const Src = extern struct { q: u64, q_ld: i64, q_hs: i64, k: u64, k_ld: i64, k_hs: i64, v: u64 = 0, v_ld: i64 = 0, v_hs: i64 = 0 }; /// q/k norms [2, dim] (or q's with norm_k), fp32 or bf16; RoPE on `rot` dims: cos/sin tables or inv [rot / 2]. -pub const Shape = struct { norms: ?u64 = null, norm_k: ?u64 = null, norms_bf16: bool = false, cos: ?u64 = null, sin: ?u64 = null, inv: ?u64 = null, rot: usize = 0, eps: f32 = 1e-6 }; +pub const Shape = struct { + norms: ?u64 = null, + norm_k: ?u64 = null, + norms_bf16: bool = false, + cos: ?u64 = null, + sin: ?u64 = null, + inv: ?u64 = null, + rot: usize = 0, + eps: f32 = 1e-6, + /// MLX's rows: q and k = bf16(w * bf16(x * inv)) before RoPE. + round: bool = false, + /// v by a unit RMS norm, bf16(x * inv); else v as is. + v_norm: bool = false, +}; /// qk_rows.cu's QkW. const W = extern struct { q: u64, k: u64, bf16: i32 }; @@ -45,12 +58,15 @@ pub const QkRows = struct { fn_: Function, /// tf_qk_rows_w: dim 128 with no or full rotation, a warp a head. warp_fn: Function, + /// 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, 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") }; + 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") }; } pub fn unload(g: *QkRows) void { @@ -67,7 +83,12 @@ pub const QkRows = struct { /// As `rows`, the heads read from `src` (e.g. [q | gate] rows: q_hs = 2 dim); no v source: no v out or vc. pub fn rowsFrom(g: *const QkRows, s: Stream, src: Src, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { - if (hs.dim % 32 != 0 or hs.dim > 256 or shape.rot % 2 != 0 or shape.rot > hs.dim) return error.HeadShape; + return g.rowsVia(null, 0, s, src, hs, shape, pos, slot, out, n); + } + + /// rowsFrom through `f` (null: the kernel the shape selects) on tf_qk_rows_round's arguments, `last` its final int. + pub fn rowsVia(g: *const QkRows, f: ?Function, last: i32, s: Stream, src: Src, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + if (hs.dim % 32 != 0 or hs.dim > 512 or shape.rot % 2 != 0 or shape.rot > hs.dim) return error.HeadShape; if ((shape.cos == null) != (shape.sin == null) or (shape.cos != null and shape.inv != null)) return error.HeadShape; const roped = shape.cos != null or shape.inv != null; if (roped and shape.rot == 0) return error.HeadShape; @@ -97,6 +118,11 @@ pub const QkRows = struct { 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); + if (f != null or hs.dim > 256 or shape.round or shape.v_norm) { + a.add(if (f != null) last else @as(i32, @intFromBool(shape.v_norm))); + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast(heads), .z = 1 }; + return launch_.launch(f orelse if (shape.round) g.round_fn else g.once_fn, .{ .grid = grid, .block = .{ .x = @intCast(hs.dim) } }, s, &a); + } if (hs.dim == 128 and (!roped or shape.rot == 128)) { const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast((heads + 3) / 4), .z = 1 }; return launch_.launch(g.warp_fn, .{ .grid = grid, .block = .{ .x = 128 } }, s, &a); diff --git a/zig/src/cuda/window_attention.zig b/zig/src/cuda/window_attention.zig index 40569b7ef9..7a17eb5acf 100644 --- a/zig/src/cuda/window_attention.zig +++ b/zig/src/cuda/window_attention.zig @@ -10,8 +10,20 @@ const Module = module.Module; const Function = module.Function; const Stream = stream_.Stream; -/// The heads and the keys a row sees: (pos - win, pos] (win 0: all), cache row slot * ring + key % ring. -pub const Window = struct { h: usize, hk: usize, d: usize, win: usize, ring: usize, scale: f32 }; +/// A row's heads and keys (pos - win, pos] (win 0: all) at cache row slot * ring + key % ring; p32: wide P in halves. +pub const Window = struct { h: usize, hk: usize, d: usize, win: usize, ring: usize, scale: f32, p32: bool = false }; + +/// Head dims window_rows.cu serves, with any query heads a key head up to 16. +pub fn wide(d: usize) bool { + return d == 256 or d == 512; +} + +/// window_rows.cu's decode chunk: keys [c * chunk_keys, (c + 1) * chunk_keys), merged in chunk order. +pub const chunk_keys = 256; +/// The head dim 512 decode partials' staged K and V tiles: two of 32 keys. +const staged_smem: u32 = 2 * 32 * 512 * 2; +/// Prompt rows' dynamic shared bytes at 256 and 512: 4 or 2 slots of 32 keys, then the warps' partial scores. +const rows_smem = [_]u32{ 4 * 32 * 32 * 16 + 8 * 16 * 32 * 4, 2 * 32 * 64 * 16 + 8 * 16 * 32 * 4 }; /// The prompt kernel's instances: head dim 128, 8 warps, 4 or 8 query heads a block (`hpc`), 8 staging slots. const prompt_name = "_ZN16tf_window_prompt12wattn_kernelILi128ELi8ELi{d}ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifii"; @@ -25,6 +37,11 @@ pub const WindowAttention = struct { merge_fn: Function, tc_fn: Function, prompt_fns: [2]Function, // 4, then 8 query heads a block + rows: Module, // window_rows.cu + wide_fns: [rows_smem.len]Wide, + + /// window_rows.cu's kernels at one head dim: decode partials and merge, prompt rows (bf16 P, then P32). + pub const Wide = struct { part: Function, merge: Function, prompt: [2]Function }; pub fn load(d: *const Driver) !WindowAttention { if (!kernels.available) return error.BuiltWithoutKernels; @@ -35,17 +52,40 @@ pub const WindowAttention = struct { p[i] = try mod.function(std.fmt.comptimePrint(prompt_name, .{hpc})); try p[i].allowDynamicShared(prompt_smem); } + var rows = try Module.load(d, kernels.window_rows); + errdefer rows.unload(); + var wf: [rows_smem.len]Wide = undefined; + inline for (.{ 256, 512 }, &wf, rows_smem) |dim, *f, smem| { + f.* = .{ .part = try rows.function(std.fmt.comptimePrint("tf_wrows_part_{d}", .{dim})), .merge = try rows.function(std.fmt.comptimePrint("tf_wrows_merge_{d}", .{dim})), .prompt = .{ try rows.function(std.fmt.comptimePrint("tf_wrows_prompt_{d}", .{dim})), try rows.function(std.fmt.comptimePrint("tf_wrows_prompt_{d}_p32", .{dim})) } }; + for (f.prompt) |pf| try pf.allowDynamicShared(smem); + if (dim == 512) try f.part.allowDynamicShared(staged_smem); + } return .{ .mod = mod, .part_fn = try mod.function("tf_window_part"), .merge_fn = try mod.function("tf_window_merge"), .tc_fn = tc: { const f = try mod.function("tf_window_decode_tc"); try f.allowDynamicShared(tc_smem); break :tc f; - }, .prompt_fns = p }; + }, .prompt_fns = p, .rows = rows, .wide_fns = wf }; } pub fn unload(g: *WindowAttention) void { + g.rows.unload(); g.mod.unload(); } + /// The dynamic shared bytes window_rows.cu's prompt rows take at head dim d. + pub fn promptShared(d: usize) u32 { + return rows_smem[@intFromBool(d == 512)]; + } + + pub fn wideFns(g: *const WindowAttention, d: usize) *const Wide { + return &g.wide_fns[@intFromBool(d == 512)]; + } + + /// Chunks a wide decode row's keys can touch: the window's, or (win 0) the ring's or linear cache's positions. + pub fn chunks(w: Window) usize { + return if (w.win > 0) (w.win + chunk_keys - 2) / chunk_keys + 1 else (w.ring + chunk_keys - 1) / chunk_keys; + } + /// Key tiles a decode row's window can touch (its partials a key head). pub fn tiles(w: Window) usize { return if (w.win > 0) (w.win + 62) / 64 + 1 else (w.ring + 63) / 64 + 1; @@ -53,6 +93,10 @@ pub const WindowAttention = struct { /// The fp32 partials `decode` writes for up to `rows` rows: o, then m and l. pub fn scratchBytes(w: Window, rows: usize) usize { + if (wide(w.d)) { + const parts = rows * w.hk * chunks(w) * (w.h / w.hk); + return parts * w.d * 4 + 2 * parts * 4; + } if (w.d == 128) return 16; // the tensor-core form keeps its partials in shared memory const parts = rows * w.hk * tiles(w) * (w.h / w.hk); return parts * w.d * 4 + 2 * parts * 4; @@ -60,6 +104,7 @@ pub const WindowAttention = struct { /// n rows q [n, h, d] at pos/slot (int32 [n]) -> out [n, h, d]; `scratch`: scratchBytes(w, n) (d != 128). pub fn decode(g: *const WindowAttention, s: Stream, q: u64, kc: u64, vc: u64, out: u64, pos: u64, slot: ?u64, w: Window, scratch: u64, n: usize) !void { + if (wide(w.d)) return decodeWith(g.wideFns(w.d), s, q, kc, vc, out, pos, slot, w, scratch, n); if (w.hk == 0 or w.h % w.hk != 0 or w.h / w.hk > 16 or w.d > 128 or w.ring == 0) return error.WindowShape; if (w.win == 0 and w.ring > 4096) return error.WindowShape; // an unwindowed row's tiles come from the ring if (n == 0) return; @@ -96,8 +141,39 @@ pub const WindowAttention = struct { try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = 1 }, .block = .{ .x = 256 } }, s, &b); } + /// Wide decode rows by `f` (wideFns, or the contract's controls): chunk partials (scratchBytes), their merge. + pub fn decodeWith(f: *const Wide, s: Stream, q: u64, kc: u64, vc: u64, out: u64, pos: u64, slot: ?u64, w: Window, scratch: u64, n: usize) !void { + if (!wide(w.d) or w.hk == 0 or w.h % w.hk != 0 or w.h / w.hk > 16 or w.ring == 0 or w.ring > std.math.maxInt(i32)) return error.WindowShape; + if (n == 0) return; + const t = chunks(w); + const parts = n * w.hk * t * (w.h / w.hk); + const po = scratch; + const pm = po + parts * w.d * 4; + const pl = pm + parts * 4; + var a: launch_.Args = .{}; + for ([_]u64{ q, kc, vc, pos, slot orelse 0, po, pm, pl }) |v| a.add(v); + a.add(@as(i32, @intCast(w.h))); + a.add(@as(i32, @intCast(w.hk))); + a.add(@as(i32, @intCast(w.win))); + a.add(@as(i64, @intCast(w.ring))); + a.add(w.scale); + const pairs = if (w.d == 512) (w.h / w.hk + 3) / 4 else 1; // head dim 512: a block 4 query heads + try launch_.launch(f.part, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk * pairs), .z = @intCast(t) }, .block = .{ .x = 256 }, .shared = if (w.d == 512) staged_smem else 0 }, s, &a); + var b: launch_.Args = .{}; + for ([_]u64{ po, pm, pl, out }) |v| b.add(v); + b.add(@as(i32, @intCast(w.h))); + b.add(@as(i32, @intCast(w.hk))); + b.add(@as(i32, @intCast(t))); + for ([_]u64{ q, kc, vc, pos, slot orelse 0 }) |v| b.add(v); + b.add(@as(i32, @intCast(w.win))); + b.add(@as(i64, @intCast(w.ring))); + b.add(w.scale); + try launch_.launch(f.merge, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = @intCast(w.h / w.hk * w.d / 256) }, .block = .{ .x = 256 } }, s, &b); + } + /// n prompt rows of one stream at p0, p0 + 1, ... (d 128); kc/vc that stream's ring, key p at row p % ring. pub fn prompt(g: *const WindowAttention, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: Window) !void { + if (wide(w.d)) return g.promptRows(s, q, kc, vc, out, p0, n, w, null); const grp = w.h / w.hk; if (w.d != 128 or w.hk == 0 or w.h % w.hk != 0 or grp % 4 != 0) return error.WindowShape; if (n == 0) return; @@ -112,4 +188,28 @@ pub const WindowAttention = struct { const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((n + rows - 1) / rows), .y = @intCast(w.hk * (grp / hpc)), .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; try launch_.launch(g.prompt_fns[if (hpc == 8) 1 else 0], cfg, s, &a); } + + /// Wide prompt rows, G up to 16; row i's last key hi[i] (int32, in [p0 + i, p0 + n)) or its position (null). + pub fn promptRows(g: *const WindowAttention, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: Window, hi: ?u64) !void { + if (!wide(w.d)) return error.WindowShape; + return promptWith(g.wideFns(w.d), s, q, kc, vc, out, p0, n, w, hi); + } + + /// promptRows by `f` (wideFns, or the contract's controls). + pub fn promptWith(f: *const Wide, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: Window, hi: ?u64) !void { + const grp = if (w.hk > 0) w.h / w.hk else 0; + if (!wide(w.d) or w.hk == 0 or w.h % w.hk != 0 or grp > 16 or p0 + n > std.math.maxInt(i32)) return error.WindowShape; + if (n == 0) return; + const ng = 1024 / w.d; // 16-row groups a block: 8 warps, a 128-dim slice each + const hb = if (grp % ng == 0) ng else 1; + const rows = 16 * (ng / hb); + var a: launch_.Args = .{}; + for ([_]u64{ q, kc, vc, out, hi orelse 0 }) |v| a.add(v); + for ([_]usize{ p0, n, w.h, w.hk, grp }) |v| a.add(@as(i32, @intCast(v))); + a.add(w.scale); + a.add(@as(i32, @intCast(w.win))); + a.add(@as(i32, @intCast(w.ring))); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((n + rows - 1) / rows), .y = @intCast(w.hk * (grp / hb)), .z = 1 }, .block = .{ .x = 256 }, .shared = promptShared(w.d) }; + try launch_.launch(f.prompt[@intFromBool(w.p32)], cfg, s, &a); + } }; diff --git a/zig/src/families/gemma4/cuda.zig b/zig/src/families/gemma4/cuda.zig new file mode 100644 index 0000000000..9fcd22c6bf --- /dev/null +++ b/zig/src/families/gemma4/cuda.zig @@ -0,0 +1,43 @@ +//! Gemma 4 on CUDA (#1048 step 8); cuda_config/cuda_widths copy config/widths (a file belongs to one module). + +const std = @import("std"); + +pub const config = @import("cuda_config.zig"); +pub const widths = @import("cuda_widths.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const forward = @import("cuda_forward.zig"); +pub const snap = @import("cuda_snap.zig"); +pub const backend = @import("cuda_backend.zig"); +pub const assistant = @import("cuda_assistant.zig"); +pub const native = @import("cuda_native.zig").Family("gemma4_unified"); +pub const native31 = @import("cuda_native.zig").Family("gemma4"); + +test "cuda_config.zig and cuda_widths.zig are config.zig and widths.zig but for their first lines (doc and imports)" { + // the copies' width tests log the refusals they plant, as Metal's do: warnings, meant + std.testing.log_level = .err; + const pairs = [_][2][]const u8{ .{ @embedFile("config.zig"), @embedFile("cuda_config.zig") }, .{ @embedFile("widths.zig"), @embedFile("cuda_widths.zig") } }; + for (pairs) |p| { + var a = std.mem.splitScalar(u8, p[0], '\n'); + var b = std.mem.splitScalar(u8, p[1], '\n'); + var line: usize = 0; + while (a.next()) |x| : (line += 1) { + const y = b.next() orelse return error.CopyShorter; + if (line >= 4 and !std.mem.eql(u8, x, y)) return error.CopyDiffers; + } + try std.testing.expect(b.next() == null); + } +} + +test { + _ = @import("cuda_config.zig"); + _ = @import("cuda_widths.zig"); + _ = @import("cuda_weights.zig"); + _ = @import("cuda_kernels.zig"); + _ = @import("cuda_forward.zig"); + _ = @import("cuda_snap.zig"); + _ = @import("cuda_backend.zig"); + _ = @import("cuda_assistant.zig"); + _ = @import("cuda_assistant_weights.zig"); + _ = @import("cuda_native.zig"); +} diff --git a/zig/src/families/gemma4/cuda_assistant.zig b/zig/src/families/gemma4/cuda_assistant.zig new file mode 100644 index 0000000000..87ad98bc25 --- /dev/null +++ b/zig/src/families/gemma4/cuda_assistant.zig @@ -0,0 +1,355 @@ +//! Gemma 4's MTP assistant on CUDA (#1048 step 10): chains of drafts over the target's own last caches, as Metal's. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("cuda_config.zig"); +const fwd = @import("cuda_forward.zig"); +const kern = @import("cuda_kernels.zig"); +const be = @import("cuda_backend.zig"); + +const dr = lanes.drafter; +const Stream = lanes.Stream; +const Affine4 = cuda.qlinear.Affine4; +const Pick = cuda.tree_round.Pick; +const Window = cuda.window_attention.Window; + +/// Drafts a round at most (the reference's default; the head was trained for 3). +pub const depth = 6; +/// Streams one hold takes. +const max_items = 32; +const aw = @import("cuda_assistant_weights.zig"); +pub const Config = aw.Config; +pub const isAssistant = aw.isAssistant; +const Weights = aw.Weights; +const load = aw.load; + +/// One stream's state: the target's final-normed last kept row, its position, the token after it. +const Lane = struct { + h: cuda.DeviceBuffer, + position: u64 = 0, + token: u32 = 0, + rows: bool = false, // a row was absorbed + held: [depth]u32 = @splat(0), + held_len: u32 = 0, +}; + +/// A stream's chain in a hold: its lane, its target cache slot, the pending token, the kept row's position. +const Item = struct { lane: *Lane, slot: usize, pending: u32, position: u64, depth: u32 }; + +pub const Drafter = struct { + gpa: std.mem.Allocator, + io: std.Io, + d: *const cuda.Driver, + m: *fwd.Model, + k: *const kern.Kernels, + gb: *be.Gemma, + s: cuda.Stream, + c: Config, + w: Weights, + lanes: std.AutoHashMapUnmanaged(*Stream, *Lane) = .empty, + /// The target's layers each kind reads: its last sliding and its last global layer. + sources: [2]usize, + tap: [1]u32, + step_ms: f64 = 0, + /// The target's forward ms by window width (1 to 16), timed at load: the depth rule prices drafts by them. + costs: [be.max_window]lanes.config.Cost = undefined, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + x: Scratch, + + const Scratch = struct { ids: u64, pos: u64, slot: u64, picks: u64, e: u64, hin: u64, cat: u64, xs: u64, h: [2]u64, hs: u64, hn: u64, normed: u64, nmlp: u64, qraw: u64, q: u64, kj: u64, att: u64, o: u64, gate: u64, up: u64, act: u64, y: u64, part: u64, logits: u64 }; + + pub fn init(gpa: std.mem.Allocator, io: std.Io, gb: *be.Gemma, dir: []const u8) !*Drafter { + const m = gb.m; + const geo = m.geo; + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + const c = try Config.parse(gpa, bytes, geo); + const x = try gpa.create(Drafter); + errdefer gpa.destroy(x); + x.* = .{ .gpa = gpa, .io = io, .d = gb.d, .m = m, .k = gb.k, .gb = gb, .s = gb.s, .c = c, .w = undefined, .sources = .{ 0, 0 }, .tap = .{@intCast(m.layers - 1)}, .x = undefined }; + for (0..m.layers) |i| x.sources[@intFromBool(cfg.global(i))] = i; + x.w = try load(gpa, io, x.d, x.s, x.k, dir, &c, geo); + errdefer x.w.deinit(gpa); + errdefer x.freeScratch(); + const n = max_items; + const hid = c.hidden; + const qw = c.heads * cfg.global_dim; + const wide: Window = .{ .h = c.heads, .hk = geo.global_kv_heads, .d = cfg.global_dim, .win = 0, .ring = m.context, .scale = 1 }; + const narrow: Window = .{ .h = c.heads, .hk = geo.sliding_kv_heads, .d = cfg.sliding_dim, .win = cfg.window, .ring = fwd.ring, .scale = 1 }; + const part = @max(cuda.window_attention.WindowAttention.scratchBytes(wide, 16), cuda.window_attention.WindowAttention.scratchBytes(narrow, 16)); + x.x = .{ + .ids = try x.alloc(n * 4), + .pos = try x.alloc(n * 4), + .slot = try x.alloc(n * 4), + .picks = try x.alloc(depth * n * @sizeOf(Pick)), + .e = try x.alloc(n * geo.hidden * 2), + .hin = try x.alloc(n * geo.hidden * 2), + .cat = try x.alloc(n * 2 * geo.hidden * 2), + .xs = try x.alloc(n * (@max(2 * geo.hidden, @max(c.intermediate, qw)) / 64) * 4), + .h = .{ try x.alloc(n * hid * 2), try x.alloc(n * hid * 2) }, + .hs = try x.alloc(n * geo.hidden * 2), + .hn = try x.alloc(n * hid * 2), + .normed = try x.alloc(n * hid * 2), + .nmlp = try x.alloc(n * hid * 2), + .qraw = try x.alloc(n * qw * 2), + .q = try x.alloc(n * qw * 2), + .kj = try x.alloc(n * cfg.global_dim * 2), + .att = try x.alloc(n * qw * 2), + .o = try x.alloc(n * hid * 2), + .gate = try x.alloc(n * c.intermediate * 2), + .up = try x.alloc(n * c.intermediate * 2), + .act = try x.alloc(n * c.intermediate * 2), + .y = try x.alloc(n * hid * 2), + .part = try x.alloc(part), + .logits = try x.alloc(n * cfg.vocab * 2), + }; + try x.warm(); + return x; + } + + fn alloc(x: *Drafter, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(x.d, @max(bytes, 256)); + errdefer { + var y = b; + y.free(); + } + try x.buffers.append(x.gpa, b); + return b.ptr; + } + + fn freeScratch(x: *Drafter) void { + for (x.buffers.items) |*b| b.free(); + x.buffers.deinit(x.gpa); + } + + pub fn deinit(x: *Drafter) void { + x.s.synchronize() catch {}; + var it = x.lanes.valueIterator(); + while (it.next()) |l| dropLane(x, l.*); + x.lanes.deinit(x.gpa); + x.freeScratch(); + x.w.deinit(x.gpa); + x.gpa.destroy(x); + } + + pub fn drafter(x: *Drafter) dr.Drafter { + return .{ .ptr = x, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Drafter { + return @ptrCast(@alignCast(ptr)); + } + + /// The target's last layer: its output before the final norm, which absorb applies. + fn tapsFn(ptr: *anyopaque) []const u32 { + return &cast(ptr).tap; + } + + /// Chains of up to 6, one assistant forward a draft: the rule prices each draft and may choose none. + fn factsFn(ptr: *anyopaque) dr.Facts { + const x = cast(ptr); + return .{ .depth = depth, .step_ms = x.step_ms, .batched = true, .stream_bytes = x.m.geo.hidden * 2, .window_costs = &x.costs }; + } + + fn newLane(x: *Drafter) !*Lane { + const l = try x.gpa.create(Lane); + errdefer x.gpa.destroy(l); + l.* = .{ .h = try cuda.DeviceBuffer.alloc(x.d, x.m.geo.hidden * 2) }; + return l; + } + + fn dropLane(x: *Drafter, l: *Lane) void { + l.h.free(); + x.gpa.destroy(l); + } + + fn openFn(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = cast(ptr); + if (x.lanes.contains(s)) return error.DraftStreamOpen; + const l = try x.newLane(); + errdefer x.dropLane(l); + try x.lanes.put(x.gpa, s, l); + } + + fn releaseFn(ptr: *anyopaque, s: *Stream) void { + const x = cast(ptr); + if (x.lanes.fetchRemove(s)) |e| { + x.s.synchronize() catch {}; // a queued chain may still read the lane's row + x.dropLane(e.value); + } + } + + fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const x = cast(ptr); + for (streams, out) |s, o| { + const l = x.lanes.get(s) orelse return error.DraftStreamUnknown; + if (o.len > l.held_len) return error.DraftHeldFewer; + @memcpy(o, l.held[0..o.len]); + } + } + + /// Each stream's last absorbed row through the target's final norm into its lane, queued on the target's stream. + fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { + const x = cast(ptr); + const width = x.m.geo.hidden * 2; + for (items) |it| { + const l = x.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + const f = it.features; + if (f.space != .device or f.dtype != .bf16 or f.row_bytes != width or f.rows == 0 or it.follow.len < f.rows) return error.DraftFeaturesUnsupported; + const last = f.rows - 1; + try x.k.norms.addRmsTail(x.s, f.buffer + f.offset + last * f.row_bytes, null, x.m.w.final_norm, x.x.hs, l.h.ptr, null, 1, x.m.geo.hidden, cfg.eps, .{ .round = true }); + l.position = it.start + last; + l.token = it.follow[last]; + l.rows = true; + l.held_len = 0; + } + } + + /// x [rows, l.k] (its group sums at xs) through `l` into out [rows, l.n], in verify tiles. + fn project(x: *Drafter, l: Affine4, in: u64, xs: u64, out: u64, rows: usize) !void { + const tile = x.k.affine.verifyRows(); + var at: usize = 0; + while (at < rows) : (at += tile) { + const n = @min(tile, rows - at); + try x.k.affine.verify(x.s, in + at * l.k * 2, xs + at * (l.k / 64) * 4, l, out + at * l.n * 2, n); + } + } + + fn copy2d(x: *Drafter, dst: u64, dst_pitch: usize, src: u64, src_pitch: usize, width: usize, rows: usize) !void { + const p: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = src_pitch, .dst = dst, .dst_pitch = dst_pitch, .width = width, .height = rows }; + try x.d.check(x.d.api.cuMemcpy2DAsync_v2(&p, x.s.handle), "cuMemcpy2DAsync"); + } + + /// Assistant layer i on the rows: input normed in normed (sums in xs), residual h[cur] into h[1 - cur]. + fn layer(x: *Drafter, i: usize, cur: usize, rows: usize) !void { + const c = &x.c; + const geo = x.m.geo; + const l = &x.w.layers[i]; + const s = x.x; + const global = !c.sliding[i]; + const dh = c.headDim(i); + const qw = c.heads * dh; + try x.project(l.q, s.normed, s.xs, s.qraw, rows); + // q alone: k reads q's first head into scratch (no cache written), v none + const src: cuda.qk_rows.Src = .{ .q = s.qraw, .q_ld = @intCast(qw), .q_hs = @intCast(dh), .k = s.qraw, .k_ld = @intCast(qw), .k_hs = @intCast(dh) }; + const shape: cuda.qk_rows.Shape = .{ .norms = l.q_norm, .norm_k = l.q_norm, .norms_bf16 = true, .inv = if (global) x.m.inv_global else x.m.inv_sliding, .rot = dh, .eps = cfg.eps, .round = true }; + try x.k.qk.rowsFrom(x.s, src, .{ .h = c.heads, .hk = 1, .dim = dh }, shape, s.pos, null, .{ .q = s.q, .k = s.kj, .v = 0 }, rows); + const t = x.sources[@intFromBool(global)]; + const len = x.m.cacheLen(t); + const w: Window = .{ .h = c.heads, .hk = geo.kvHeads(t), .d = dh, .win = if (global) 0 else cfg.window, .ring = len, .scale = 1 }; + var r: usize = 0; + while (r < rows) : (r += 16) { + const n = @min(16, rows - r); + try x.k.win.decode(x.s, s.q + r * qw * 2, x.m.kc[t], x.m.vc[t], s.att + r * qw * 2, s.pos + r * 4, s.slot + r * 4, w, s.part, n); + } + try x.k.act.groupSums(x.s, s.att, qw, s.xs, rows, qw); + try x.project(l.o, s.att, s.xs, s.o, rows); + try x.k.norms.addRmsTail(x.s, s.o, s.h[cur], l.pre_feedforward, s.hn, s.nmlp, s.xs, rows, c.hidden, cfg.eps, .{ .round = true, .branch = l.post_attention }); + try x.project(l.gate, s.nmlp, s.xs, s.gate, rows); + try x.project(l.up, s.nmlp, s.xs, s.up, rows); + try x.k.act.geglu(x.s, s.gate, c.intermediate, s.up, c.intermediate, s.act, s.xs, rows, c.intermediate); + try x.project(l.down, s.act, s.xs, s.y, rows); + const next = if (i + 1 < c.layers) x.w.layers[i + 1].input_norm else x.w.norm; + try x.k.norms.addRmsTail(x.s, s.y, s.hn, next, s.h[1 - cur], s.normed, s.xs, rows, c.hidden, cfg.eps, .{ .round = true, .branch = l.post_feedforward, .scalar = l.scalar }); + } + + /// Every item's chain of `steps` drafts: step s's picks at picks[s * max_items + row], queued on the stream. + fn chain(x: *Drafter, items: []const Item, steps: u32) !void { + const geo = x.m.geo; + const s = x.x; + const rows = items.len; + const hid = geo.hidden; + var host: [3 * max_items]i32 = undefined; + for (items, 0..) |it, r| { + host[r] = @intCast(it.pending); + host[max_items + r] = @intCast(it.position); + host[2 * max_items + r] = @intCast(it.slot); + } + for ([_]u64{ s.ids, s.pos, s.slot }, 0..) |dst, j| try x.d.check(x.d.api.cuMemcpyHtoDAsync_v2(dst, &host[j * max_items], rows * 4, x.s.handle), "cuMemcpyHtoDAsync"); + for (items, 0..) |it, r| try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(s.hin + r * hid * 2, it.lane.h.ptr, hid * 2, x.s.handle), "cuMemcpyDtoDAsync"); + for (0..steps) |step| { + if (step > 0) try x.nextIds(step - 1, rows); + try x.k.embed.affine(x.s, s.ids, x.m.w.embed_w, x.m.w.embed_s, x.m.w.embed_b, s.e, rows, hid, .{ .bits = @intCast(x.m.w.embed_bits), .group = cfg.group, .scale = fwd.embedScale(hid) }); + try x.copy2d(s.cat, 4 * hid, s.e, 2 * hid, 2 * hid, rows); + try x.copy2d(s.cat + 2 * hid, 4 * hid, s.hin, 2 * hid, 2 * hid, rows); + try x.k.act.groupSums(x.s, s.cat, 2 * hid, s.xs, rows, 2 * hid); + try x.project(x.w.pre, s.cat, s.xs, s.h[0], rows); + try x.k.norms.addRmsTail(x.s, s.h[0], null, x.w.layers[0].input_norm, s.hn, s.normed, s.xs, rows, x.c.hidden, cfg.eps, .{ .round = true }); + var cur: usize = 0; + for (0..x.c.layers) |i| { + try x.layer(i, cur, rows); + cur = 1 - cur; + } + try x.project(x.w.head, s.normed, s.xs, s.logits, rows); + try x.k.round.argmax(x.s, s.logits, s.picks + step * max_items * @sizeOf(Pick), rows, cfg.vocab, cfg.vocab); + if (step + 1 < steps) try x.project(x.w.post, s.normed, s.xs, s.hin, rows); + } + } + + /// Step `step`'s picks as the next step's ids, read on the host first: a row with no finite logit stops the hold. + fn nextIds(x: *Drafter, step: usize, rows: usize) !void { + var picks: [max_items]Pick = undefined; + try x.s.synchronize(); + try x.d.check(x.d.api.cuMemcpyDtoH_v2(&picks, x.x.picks + step * max_items * @sizeOf(Pick), rows * @sizeOf(Pick)), "cuMemcpyDtoH"); + for (picks[0..rows]) |p| if (p.token >= cfg.vocab) return error.DraftNonfinite; + try x.copy2d(x.x.ids, 4, x.x.picks + step * max_items * @sizeOf(Pick), @sizeOf(Pick), 4, rows); + } + + /// Each stream's chain of `depth` drafts after `pending`, all streams' rows through each step together. + fn holdFn(ptr: *anyopaque, holds: []const dr.Hold) anyerror!void { + const x = cast(ptr); + if (holds.len > max_items) return error.DraftTooManyStreams; + var items: [max_items]Item = undefined; + var steps: u32 = 0; + for (holds, items[0..holds.len]) |h, *it| { + const l = x.lanes.get(h.stream) orelse return error.DraftStreamUnknown; + const t = x.gb.slotOf(h.stream) orelse return error.DraftStreamUnknown; + if (h.depth == 0 or h.depth > depth) return error.DraftDepth; + if (h.lanes != null) return error.DraftTreesUnsupported; + if (!l.rows or l.position + 2 != h.position or t.len < h.position - 1) return error.DraftContextGap; + it.* = .{ .lane = l, .slot = t.slot, .pending = h.pending, .position = l.position, .depth = h.depth }; + steps = @max(steps, h.depth); + } + try x.chain(items[0..holds.len], steps); + var picks: [depth * max_items]Pick = undefined; + try x.s.synchronize(); + try x.d.check(x.d.api.cuMemcpyDtoH_v2(&picks, x.x.picks, steps * max_items * @sizeOf(Pick)), "cuMemcpyDtoH"); + for (items[0..holds.len], 0..) |it, r| { + for (it.lane.held[0..it.depth], 0..) |*tok, step| { + const p = picks[step * max_items + r]; + if (p.nonfinite != 0 or p.token >= cfg.vocab) return error.DraftNonfinite; + tok.* = p.token; + } + it.lane.held_len = it.depth; + } + } + + /// Chains of 1 and 6 drafts over slot 0's caches (no stream holds them yet), timed: one draft's cost. + fn warm(x: *Drafter) !void { + const l = try x.newLane(); + defer x.dropLane(l); + var ms: [2]f64 = undefined; + for ([_]u32{ depth, 1, depth, 1, depth }, 0..) |steps, i| { + const t0 = std.Io.Clock.awake.now(x.io); + try x.chain(&.{.{ .lane = l, .slot = 0, .pending = 2, .position = 128, .depth = steps }}, steps); + try x.s.synchronize(); + if (i >= 3) ms[i - 3] = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(x.io)).toNanoseconds())) / 1e6; + } + x.step_ms = @max(ms[1] - ms[0], 0) / @as(f64, depth - 1); + // the target's windows of 1 to 16 rows on slot 0 (no stream holds it yet), the best of three each + const tokens: [be.max_window]u32 = @splat(2); + for (&x.costs, 1..) |*c, n| { + var best: f64 = std.math.inf(f64); + for (0..4) |i| { + const t0 = std.Io.Clock.awake.now(x.io); + try x.gb.run(&.{.{ .slot = 0, .p0 = 128, .tokens = tokens[0..n] }}, false, 0, n); + try x.s.synchronize(); + const t = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(x.io)).toNanoseconds())) / 1e6; + if (i > 0) best = @min(best, t); + } + c.* = .{ .width = @intCast(n), .ms = best }; + } + } +}; diff --git a/zig/src/families/gemma4/cuda_assistant_weights.zig b/zig/src/families/gemma4/cuda_assistant_weights.zig new file mode 100644 index 0000000000..f8fe628fe9 --- /dev/null +++ b/zig/src/families/gemma4/cuda_assistant_weights.zig @@ -0,0 +1,294 @@ +//! Gemma 4's MTP assistant weights on CUDA: config admitted as Metal's, bf16 linears and the tied head made 4-bit. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const cfg = @import("cuda_config.zig"); +const kern = @import("cuda_kernels.zig"); + +const Affine4 = cuda.qlinear.Affine4; + +/// Layers an admitted assistant has at most. +pub const max_layers = 8; +/// The stock assistant's sizes (google/gemma-4-12B-it-assistant); another config is refused with its reason. +const stock = .{ .hidden = 1024, .layers = 4, .intermediate = 8192, .heads = 16 }; + +pub const Config = struct { + layers: usize, + hidden: usize, + intermediate: usize, + heads: usize, + /// Each layer's kind: true for a sliding layer (the target's last sliding cache), false for a full one. + sliding: [max_layers]bool = @splat(true), + + pub fn headDim(c: *const Config, i: usize) usize { + return if (c.sliding[i]) cfg.sliding_dim else cfg.global_dim; + } + + /// Admission as Metal's assistant_weights.Config.parse, refusal for refusal. + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8, geo: cfg.Geometry) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const o = if (parsed.value == .object) parsed.value.object else return error.BadDraftConfig; + if (!isString(o.get("model_type"), "gemma4_unified_assistant")) return error.UnsupportedAssistantType; + if (isTrue(o, "use_ordered_embeddings")) return error.AssistantOrderedEmbeddingsNotServed; + if (try integer(o, "backbone_hidden_size") != geo.hidden) return error.AssistantForAnotherTarget; + if (o.get("quantization") != null or o.get("quantization_config") != null) return error.UnsupportedAssistantQuantization; + const t = if (o.get("text_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + if (!isString(t.get("model_type"), "gemma4_unified_text")) return error.BadDraftConfig; + var c: Config = .{ .layers = @intCast(try integer(t, "num_hidden_layers")), .hidden = @intCast(try integer(t, "hidden_size")), .intermediate = @intCast(try integer(t, "intermediate_size")), .heads = @intCast(try integer(t, "num_attention_heads")) }; + if (c.hidden != stock.hidden or c.layers != stock.layers or c.intermediate != stock.intermediate or c.heads != stock.heads) return error.UnsupportedAssistantGeometry; + if (try integer(t, "num_key_value_heads") != geo.sliding_kv_heads or try integer(t, "num_global_key_value_heads") != geo.global_kv_heads) return error.AssistantForAnotherTarget; + if (try integer(t, "num_kv_shared_layers") != c.layers) return error.UnsupportedAssistantGeometry; + const dims = .{ .{ "head_dim", cfg.sliding_dim }, .{ "global_head_dim", cfg.global_dim }, .{ "sliding_window", cfg.window }, .{ "vocab_size", cfg.vocab } }; + inline for (dims) |d| if (try integer(t, d[0]) != d[1]) return error.UnsupportedAssistantGeometry; + if (isTrue(t, "enable_moe_block") or isTrue(t, "use_double_wide_mlp") or isTrue(t, "attention_bias")) return error.UnsupportedAssistantGeometry; + if (t.get("hidden_size_per_layer_input")) |v| if (v != .integer or v.integer != 0) return error.UnsupportedAssistantGeometry; + if (!isTrue(t, "attention_k_eq_v") or !isTrue(t, "tie_word_embeddings")) return error.UnsupportedAssistantGeometry; + if (!isString(t.get("hidden_activation"), "gelu_pytorch_tanh")) return error.UnsupportedAssistantGeometry; + if (t.get("final_logit_softcapping")) |v| if (v != .null) return error.UnsupportedAssistantGeometry; + if (try number(t, "rms_norm_eps") != 1e-6) return error.UnsupportedAssistantGeometry; + const ropes = if (t.get("rope_parameters")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + try rope(ropes, "sliding_attention", "default", cfg.sliding_theta, null); + try rope(ropes, "full_attention", "proportional", cfg.global_theta, 0.25); + const kinds = t.get("layer_types") orelse return error.BadDraftConfig; + if (kinds != .array or kinds.array.items.len != c.layers) return error.BadDraftConfig; + for (kinds.array.items, 0..) |k, i| { + c.sliding[i] = isString(k, "sliding_attention"); + if (!c.sliding[i] and !isString(k, "full_attention")) return error.UnsupportedAssistantGeometry; + } + return c; + } +}; + +fn isString(v: ?std.json.Value, expected: []const u8) bool { + const s = v orelse return false; + return s == .string and std.mem.eql(u8, s.string, expected); +} + +fn isTrue(o: std.json.ObjectMap, key: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) !i64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return if (v == .integer) v.integer else error.BadDraftConfig; +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +fn rope(o: std.json.ObjectMap, kind: []const u8, kind_type: []const u8, theta: f64, partial: ?f64) !void { + const v = o.get(kind) orelse return error.BadDraftConfig; + const r = if (v == .object) v.object else return error.BadDraftConfig; + if (!isString(r.get("rope_type"), kind_type) or try number(r, "rope_theta") != theta) return error.UnsupportedAssistantRotary; + if (partial) |p| if (try number(r, "partial_rotary_factor") != p) return error.UnsupportedAssistantRotary; +} + +/// Whether `dir` holds a Gemma 4 assistant (its config's model_type names one). +pub fn isAssistant(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !bool { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const t = if (parsed.value == .object) parsed.value.object.get("model_type") else null; + return t != null and t.? == .string and std.mem.startsWith(u8, t.?.string, "gemma4") and std.mem.endsWith(u8, t.?.string, "assistant"); +} + +/// A group of 64 bf16 values as MLX quantizes it to 4 bits (Metal's dflash_weights.quantizeGroup). +fn quantizeGroup(x: []const u16, words: []u32, scale: *u16, bias: *u16) void { + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (x) |v| { + lo = @min(lo, bf(v)); + hi = @max(hi, bf(v)); + } + const bins: f32 = 15; + const scale0 = @max(hi - lo, 1e-7) / bins; + const from_min = @abs(lo) > @abs(hi); + const edge = if (from_min) lo else hi; + const q0 = @round(@abs(edge) / scale0); + const s_f: f32 = if (q0 == 0) scale0 else (if (from_min) @abs(edge) else -@abs(edge)) / q0; + scale.* = toBf(s_f); + bias.* = toBf(if (q0 == 0) lo else edge); + const s = bf(scale.*); + const b = bf(bias.*); + @memset(words, 0); + for (x, 0..) |v, i| { + const q = std.math.clamp(@round((bf(v) - b) / s), 0, bins); + words[i / 8] |= @as(u32, @intFromFloat(q)) << @intCast((i % 8) * 4); + } +} + +fn bf(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// Round to nearest even, as MLX casts. +fn toBf(f: f32) u16 { + const b: u32 = @bitCast(f); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +const Job = struct { + src: []const u16, + k: usize, + words: []u32, + scales: []u16, + biases: []u16, + fn run(j: Job, from: usize, to: usize) void { + const groups = j.k / 64; + for (from..to) |r| for (0..groups) |g| { + quantizeGroup(j.src[r * j.k + g * 64 ..][0..64], j.words[(r * groups + g) * 8 ..][0..8], &j.scales[r * groups + g], &j.biases[r * groups + g]); + }; + } +}; + +pub const Layer = struct { + q: Affine4, + o: Affine4, + gate: Affine4, + up: Affine4, + down: Affine4, + q_norm: u64, + input_norm: u64, + post_attention: u64, + pre_feedforward: u64, + post_feedforward: u64, + scalar: u64, +}; + +pub const Weights = struct { + pre: Affine4, // [embedding | hidden] to the assistant's width + post: Affine4, // its output back to the target's + head: Affine4, // the tied head: embed_tokens as a 4-bit linear over the vocabulary + norm: u64, + layers: [max_layers]Layer, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(gpa); + } +}; + +const Loader = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + s: cuda.Stream, + k: *const kern.Kernels, + ck: *core.checkpoint.Checkpoint, + w: *Weights, + + fn keep(x: *Loader, b: cuda.DeviceBuffer) !u64 { + errdefer { + var y = b; + y.free(); + } + try x.w.buffers.append(x.gpa, b); + return b.ptr; + } + + fn vector(x: *Loader, name: []const u8, len: usize) !u64 { + return x.keep(try cuda.DeviceBuffer.fromHost(x.d, (try x.ck.expect(name, .bf16, &.{len})).bytes)); + } + + /// An n-by-k bf16 matrix quantized to MLX affine 4/64 on every core, packed into qlinear's tiles. + fn affine4(x: *Loader, name: []const u8, n: usize, k: usize) !Affine4 { + const t = try x.ck.expect(name, .bf16, &.{ n, k }); + const src = std.mem.bytesAsSlice(u16, @as([]align(1) const u8, t.bytes)); + const host = try x.gpa.alloc(u16, n * k); + defer x.gpa.free(host); + for (host, src) |*h, v| h.* = v; + const words = try x.gpa.alloc(u32, n * k / 8); + defer x.gpa.free(words); + const sb = try x.gpa.alloc(u16, 2 * n * k / 64); + defer x.gpa.free(sb); + const job: Job = .{ .src = host, .k = k, .words = words, .scales = sb[0 .. n * k / 64], .biases = sb[n * k / 64 ..] }; + const threads = @min(std.Thread.getCpuCount() catch 1, 16); + const per = (n + threads - 1) / threads; + var handles: [16]std.Thread = undefined; + var spawned: usize = 0; + var at: usize = 0; + while (at < n) : (at += per) { + handles[spawned] = try std.Thread.spawn(.{}, Job.run, .{ job, at, @min(n, at + per) }); + spawned += 1; + } + for (handles[0..spawned]) |h| h.join(); + var raw_w = try cuda.DeviceBuffer.fromHost(x.d, std.mem.sliceAsBytes(words)); + defer raw_w.free(); + var raw_sb = try cuda.DeviceBuffer.fromHost(x.d, std.mem.sliceAsBytes(sb)); + defer raw_sb.free(); + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try x.keep(try cuda.DeviceBuffer.alloc(x.d, l.words)), .s = try x.keep(try cuda.DeviceBuffer.alloc(x.d, l.scales)), .b = try x.keep(try cuda.DeviceBuffer.alloc(x.d, l.scales)), .n = n, .k = k, .npad = l.npad }; + try x.k.affine.pack(x.s, raw_w.ptr, raw_sb.ptr, raw_sb.ptr + n * k / 64 * 2, q); + try x.s.synchronize(); // the raw buffers are freed on return + return q; + } +}; + +/// Every tensor checked by name, dtype and shape; the bf16 linears and the tied head made 4-bit as Metal's. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, k: *const kern.Kernels, dir: []const u8, c: *const Config, geo: cfg.Geometry) !Weights { + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var w: Weights = .{ .pre = undefined, .post = undefined, .head = undefined, .norm = 0, .layers = undefined }; + errdefer w.deinit(gpa); + var x: Loader = .{ .gpa = gpa, .d = d, .s = s, .k = k, .ck = &ck, .w = &w }; + const hid = c.hidden; + w.pre = try x.affine4("pre_projection.weight", hid, 2 * geo.hidden); + w.post = try x.affine4("post_projection.weight", geo.hidden, hid); + w.head = try x.affine4("model.embed_tokens.weight", cfg.vocab, hid); + w.norm = try x.vector("model.norm.weight", hid); + var buf: [160]u8 = undefined; + for (w.layers[0..c.layers], 0..) |*l, i| { + const q = c.heads * c.headDim(i); + const name = struct { + fn of(b: []u8, layer: usize, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "model.layers.{d}.{s}", .{ layer, suffix }); + } + }.of; + l.q = try x.affine4(try name(&buf, i, "self_attn.q_proj.weight"), q, hid); + l.o = try x.affine4(try name(&buf, i, "self_attn.o_proj.weight"), hid, q); + l.gate = try x.affine4(try name(&buf, i, "mlp.gate_proj.weight"), c.intermediate, hid); + l.up = try x.affine4(try name(&buf, i, "mlp.up_proj.weight"), c.intermediate, hid); + l.down = try x.affine4(try name(&buf, i, "mlp.down_proj.weight"), hid, c.intermediate); + l.q_norm = try x.vector(try name(&buf, i, "self_attn.q_norm.weight"), c.headDim(i)); + l.input_norm = try x.vector(try name(&buf, i, "input_layernorm.weight"), hid); + l.post_attention = try x.vector(try name(&buf, i, "post_attention_layernorm.weight"), hid); + l.pre_feedforward = try x.vector(try name(&buf, i, "pre_feedforward_layernorm.weight"), hid); + l.post_feedforward = try x.vector(try name(&buf, i, "post_feedforward_layernorm.weight"), hid); + l.scalar = try x.vector(try name(&buf, i, "layer_scalar"), 1); + } + return w; +} + +const fixture = @embedFile("fixtures/assistant_config.json"); + +test "the 12B's assistant is admitted against the 12B, and refused against the 31B" { + const c = try Config.parse(std.testing.allocator, fixture, cfg.unified12b); + try std.testing.expectEqual(@as(usize, 4), c.layers); + try std.testing.expect(c.sliding[0] and c.sliding[2] and !c.sliding[3]); + try std.testing.expectEqual(@as(usize, 512), c.headDim(3)); + try std.testing.expectError(error.AssistantForAnotherTarget, Config.parse(std.testing.allocator, fixture, cfg.dense31b)); +} + +test "a group quantizes as MLX's affine 4-bit: every value within half a step" { + var x: [64]u16 = undefined; + for (&x, 0..) |*v, i| v.* = toBf(@as(f32, @floatFromInt(i)) * 0.03 - 0.7); + var words: [8]u32 = undefined; + var scale: u16 = 0; + var bias: u16 = 0; + quantizeGroup(&x, &words, &scale, &bias); + for (x, 0..) |v, i| { + const q: f32 = @floatFromInt((words[i / 8] >> @intCast((i % 8) * 4)) & 0xF); + try std.testing.expect(@abs(q * bf(scale) + bf(bias) - bf(v)) <= @abs(bf(scale)) * 0.51); + } +} diff --git a/zig/src/families/gemma4/cuda_backend.zig b/zig/src/families/gemma4/cuda_backend.zig new file mode 100644 index 0000000000..0c1bbeb01a --- /dev/null +++ b/zig/src/families/gemma4/cuda_backend.zig @@ -0,0 +1,397 @@ +//! Gemma 4 on CUDA behind lanes.Backend: a slot a stream, a chunk a prompt step, one forward a round of windows. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("cuda_config.zig"); +const fwd = @import("cuda_forward.zig"); +const kern = @import("cuda_kernels.zig"); +const snap = @import("cuda_snap.zig"); + +const be = lanes.backend; +const Stream = lanes.Stream; +const Pick = cuda.tree_round.Pick; + +/// Prompt rows a step takes: the forward's chunk grid, from position 0. +pub const prompt_rows = fwd.prompt_rows; +/// Window rows a round takes over every stream. +pub const round_rows = fwd.round_rows; +/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. +pub const max_window = 16; +/// Drawn tokens a handle names (a ring; read before the round loop could wrap it). +const ring = 1024; + +const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's soft-capped bf16 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + tap: Tap = .{}, +}; + +/// A stream's rows in the last forward: rows [row, row + rows) at positions [start, start + rows) while gen holds. +const Tap = struct { row: usize = 0, start: u64 = 0, rows: usize = 0, gen: u64 = 0 }; + +pub const Gemma = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + m: *fwd.Model, + k: *const kern.Kernels, + s: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + tok: cuda.DeviceBuffer, // drawn tokens (u32): the handle ring, then a round's rows + picks: cuda.DeviceBuffer, // argmax picks (Pick), the same entries + greedy: []bool, // entry e of the ring (and the round's rows) is a pick, not a draw + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot + next: u64 = 0, + gen: u64 = 1, // forwards run: a tap is current while its gen is this + tapped: bool = false, // a drafter reads the last layer's rows (prepare_features) + + /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. + pub fn openBytes(slots: usize) usize { + return (ring + round_rows) * (4 + @sizeOf(Pick)) + round_rows * 16 + slots * (cfg.vocab * 2 + @sizeOf(cuda.sample.Rule)); + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *fwd.Model, k: *const kern.Kernels, s: cuda.Stream) !Gemma { + const free = try gpa.alloc(bool, m.slots); + errdefer gpa.free(free); + @memset(free, true); + const greedy = try gpa.alloc(bool, ring + round_rows); + errdefer gpa.free(greedy); + @memset(greedy, true); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var picks = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * @sizeOf(Pick)); + errdefer picks.free(); + var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + errdefer meta.free(); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slots); + errdefer gpa.free(bufs); + var made: usize = 0; + errdefer for (bufs[0..made]) |*b| b.free(); + for (bufs, 0..) |*b, i| { + b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) cfg.vocab * 2 else @sizeOf(cuda.sample.Rule)); + made += 1; + } + return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta, .slot_bufs = bufs }; + } + + pub fn deinit(x: *Gemma) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| l.rows.deinit(x.gpa); + x.lanes.deinit(x.gpa); + for (x.slot_bufs) |*b| b.free(); + x.gpa.free(x.slot_bufs); + x.gpa.free(x.free); + x.gpa.free(x.greedy); + x.tok.free(); + x.picks.free(); + x.meta.free(); + } + + pub fn backend(x: *Gemma) be.Backend { + return .{ .ptr = x, .vtable = &vtable }; + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true, .prepare_features = prepareFeatures, .features = features, .prefill_features = prefillFeatures }; + + /// The drafter's taps: the last layer's output (before the final norm), the only one a forward keeps. + fn prepareFeatures(ptr: *anyopaque, taps: []const u32) anyerror!void { + const x = self(ptr); + if (taps.len != 1 or taps[0] + 1 != x.m.layers) return error.GemmaTapNotServed; + x.tapped = true; + } + + /// Rows [start, start + count) of the stream's last forward, on the device until the next forward. + fn features(ptr: *anyopaque, s: *Stream, taps: []const u32, start: u64, count: u32) anyerror!be.Features { + const x = self(ptr); + if (!x.tapped or taps.len != 1 or taps[0] + 1 != x.m.layers) return error.GemmaTapNotServed; + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + return x.featuresOf(l.tap, start, count); + } + + fn featuresOf(x: *const Gemma, t: Tap, start: u64, count: u32) !be.Features { + if (t.gen != x.gen or start < t.start or start + count > t.start + t.rows or count == 0) return error.RowsNotHeld; + const width = x.m.geo.hidden * 2; + return .{ .buffer = x.m.s.h[0], .offset = (t.row + start - t.start) * width, .rows = count, .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 }; + } + + /// The prompt pass, `sink` taking each chunk's rows as the chunk is made. + fn prefillFeatures(ptr: *anyopaque, s: *Stream, taps: []const u32, sink: be.FeatureSink) anyerror!void { + const x = self(ptr); + if (!x.tapped or taps.len != 1 or taps[0] + 1 != x.m.layers) return error.GemmaTapNotServed; + const l = try x.opening(s); + errdefer x.dropLane(s); + while (true) { + const at = l.fill.?; + const done = try x.chunk(s, l); + try sink.take(sink.ptr, s, try x.featuresOf(l.tap, at, @intCast(l.tap.rows)), at); + if (done) return; + } + } + + /// One forward of the model, its rows' taps current until the next. + pub fn run(x: *Gemma, chains: []const fwd.Chain, prompt: bool, head0: usize, heads: usize) !void { + x.gen += 1; + try x.m.forward(x.k, x.s, chains, prompt, head0, heads); + } + + /// The stream's cache slot and length, for a drafter reading its caches. + pub fn slotOf(x: *const Gemma, s: *const Stream) ?struct { slot: usize, len: u64 } { + const l = x.lanes.getPtr(s) orelse return null; + return .{ .slot = l.slot, .len = l.len }; + } + + /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a forward. + pub fn facts(x: *const Gemma) lanes.Model { + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slots, round_rows)) }; + } + + fn self(ptr: *anyopaque) *Gemma { + return @ptrCast(@alignCast(ptr)); + } + + fn upload(x: *Gemma, dst: u64, bytes: []const u8) !void { + try x.d.check(x.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, x.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Rows [row0, row0 + rows) of bf16 `logits` into entries [at, at + rows), row r at `position0 + r`. + fn drawRows(x: *Gemma, l: *const Lane, logits: u64, row0: usize, rows: usize, position0: u64, at: usize, meta_at: usize) !void { + const v = cfg.vocab; + const src = logits + row0 * v * 2; + @memset(x.greedy[at..][0..rows], !l.sampled); + if (!l.sampled) return x.k.round.argmax(x.s, src, x.picks.ptr + at * @sizeOf(Pick), rows, v, v); + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try x.upload(mp, std.mem.asBytes(&meta)); + try x.k.draws.draw(x.s, src, v, l.rule, mp, 0, x.tok.ptr + at * 4, rows, null, null); + } + + /// Entries [at, at + out.len) as tokens, after the stream's work; a pick over non-finite logits is an error. + fn tokens(x: *Gemma, at: usize, out: []u32) !void { + try x.s.synchronize(); + var picks: [round_rows]Pick = undefined; + try x.tok.download(at * 4, std.mem.sliceAsBytes(out)); + try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..out.len])); + for (out, picks[0..out.len], x.greedy[at..][0..out.len]) |*t, p, g| if (g) { + if (p.nonfinite != 0) return error.NonFiniteLogits; + t.* = p.token; + }; + } + + pub fn saveState(x: *Gemma, owner: ?*anyopaque, at: u32) !*snap.Saved { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if ((l.fill orelse l.len) != at) return error.NotAtMark; // a state is copied only where the pass stands + return snap.save(x.gpa, x.m, x.s, l.slot, at); + } + + pub fn restoreState(x: *Gemma, owner: ?*anyopaque, k: *const snap.Saved) !void { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + try snap.restore(x.m, x.s, l.slot, k); + l.len = k.at; + } + + pub fn dropState(x: *Gemma, k: *snap.Saved) void { + snap.drop(x.gpa, x.s, k); + } + + fn freeSlot(x: *const Gemma) ?usize { + return std.mem.indexOfScalar(bool, x.free, true); + } + + /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. + fn opening(x: *Gemma, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (p.len + @as(usize, s.max_new) + max_window > x.m.context) return error.PastContext; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + const slot = x.freeSlot() orelse { + x.lanes.removeByPtr(got.key_ptr); + return error.NoFreeSlot; + }; + x.free[slot] = false; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr }; + } + const l = got.value_ptr; + l.sampled = false; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try x.upload(l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + l.fill = x.resumeAt(s, l, p.len); + s.fill_at.store(@intCast(l.fill.?), .monotonic); + return l; + } + + /// The kept state's position, restored into the slot (0: none, or not this prompt's). + fn resumeAt(x: *Gemma, s: *Stream, l: *const Lane, len: usize) usize { + const saved = s.reuse.saved orelse return 0; + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len) { + s.reuse_failed = true; + return 0; + } + snap.restore(x.m, x.s, l.slot, k) catch { + s.reuse_failed = true; + return 0; + }; + s.cached = k.at; + return k.at; + } + + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written. + fn chunk(x: *Gemma, s: *Stream, l: *Lane) !bool { + const p = s.prompt(); + const at = l.fill.?; + if (s.isCancelled()) return error.Cancelled; + var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows); + for (s.reuse.marks) |mk| if (mk > at and mk < end) { + end = mk; + }; + const done = end == p.len; + try x.run(&.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, end - at - 1, @intFromBool(done)); + l.tap = .{ .row = 0, .start = at, .rows = end - at, .gen = x.gen }; + if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, cfg.vocab * 2, x.s.handle), "cuMemcpyDtoDAsync"); + s.fill_at.store(@intCast(end), .monotonic); + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l); + } + + fn handle(x: *Gemma) usize { + const h = x.next; + x.next += 1; + return @intCast(h % ring); + } + + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.next; + try x.drawRows(l, l.last, 0, 1, position, x.handle(), 0); + return h; + } + + fn value(x: *Gemma, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| read(x, h), + }; + } + + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const t = try x.value(feed); + if (position != l.len + 1) return error.PositionMismatch; + try x.run(&.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{t} }}, false, 0, 1); + l.tap = .{ .row = 0, .start = l.len, .rows = 1, .gen = x.gen }; + l.len += 1; + const h = x.next; + try x.drawRows(l, x.m.s.logits, 0, 1, position, x.handle(), 0); + return h; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + const x = self(ptr); + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.tokens(@intCast(h % ring), &out); + return out[0]; + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + const chains = try x.gpa.alloc(fwd.Chain, windows.len); + defer x.gpa.free(chains); + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (w.rows() > max_window) return error.WindowTooWide; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items }; + total += l.rows.items.len; + } + if (total > round_rows) return error.RoundTooWide; + try x.run(chains, false, 0, total); + var at: usize = 0; + for (windows, 0..) |w, i| { + const l = x.lanes.getPtr(w.stream).?; + l.tap = .{ .row = at, .start = l.base, .rows = l.rows.items.len, .gen = x.gen }; + try x.drawRows(l, x.m.s.logits, at, l.rows.items.len, w.positions[0], ring + at, i); + at += l.rows.items.len; + } + var got: [round_rows]u32 = undefined; + try x.tokens(ring, got[0..total]); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + l.len = l.base + n; // a chain keeps every row unless rolled back + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + for (windows, paths) |w, path| { + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote + } + } + + /// No draft head is loaded (#1048 steps 10/11): absorbing is a no-op, a request for drafts is refused by name. + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + fn dropLane(x: *Gemma, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + } + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } +}; diff --git a/zig/src/families/gemma4/cuda_config.zig b/zig/src/families/gemma4/cuda_config.zig new file mode 100644 index 0000000000..e61d48916a --- /dev/null +++ b/zig/src/families/gemma4/cuda_config.zig @@ -0,0 +1,397 @@ +//! config.zig for CUDA builds, rule for rule (cuda_widths.zig in place of widths.zig): Gemma 4 dense admission. +const std = @import("std"); +pub const widths = @import("cuda_widths.zig"); + +pub const vocab = 262144; +pub const sliding_dim = 256; +pub const global_dim = 512; +pub const window = 1024; +pub const group = widths.group; +pub const eps: f32 = 1e-6; +pub const softcap: f32 = 30.0; +pub const sliding_theta: f64 = 10000.0; +pub const global_theta: f64 = 1000000.0; +/// Pairs the global layers' proportional RoPE rotates: partial_rotary_factor 0.25 of 512 dims, in pairs. +pub const global_rotated = 64; + +/// A model's sizes; everything else (head dims, window, vocab, RoPE) is the same in each admitted one. +pub const Geometry = struct { + name: []const u8, + hidden: usize, + layers: usize, + intermediate: usize, + query_heads: usize, + sliding_kv_heads: usize, + global_kv_heads: usize, + + /// The 12B's unified checkpoint: its image input is the patch embedder, not the 31B's tower. + pub fn unified(g: Geometry) bool { + return std.mem.eql(u8, g.name, unified12b.name); + } + + /// Layer `index`'s key heads. + pub fn kvHeads(g: Geometry, index: usize) usize { + return if (global(index)) g.global_kv_heads else g.sliding_kv_heads; + } + + /// The widest q a layer writes: every query head at the global head dim. + pub fn queryWidth(g: Geometry) usize { + return g.query_heads * global_dim; + } + + /// The widest k (or v) a layer writes, of a sliding and a global layer. + pub fn kvWidth(g: Geometry) usize { + return @max(g.sliding_kv_heads * sliding_dim, g.global_kv_heads * global_dim); + } +}; + +/// Gemma 4 31B (model_type gemma4). +pub const dense31b: Geometry = .{ .name = "31B", .hidden = 5376, .layers = 60, .intermediate = 21504, .query_heads = 32, .sliding_kv_heads = 16, .global_kv_heads = 4 }; +/// Gemma 4 12B's text model (model_type gemma4_unified, text_config gemma4_unified_text). +pub const unified12b: Geometry = .{ .name = "12B", .hidden = 3840, .layers = 48, .intermediate = 15360, .query_heads = 16, .sliding_kv_heads = 8, .global_kv_heads = 1 }; +/// Layers an admitted geometry has at most: the size of per-layer arrays. +pub const max_layers = 60; + +comptime { + for ([_]Geometry{ dense31b, unified12b }) |g| std.debug.assert(g.layers <= max_layers); +} + +pub const Error = error{ + BadGemmaConfig, + BadGemmaContext, + BadGemmaEos, + GemmaMoeNotServed, + GemmaUnifiedNotServed, + GemmaPerLayerInputsNotServed, + GemmaSharedKvNotServed, + GemmaKeysNotValues, + UnsupportedGemmaGeometry, + UnsupportedGemmaRotary, + UnsupportedGemmaQuantization, +} || widths.Error || std.json.ParseError(std.json.Scanner); + +pub const Config = struct { + geo: Geometry, + /// Each linear's MLX affine width (the config's quantization table). + widths: widths.Widths = .{}, + context: usize, + eos: [4]u32 = .{ 0, 0, 0, 0 }, + eos_count: usize = 0, + + 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 bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } +}; + +/// Whether layer `index` is a global (full attention) layer: every sixth, and the last. +pub fn global(index: usize) bool { + return index % 6 == 5; +} + +/// A layer's head dim and key heads. +pub fn headDim(index: usize) usize { + return if (global(index)) global_dim else sliding_dim; +} + +/// RoPE inverse frequencies for a layer kind, as float32 [dim / 2]: zero for pairs proportional RoPE leaves alone. +pub fn inverseFrequencies(is_global: bool, out: []f32) void { + const dim: f64 = if (is_global) global_dim else sliding_dim; + const theta = if (is_global) global_theta else sliding_theta; + for (out, 0..) |*f, i| { + const rotated = !is_global or i < global_rotated; + f.* = if (rotated) @floatCast(1.0 / std.math.pow(f64, theta, 2.0 * @as(f64, @floatFromInt(i)) / dim)) else 0.0; + } +} + +fn object(v: std.json.Value) Error!std.json.ObjectMap { + return if (v == .object) v.object else error.BadGemmaConfig; +} + +fn field(o: std.json.ObjectMap, key: []const u8) Error!std.json.Value { + return o.get(key) orelse error.BadGemmaConfig; +} + +fn number(v: std.json.Value) Error!f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadGemmaConfig, + }; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) Error!i64 { + const v = try field(o, key); + return if (v == .integer) v.integer else error.BadGemmaConfig; +} + +fn optionalInt(o: std.json.ObjectMap, key: []const u8) i64 { + const v = o.get(key) orelse return 0; + return if (v == .integer) v.integer else 0; +} + +fn isString(v: ?std.json.Value, expected: []const u8) bool { + const s = v orelse return false; + return s == .string and std.mem.eql(u8, s.string, expected); +} + +fn isTrue(o: std.json.ObjectMap, key: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool; +} + +/// The widths of `quantization`, and of `quantization_config` when present, which must name the same ones. +fn quantization(gpa: std.mem.Allocator, root: std.json.ObjectMap, layers: usize) Error!widths.Widths { + const w = try widths.read(gpa, try object(root.get("quantization") orelse return error.UnsupportedGemmaQuantization), layers); + if (root.get("quantization_config")) |q| { + if (!std.meta.eql(w, try widths.read(gpa, try object(q), layers))) return error.UnsupportedGemmaQuantization; + } + return w; +} + +fn rope(o: std.json.ObjectMap, kind: []const u8, kind_type: []const u8, theta: f64, partial: ?f64) Error!void { + const r = try object(try field(o, kind)); + if (!isString(r.get("rope_type"), kind_type)) return error.UnsupportedGemmaRotary; + if (try number(try field(r, "rope_theta")) != theta) return error.UnsupportedGemmaRotary; + if (partial) |p| if (try number(try field(r, "partial_rotary_factor")) != p) return error.UnsupportedGemmaRotary; +} + +fn eosIds(v: std.json.Value, c: *Config) Error!void { + switch (v) { + .integer => |i| { + if (i < 0 or i >= vocab) return error.BadGemmaEos; + if (c.eos_count < c.eos.len) { + c.eos[c.eos_count] = @intCast(i); + c.eos_count += 1; + } + }, + .array => |a| for (a.items) |x| try eosIds(x, c), + else => return error.BadGemmaEos, + } +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) Error!Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = try object(parsed.value); + const unified = isString(root.get("model_type"), "gemma4_unified"); + if (!unified and !isString(root.get("model_type"), "gemma4")) return error.BadGemmaConfig; + const text = try object(try field(root, "text_config")); + if (unified and !isString(text.get("model_type"), "gemma4_unified_text")) return error.GemmaUnifiedNotServed; + const geo = if (unified) unified12b else dense31b; + if (isTrue(text, "enable_moe_block")) return error.GemmaMoeNotServed; + if (optionalInt(text, "hidden_size_per_layer_input") != 0) return error.GemmaPerLayerInputsNotServed; + if (optionalInt(text, "num_kv_shared_layers") != 0) return error.GemmaSharedKvNotServed; + if (!isTrue(text, "attention_k_eq_v")) return error.GemmaKeysNotValues; + if (isTrue(text, "use_double_wide_mlp") or isTrue(text, "attention_bias")) return error.UnsupportedGemmaGeometry; + if (!isString(text.get("hidden_activation"), "gelu_pytorch_tanh")) return error.UnsupportedGemmaGeometry; + if (!isTrue(text, "tie_word_embeddings")) return error.UnsupportedGemmaGeometry; + const dims = .{ + .{ "hidden_size", geo.hidden }, .{ "vocab_size", vocab }, .{ "num_hidden_layers", geo.layers }, + .{ "intermediate_size", geo.intermediate }, .{ "num_attention_heads", geo.query_heads }, .{ "num_key_value_heads", geo.sliding_kv_heads }, + .{ "head_dim", sliding_dim }, .{ "global_head_dim", global_dim }, .{ "num_global_key_value_heads", geo.global_kv_heads }, + .{ "sliding_window", window }, + }; + inline for (dims) |d| if (try integer(text, d[0]) != @as(i64, @intCast(d[1]))) return error.UnsupportedGemmaGeometry; + if (try number(try field(text, "rms_norm_eps")) != 1e-6) return error.UnsupportedGemmaGeometry; + if (try number(try field(text, "final_logit_softcapping")) != softcap) return error.UnsupportedGemmaGeometry; + const types = try field(text, "layer_types"); + if (types != .array or types.array.items.len != geo.layers) return error.UnsupportedGemmaGeometry; + for (types.array.items, 0..) |kind, i| { + if (!isString(kind, if (global(i)) "full_attention" else "sliding_attention")) return error.UnsupportedGemmaGeometry; + } + const ropes = try object(try field(text, "rope_parameters")); + try rope(ropes, "sliding_attention", "default", sliding_theta, null); + try rope(ropes, "full_attention", "proportional", global_theta, 0.25); + const ws = try quantization(gpa, root, geo.layers); + const context = try integer(text, "max_position_embeddings"); + if (context < 1 or context > 262144) return error.BadGemmaContext; + var c: Config = .{ .geo = geo, .widths = ws, .context = @intCast(context) }; + try eosIds(root.get("eos_token_id") orelse try field(text, "eos_token_id"), &c); + if (c.eos_count == 0) return error.BadGemmaEos; + return c; +} + +const fixture = @embedFile("fixtures/config.json"); +/// The 4-bit 12B's config as tools/zig/gemma4_12b_convert.py writes it: unified model types (the text model's part). +const fixture_12b = @embedFile("fixtures/config_12b.json"); + +fn boundaryFixture(base: []const u8, key: []const u8, value: []const u8, root_eos: bool) ![]u8 { + const a = std.testing.allocator; + var p = try std.json.parseFromSlice(std.json.Value, a, base, .{}); + defer p.deinit(); + const v = try std.json.parseFromSlice(std.json.Value, a, value, .{}); + defer v.deinit(); + if (root_eos) { + try p.value.object.put(p.arena.allocator(), key, v.value); + } else { + if (std.mem.eql(u8, key, "eos_token_id")) _ = p.value.object.swapRemove(key); + const text = p.value.object.getPtr("text_config").?; + try text.object.put(p.arena.allocator(), key, v.value); + } + return std.json.Stringify.valueAlloc(a, p.value, .{}); +} + +test "Gemma context boundaries refuse before unsigned conversion in both geometries" { + const a = std.testing.allocator; + for ([_][]const u8{ fixture, fixture_12b }) |base| { + for ([_][]const u8{ "-9223372036854775808", "-1", "0", "262145", "9223372036854775807" }) |value| { + const bytes = try boundaryFixture(base, "max_position_embeddings", value, false); + defer a.free(bytes); + try std.testing.expectError(error.BadGemmaContext, parse(a, bytes)); + } + for ([_][]const u8{ "1", "262144" }) |value| { + const bytes = try boundaryFixture(base, "max_position_embeddings", value, false); + defer a.free(bytes); + const c = try parse(a, bytes); + try std.testing.expectEqual(try std.fmt.parseInt(usize, value, 10), c.context); + } + } +} + +test "Gemma EOS scalar and array boundaries are checked at root and text fallback" { + const a = std.testing.allocator; + for ([_][]const u8{ fixture, fixture_12b }) |base| { + for ([_]bool{ false, true }) |root_eos| { + for ([_][]const u8{ + "-1", "262144", "4294967296", "9223372036854775807", + "[-1]", "[1,4294967296]", "[0,1,2,3,-1]", "[]", + "null", "\"1\"", + }) |value| { + const bytes = try boundaryFixture(base, "eos_token_id", value, root_eos); + defer a.free(bytes); + try std.testing.expectError(error.BadGemmaEos, parse(a, bytes)); + } + for ([_][]const u8{ "0", "262143", "[0,262143]" }) |value| { + const bytes = try boundaryFixture(base, "eos_token_id", value, root_eos); + defer a.free(bytes); + const c = try parse(a, bytes); + try std.testing.expect(c.eos_count > 0); + for (c.eos[0..c.eos_count]) |id| try std.testing.expect(id < vocab); + } + } + } +} + +fn edited(gpa: std.mem.Allocator, from: []const u8, to: []const u8) ![]u8 { + return std.mem.replaceOwned(u8, gpa, fixture, from, to); +} + +fn edited12b(gpa: std.mem.Allocator, edits: []const [2][]const u8) ![]u8 { + var bytes = try gpa.dupe(u8, fixture_12b); + for (edits) |e| { + const next = try std.mem.replaceOwned(u8, gpa, bytes, e[0], e[1]); + gpa.free(bytes); + bytes = next; + } + return bytes; +} + +test "the 31B 4-bit checkpoint's config is admitted, with its context and end tokens" { + const c = try parse(std.testing.allocator, fixture); + try std.testing.expectEqualStrings("31B", c.geo.name); + try std.testing.expectEqual(dense31b, c.geo); + try std.testing.expectEqual(@as(usize, 262144), c.context); + try std.testing.expectEqualSlices(u32, &.{ 1, 106, 50 }, c.eos[0..c.eos_count]); +} + +test "the 12B's unified text model is admitted with its own geometry" { + const c = try parse(std.testing.allocator, fixture_12b); + try std.testing.expectEqualStrings("12B", c.geo.name); + try std.testing.expectEqual(unified12b, c.geo); + try std.testing.expectEqual(@as(usize, 48), c.geo.layers); + try std.testing.expectEqual(@as(usize, 262144), c.context); + try std.testing.expectEqualSlices(u32, &.{ 1, 106, 50 }, c.eos[0..c.eos_count]); +} + +test "a model type with the other geometry's sizes is refused" { + const gpa = std.testing.allocator; + const dense12b = try edited12b(gpa, &.{.{ "\"model_type\": \"gemma4_unified\",", "\"model_type\": \"gemma4\"," }}); + defer gpa.free(dense12b); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, dense12b)); + const unified31b = try edited(gpa, "\"model_type\": \"gemma4\"", "\"model_type\": \"gemma4_unified\""); + defer gpa.free(unified31b); + const unified31b_text = try std.mem.replaceOwned(u8, gpa, unified31b, "\"gemma4_text\"", "\"gemma4_unified_text\""); + defer gpa.free(unified31b_text); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, unified31b_text)); + const sixty = try edited12b(gpa, &.{.{ "\"num_hidden_layers\": 48", "\"num_hidden_layers\": 60" }}); + defer gpa.free(sixty); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, sixty)); + const wide = try edited12b(gpa, &.{.{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 16" }}); + defer gpa.free(wide); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, wide)); +} + +test "each other unified layout is refused by name" { + const gpa = std.testing.allocator; + const cases = .{ + .{ "\"model_type\": \"gemma4_unified_text\"", "\"model_type\": \"gemma4_text\"", error.GemmaUnifiedNotServed }, + .{ "\"hidden_size_per_layer_input\": 0", "\"hidden_size_per_layer_input\": 256", error.GemmaPerLayerInputsNotServed }, + .{ "\"num_kv_shared_layers\": 0", "\"num_kv_shared_layers\": 18", error.GemmaSharedKvNotServed }, + .{ "\"enable_moe_block\": false", "\"enable_moe_block\": true", error.GemmaMoeNotServed }, + .{ "\"bits\": 4", "\"bits\": 7", error.UnsupportedGemmaQuantization }, + }; + inline for (cases) |c| { + const bytes = try edited12b(gpa, &.{.{ c[0], c[1] }}); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture_12b)); + try std.testing.expectError(c[2], parse(gpa, bytes)); + } +} + +test "each other Gemma 4 layout is refused by name" { + const gpa = std.testing.allocator; + const cases = .{ + .{ "\"enable_moe_block\": false", "\"enable_moe_block\": true", error.GemmaMoeNotServed }, + .{ "\"model_type\": \"gemma4\"", "\"model_type\": \"gemma4_unified\"", error.GemmaUnifiedNotServed }, + .{ "\"attention_k_eq_v\": true", "\"attention_k_eq_v\": false", error.GemmaKeysNotValues }, + .{ "\"num_kv_shared_layers\": 0", "\"num_kv_shared_layers\": 18", error.GemmaSharedKvNotServed }, + .{ "\"hidden_size_per_layer_input\": 0", "\"hidden_size_per_layer_input\": 256", error.GemmaPerLayerInputsNotServed }, + .{ "\"sliding_window\": 1024", "\"sliding_window\": 512", error.UnsupportedGemmaGeometry }, + }; + inline for (cases) |c| { + const bytes = try edited(gpa, c[0], c[1]); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture)); + try std.testing.expectError(c[2], parse(gpa, bytes)); + } +} + +test "other widths are admitted per module, and quantization_config must agree" { + const gpa = std.testing.allocator; + const eight = try edited12b(gpa, &.{.{ "\"bits\": 4", "\"bits\": 8" }}); // both tables + defer gpa.free(eight); + const c8 = try parse(gpa, eight); + try std.testing.expectEqual(@as(u8, 8), c8.widths.of(47, .down).bits); + try std.testing.expect(!c8.widths.allFast(c8.geo.layers)); + try std.testing.expect((try parse(gpa, fixture_12b)).widths.allFast(48)); + const entry = "\"group_size\": 64, \"language_model.model.layers.0.mlp.down_proj\": {\"bits\": 8, \"group_size\": 64},"; + const one = try edited(gpa, "\"quantization\": {\n \"group_size\": 64,", "\"quantization\": {" ++ entry); + defer gpa.free(one); + try std.testing.expect(!std.mem.eql(u8, one, fixture)); + try std.testing.expectError(error.UnsupportedGemmaQuantization, parse(gpa, one)); // quantization_config disagrees + const both = try edited(gpa, "\"group_size\": 64,", entry); + defer gpa.free(both); + const c = try parse(gpa, both); + try std.testing.expectEqual(@as(u8, 8), c.widths.of(0, .down).bits); + try std.testing.expectEqual(@as(u8, 4), c.widths.of(1, .down).bits); +} + +test "layer kinds, head dims and the proportional RoPE table" { + try std.testing.expect(!global(0) and global(5) and global(59) and !global(58)); + try std.testing.expectEqual(@as(usize, 512), headDim(11)); + try std.testing.expectEqual(@as(usize, 16), dense31b.kvHeads(1)); + try std.testing.expectEqual(@as(usize, 1), unified12b.kvHeads(47)); + try std.testing.expectEqual(@as(usize, 4096), dense31b.kvWidth()); + try std.testing.expectEqual(@as(usize, 2048), unified12b.kvWidth()); + var g: [global_dim / 2]f32 = undefined; + inverseFrequencies(true, &g); + try std.testing.expectEqual(@as(f32, 1.0), g[0]); + try std.testing.expect(g[63] > 0 and g[64] == 0 and g[255] == 0); + var s: [sliding_dim / 2]f32 = undefined; + inverseFrequencies(false, &s); + try std.testing.expect(s[127] > 0 and s[127] < s[1]); +} diff --git a/zig/src/families/gemma4/cuda_forward.zig b/zig/src/families/gemma4/cuda_forward.zig new file mode 100644 index 0000000000..52d67de196 --- /dev/null +++ b/zig/src/families/gemma4/cuda_forward.zig @@ -0,0 +1,322 @@ +//! Gemma 4 on CUDA: prompt chunks and decode windows through the shared operators in the Metal family's steps. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("cuda_config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Window = cuda.window_attention.Window; + +/// Rows a prompt chunk takes (Metal's chunk_rows: cut where Metal cuts, so fresh and resumed passes keep their bits). +pub const prompt_rows = 512; +/// Rows a decode forward takes over every stream's window (Metal's prompt_rows for windows). +pub const round_rows = 128; +/// Rows a decode attention call takes: its partials' scratch is sized for this many (Metal's attend_rows). +const attend_rows = 16; +/// A sliding layer's ring: the window and a prompt chunk, so a chunk's writes never reach keys its first row reads. +pub const ring = cfg.window + prompt_rows; + +/// A stream's rows in one forward: its cache slot, first position, tokens. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; + +/// sqrt(hidden) as MLX multiplies by it: a weak-typed scalar, rounded to bf16 first (Metal forward.embedScale). +pub fn embedScale(hidden: usize) f32 { + const x: f32 = @sqrt(@as(f32, @floatFromInt(hidden))); + const b: u32 = @bitCast(x); + return @bitCast((b + 0x7FFF + ((b >> 16) & 1)) & 0xFFFF0000); +} + +/// fp32 inverse RoPE frequencies [dim / 2]: sliding as Metal's host table; global 1 / f32(1e6^(2i/512)), 0 past 64. +pub fn inverseFrequencies(is_global: bool, out: []f32) void { + if (!is_global) return cfg.inverseFrequencies(false, out); + for (out, 0..) |*f, i| { + if (i >= cfg.global_rotated) { + f.* = 0; + continue; + } + const p: f32 = @floatCast(std.math.pow(f64, cfg.global_theta, @as(f64, @as(f32, @floatFromInt(2 * i)) / 512.0))); + f.* = @floatCast(1.0 / @as(f64, p)); + } +} + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + geo: cfg.Geometry, + layers: usize, + context: usize, + slots: usize, + kc: []u64, // per layer: keys [slots, len, kv heads, head dim] bf16 (len: the ring, or the context if global) + vc: []u64, + inv_sliding: u64, // fp32 [128] + inv_global: u64, // fp32 [256] + s: Scratch, + buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated + /// Decode forwards of n rows as CUDA graphs, captured at their first use ([n]; null: not yet, or graphs off). + graphs: [round_rows + 1]?cuda.graph.Exec = @splat(null), + graphed: bool = true, + + const Scratch = struct { ids: u64, pos: u64, slot: u64, h: [2]u64, n: u64, xs: u64, q: u64, k: u64, v: u64, qr: u64, kr: u64, vr: u64, att: u64, o: u64, gate: u64, up: u64, act: u64, y: u64, part: u64, logits: u64 }; + + fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += @max(bytes, 256); + return 0; + } + 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; + } + + fn zeros(m: *Model, s: cuda.Stream, bytes: usize) !u64 { + const p = try m.alloc(bytes); + if (m.planned == null) try (DeviceBuffer{ .d = m.d, .ptr = p, .len = @max(bytes, 256), .borrowed = true }).fill8(0, s.handle); + return p; + } + + /// A layer's cache positions: the ring, or every position of the context (global layers). + pub fn cacheLen(m: *const Model, i: usize) usize { + return if (cfg.global(i)) m.context else ring; + } + + /// Bytes one stream's caches take. + pub fn streamBytes(geo: cfg.Geometry, layers: usize, context: usize) usize { + var total: usize = 0; + for (0..layers) |i| total += 2 * (if (cfg.global(i)) context else ring) * geo.kvHeads(i) * cfg.headDim(i) * 2; + return total; + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !Model { + var m = blank(gpa, d, w, context, slots); + errdefer m.deinit(); + try m.build(s); + return m; + } + + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !usize { + var m = blank(gpa, d, w, context, slots); + defer m.deinit(); + m.planned = 0; + try m.build(s); + return m.planned.?; + } + + fn blank(gpa: std.mem.Allocator, d: *const cuda.Driver, w: *const wts.Weights, context: usize, slots: usize) Model { + return .{ .gpa = gpa, .d = d, .w = w, .geo = w.config.geo, .layers = w.layers.len, .context = context, .slots = slots, .kc = &.{}, .vc = &.{}, .inv_sliding = 0, .inv_global = 0, .s = undefined }; + } + + fn build(m: *Model, s: cuda.Stream) !void { + const geo = m.geo; + m.kc = try m.gpa.alloc(u64, m.layers); + m.vc = try m.gpa.alloc(u64, m.layers); + for (0..m.layers) |i| { + const bytes = m.slots * m.cacheLen(i) * geo.kvHeads(i) * cfg.headDim(i) * 2; + m.kc[i] = try m.zeros(s, bytes); + m.vc[i] = try m.zeros(s, bytes); + } + var sl: [cfg.sliding_dim / 2]f32 = undefined; + var gl: [cfg.global_dim / 2]f32 = undefined; + inverseFrequencies(false, &sl); + inverseFrequencies(true, &gl); + m.inv_sliding = try m.alloc(sl.len * 4); + m.inv_global = try m.alloc(gl.len * 4); + if (m.planned == null) { + try (DeviceBuffer{ .d = m.d, .ptr = m.inv_sliding, .len = sl.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&sl)); + try (DeviceBuffer{ .d = m.d, .ptr = m.inv_global, .len = gl.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&gl)); + } + const r = @max(prompt_rows, round_rows); + const h = geo.hidden; + const qw = geo.queryWidth(); + const kw = geo.kvWidth(); + const most = @max(h, geo.intermediate); + // the widest decode partials: attend_rows rows of a global layer over the context in 256-key chunks + const wide: Window = .{ .h = geo.query_heads, .hk = geo.global_kv_heads, .d = cfg.global_dim, .win = 0, .ring = m.context, .scale = 1 }; + const narrow: Window = .{ .h = geo.query_heads, .hk = geo.sliding_kv_heads, .d = cfg.sliding_dim, .win = cfg.window, .ring = ring, .scale = 1 }; + const part = @max(cuda.window_attention.WindowAttention.scratchBytes(wide, attend_rows), cuda.window_attention.WindowAttention.scratchBytes(narrow, attend_rows)); + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .slot = try m.alloc(r * 4), + .h = .{ try m.alloc(r * h * 2), try m.alloc(r * h * 2) }, + .n = try m.alloc(r * h * 2), + .xs = try m.alloc(r * most / 64 * 4), + .q = try m.alloc(r * qw * 2), + .k = try m.alloc(r * kw * 2), + .v = try m.alloc(r * kw * 2), + .qr = try m.alloc(r * qw * 2), + .kr = try m.alloc(r * kw * 2), + .vr = try m.alloc(r * kw * 2), + .att = try m.alloc(r * qw * 2), + .o = try m.alloc(r * h * 2), + .gate = try m.alloc(r * geo.intermediate * 2), + .up = try m.alloc(r * geo.intermediate * 2), + .act = try m.alloc(r * geo.intermediate * 2), + .y = try m.alloc(r * h * 2), + .part = try m.alloc(part), + .logits = try m.alloc(round_rows * cfg.vocab * 2), + }; + if (m.planned == null) try s.synchronize(); + } + + pub fn deinit(m: *Model) void { + for (&m.graphs) |*g| if (g.*) |*e| e.deinit(); + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.kc); + m.gpa.free(m.vc); + } + + fn upload(m: *const 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"); + } + + /// x [rows, k] through `l` into out [rows, n]: a prompt chunk's GEMM, or decode rows on their group sums `xs`. + fn project(m: *const Model, k: *const kern.Kernels, s: cuda.Stream, l: wts.Proj, x: u64, xs: u64, out: u64, rows: usize, prompt: bool) !void { + _ = m; + if (prompt) return switch (l) { + .a4 => |q| k.affine.prompt(s, x, q, out, rows), + .an => |q| k.wide.prompt(s, x, q, out, rows), + }; + const tile: usize = switch (l) { + .a4 => k.affine.verifyRows(), + .an => 16, + }; + var at: usize = 0; + while (at < rows) : (at += tile) { + const n = @min(tile, rows - at); + const xa = x + at * l.k() * 2; + const sa = xs + at * (l.k() / 64) * 4; + const oa = out + at * l.n() * 2; + switch (l) { + .a4 => |q| try k.affine.verify(s, xa, sa, q, oa, n), + .an => |q| try k.wide.decode(s, xa, sa, q, oa, n), + } + } + } + + /// Every chain's rows (a prompt chunk, or decode windows); rows [head0, head0 + heads)'s capped logits at s.logits. + pub fn forward(m: *Model, k: *const kern.Kernels, s: cuda.Stream, chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void { + var rows: usize = 0; + for (chains) |c| { + if (c.tokens.len == 0 or c.slot >= m.slots or c.p0 + c.tokens.len > m.context) return error.BadChain; + rows += c.tokens.len; + } + if (rows > @as(usize, if (prompt) prompt_rows else round_rows) or (prompt and chains.len != 1)) return error.BadWindow; + if (heads > round_rows or head0 + heads > rows) return error.BadWindow; + const ids = try m.gpa.alloc(i32, 3 * rows); + defer m.gpa.free(ids); + var at: usize = 0; + for (chains) |c| { + for (c.tokens, 0..) |t, i| { + ids[at + i] = @intCast(t); + ids[rows + at + i] = @intCast(c.p0 + i); + ids[2 * rows + at + i] = @intCast(c.slot); + } + at += c.tokens.len; + } + try m.upload(s, m.s.ids, std.mem.sliceAsBytes(ids[0..rows])); + try m.upload(s, m.s.pos, std.mem.sliceAsBytes(ids[rows .. 2 * rows])); + try m.upload(s, m.s.slot, std.mem.sliceAsBytes(ids[2 * rows ..])); + // a decode forward's launches depend on its row count only (positions and slots are read on the device) + if (prompt or !m.graphed or head0 != 0 or heads != rows) return m.launch(k, s, chains, rows, prompt, head0, heads); + if (m.graphs[rows] == null) m.graphs[rows] = try m.capture(k, s, rows); + try m.graphs[rows].?.launchOn(s); + } + + /// A decode forward of `rows` rows, every row's logits, captured as a graph. + fn capture(m: *Model, k: *const kern.Kernels, s: cuda.Stream, rows: usize) !cuda.graph.Exec { + try s.synchronize(); + try cuda.graph.beginCapture(s, .thread_local); + m.launch(k, s, &.{}, rows, false, 0, rows) catch |e| { + if (cuda.graph.endCapture(s)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return e; + }; + var g = try cuda.graph.endCapture(s); + defer g.deinit(); + return g.instantiate(); + } + + /// The forward's launches over staged ids, positions and slots: no host copies, so a graph can hold them. + fn launch(m: *Model, k: *const kern.Kernels, s: cuda.Stream, chains: []const Chain, rows: usize, prompt: bool, head0: usize, heads: usize) !void { + const h = m.geo.hidden; + const x = m.s; + try k.embed.affine(s, x.ids, m.w.embed_w, m.w.embed_s, m.w.embed_b, x.h[0], rows, h, .{ .bits = @intCast(m.w.embed_bits), .group = cfg.group, .scale = embedScale(h) }); + // layer 0's input norm alone; every later one is the previous layer's dense tail + try k.norms.addRmsTail(s, x.h[0], null, m.w.layers[0].input_norm, x.h[1], x.n, x.xs, rows, h, cfg.eps, .{ .round = true }); + for (0..m.layers) |i| try m.layer(k, s, chains, i, rows, prompt); + if (heads == 0) return; + try m.project(k, s, m.w.head, x.n + head0 * h * 2, x.xs + head0 * (h / 64) * 4, x.logits, heads, false); + try k.act.softcap(s, x.logits, x.logits, heads, cfg.vocab, cfg.softcap); + } + + /// Layer i on the rows: residual in s.h[0], its normed input (and sums) in s.n (s.xs); the same on return. + fn layer(m: *Model, k: *const kern.Kernels, s: cuda.Stream, chains: []const Chain, i: usize, rows: usize, prompt: bool) !void { + const geo = m.geo; + const l = m.w.layers[i]; + const x = m.s; + const h = geo.hidden; + const global = cfg.global(i); + const dh = cfg.headDim(i); + const kvh = geo.kvHeads(i); + const qw = geo.query_heads * dh; + const kw = kvh * dh; + try m.project(k, s, l.q, x.n, x.xs, x.q, rows, prompt); + try m.project(k, s, l.k, x.n, x.xs, x.k, rows, prompt); + if (l.v) |v| try m.project(k, s, v, x.n, x.xs, x.v, rows, prompt); + const vsrc = if (global) x.k else x.v; // global layers: v is k_proj's output, unit-normed and not rotated + const src: cuda.qk_rows.Src = .{ .q = x.q, .q_ld = @intCast(qw), .q_hs = @intCast(dh), .k = x.k, .k_ld = @intCast(kw), .k_hs = @intCast(dh), .v = vsrc, .v_ld = @intCast(kw), .v_hs = @intCast(dh) }; + const shape: cuda.qk_rows.Shape = .{ .norms = l.q_norm, .norm_k = l.k_norm, .norms_bf16 = true, .inv = if (global) m.inv_global else m.inv_sliding, .rot = dh, .eps = cfg.eps, .round = true, .v_norm = true }; + const len = m.cacheLen(i); + try k.qk.rowsFrom(s, src, .{ .h = geo.query_heads, .hk = kvh, .dim = dh }, shape, x.pos, x.slot, .{ .q = x.qr, .k = x.kr, .v = x.vr, .kc = m.kc[i], .vc = m.vc[i], .len = len }, rows); + const w: Window = .{ .h = geo.query_heads, .hk = kvh, .d = dh, .win = if (global) 0 else cfg.window, .ring = len, .scale = 1, .p32 = true }; + if (prompt) { + const c = chains[0]; + const off = c.slot * len * kw * 2; + try k.win.promptRows(s, x.qr, m.kc[i] + off, m.vc[i] + off, x.att, c.p0, rows, if (global) .{ .h = w.h, .hk = w.hk, .d = w.d, .win = 0, .ring = 0, .scale = 1, .p32 = true } else w, null); + } else { + var r: usize = 0; + while (r < rows) : (r += attend_rows) { + const n = @min(attend_rows, rows - r); + try k.win.decode(s, x.qr + r * qw * 2, m.kc[i], m.vc[i], x.att + r * qw * 2, x.pos + r * 4, x.slot + r * 4, w, x.part, n); + } + } + try k.act.groupSums(s, x.att, qw, x.xs, rows, qw); + try m.project(k, s, l.o, x.att, x.xs, x.o, rows, prompt); + try k.norms.addRmsTail(s, x.o, x.h[0], l.pre_feedforward, x.h[1], x.n, x.xs, rows, h, cfg.eps, .{ .round = true, .branch = l.post_attention }); + try m.project(k, s, l.gate, x.n, x.xs, x.gate, rows, prompt); + try m.project(k, s, l.up, x.n, x.xs, x.up, rows, prompt); + try k.act.geglu(s, x.gate, geo.intermediate, x.up, geo.intermediate, x.act, x.xs, rows, geo.intermediate); + try m.project(k, s, l.down, x.act, x.xs, x.y, rows, prompt); + const next = if (i + 1 < m.layers) m.w.layers[i + 1].input_norm else m.w.final_norm; + try k.norms.addRmsTail(s, x.y, x.h[1], next, x.h[0], x.n, x.xs, rows, h, cfg.eps, .{ .round = true, .branch = l.post_feedforward, .scalar = l.scalar }); + } +}; + +test "the embedding scale is bf16(sqrt(hidden)): 62 for the 12B, 73.5 for the 31B" { + try std.testing.expectEqual(@as(f32, 62.0), embedScale(3840)); + try std.testing.expectEqual(@as(f32, 73.5), embedScale(5376)); +} + +test "RoPE tables: sliding as Metal's host table, global zero past 64 pairs and decreasing before" { + var sl: [128]f32 = undefined; + var ref: [128]f32 = undefined; + inverseFrequencies(false, &sl); + cfg.inverseFrequencies(false, &ref); + try std.testing.expectEqualSlices(f32, &ref, &sl); + var gl: [256]f32 = undefined; + inverseFrequencies(true, &gl); + try std.testing.expectEqual(@as(f32, 1), gl[0]); + for (1..64) |i| try std.testing.expect(gl[i] < gl[i - 1] and gl[i] > 0); + for (gl[64..]) |v| try std.testing.expectEqual(@as(f32, 0), v); +} diff --git a/zig/src/families/gemma4/cuda_kernels.zig b/zig/src/families/gemma4/cuda_kernels.zig new file mode 100644 index 0000000000..5d3a6fde6b --- /dev/null +++ b/zig/src/families/gemma4/cuda_kernels.zig @@ -0,0 +1,60 @@ +//! Gemma 4's launches on CUDA, every one a shared operator (#1048): projections, rows, attention, argmax, draws. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const Kernels = struct { + /// qmm_group, qmm_prefill, lane_gemv, affine4_pack, sample, tree_round, affine8, affine_prefill (other widths). + mods: [8]cuda.Module, + affine: cuda.qlinear.Affine4Kernels, + wide: cuda.qlinear.AffineNKernels, + draws: cuda.sample.Draws, + round: cuda.tree_round.Kernels, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + qk: cuda.qk_rows.QkRows, + win: cuda.window_attention.WindowAttention, + act: cuda.act_rows.ActRows, + + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + const kk = cuda.kernels; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.affine4_pack, kk.sample, kk.tree_round, kk.affine8, kk.affine_prefill }; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + k.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.capability() == 121; + k.affine = try cuda.qlinear.Affine4Kernels.resolve(d, k.mods[0], k.mods[1], k.mods[2], k.mods[3], sms, gb10); + errdefer k.affine.deinit(); + k.wide = try cuda.qlinear.AffineNKernels.resolve(d, k.mods[6], k.mods[7], k.mods[3], sms, gb10); + errdefer k.wide.deinit(); + k.draws = try cuda.sample.Draws.resolve(k.mods[4]); + k.round = try cuda.tree_round.Kernels.resolve(k.mods[5]); + 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.win = try cuda.window_attention.WindowAttention.load(d); + errdefer k.win.unload(); + k.act = try cuda.act_rows.ActRows.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.act.unload(); + k.win.unload(); + k.qk.unload(); + k.norms.unload(); + k.embed.unload(); + k.wide.deinit(); + k.affine.deinit(); + for (&k.mods) |*m| m.unload(); + } +}; diff --git a/zig/src/families/gemma4/cuda_native.zig b/zig/src/families/gemma4/cuda_native.zig new file mode 100644 index 0000000000..de549834e1 --- /dev/null +++ b/zig/src/families/gemma4/cuda_native.zig @@ -0,0 +1,219 @@ +//! Gemma 4 (affine-4/64 text) behind the native CUDA registry; images, widths, drafters (steps 9-12) refused. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const fwd = @import("cuda_forward.zig"); +const be = @import("cuda_backend.zig"); +const snap = @import("cuda_snap.zig"); +const asst = @import("cuda_assistant.zig"); + +/// The registry's entry for one model_type: gemma4_unified (12B) or gemma4 (31B). +pub fn Family(comptime name: []const u8) type { + return struct { + pub const model_type = name; + pub const formats: []const []const u8 = &.{ "mlx-q2g64", "mlx-q3g64", "mlx-q4g64", "mlx-q5g64", "mlx-q6g64", "mlx-q8g64" }; + pub const default_context: i64 = 32768; + pub const max_segments: u32 = 1; + pub const prompt_rows: u32 = be.prompt_rows; + pub const snap_lookahead: u32 = 0; + pub const takes_drafter = true; + pub const openDrafter = native.openDrafter; + pub const explains_open = true; + pub const Room = native.Room; + pub const Options = native.Options; + pub const Loaded = native.Loaded; + pub const open = native.open; + pub const snap_bytes = native.snap_bytes; + pub const snap_save = native.snap_save; + pub const snap_restore = native.snap_restore; + pub const snap_drop = native.snap_drop; + pub const explain = native.explain; + }; +} + +const native = @This(); + +/// The server's reading of the room left after the weights; null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; +/// No `vision` or `precision` field: the server refuses --vision and --precision for this family by name. +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + target: *be.Gemma, +}; + +const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: fwd.Model, gb: be.Gemma }; + +/// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses. +fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + if (fixed and asked > be.round_rows) return error.ParallelPastRound; + var n = std.math.clamp(@as(usize, asked), 1, be.round_rows); + while (try plan.bytes(n) > room) : (n -= 1) { + if (fixed) return error.ParallelPastMemory; + if (n == 1) return error.NoStreamFits; + } + return n; +} + +const Plan = struct { + own: *const Owned, + d: *const cuda.Driver, + context: usize, + + fn bytes(p: Plan, n: usize) !u64 { + const o = p.own; + return try fwd.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n) + be.Gemma.openBytes(n); + } +}; + +/// Positions a stream's caches hold: whole 1,024s. +fn positions(context: usize) usize { + return (context + 1023) / 1024 * 1024; +} + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + if (o.lanes_fixed and o.lanes > be.round_rows) return error.ParallelPastRound; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.k = try kern.Kernels.load(ctx); + errdefer own.k.deinit(); + own.s = try cuda.Stream.init(ctx.d, false); + errdefer own.s.deinit(); + own.w = try wts.load(gpa, io, ctx.d, own.s, &own.k, dir, null); + errdefer own.w.deinit(); + const context = @min(positions(o.context), own.w.config.context); + const plan: Plan = .{ .own = own, .d = ctx.d, .context = context }; + const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + if (e == error.ParallelPastMemory) { + const most = fit(room, o.lanes, false, plan) catch 0; + std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, gib(room), most, if (most == 1) "" else "s" }); + } + return e; + }; + own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, context, n); + errdefer own.m.deinit(); + own.m.graphed = !envOff("TENSORFOLD_GEMMA4_GRAPHS"); // =0: every decode forward eager + own.gb = try be.Gemma.init(gpa, ctx.d, &own.m, &own.k, own.s); + return .{ + .backend = own.gb.backend(), + .facts = own.gb.facts(), + .rows = be.max_window, + .stream_bytes = fwd.Model.streamBytes(own.w.config.geo, own.w.layers.len, context), + .free_streams = @intCast(n), + .ctx = own, + .deinit = release, + .target = &own.gb, + }; +} + +fn envOff(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return false; + return std.mem.eql(u8, std.mem.span(v), "0"); +} + +fn gib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + +fn gemma(ptr: *anyopaque) *be.Gemma { + return @ptrCast(@alignCast(ptr)); +} + +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(gemma(ptr).m, at); +} + +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + return @ptrCast(try gemma(ptr).saveState(owner, at)); +} + +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + return gemma(ptr).restoreState(owner, @ptrCast(@alignCast(saved))); +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + gemma(ptr).dropState(@ptrCast(@alignCast(saved))); +} + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => "Gemma 4 on CUDA drafts with an assistant loaded by --drafter, or from the prompt's own text", + error.GemmaDrafterNotServed => "--drafter: Gemma 4 on CUDA loads the MTP assistant (gemma4_unified_assistant) so far; a DFlash drafter is not built yet", + error.AssistantForAnotherTarget => "--drafter: this assistant was trained for another Gemma 4 target (its backbone width or key heads differ)", + error.UnsupportedAssistantGeometry, error.UnsupportedAssistantType, error.UnsupportedAssistantRotary, error.UnsupportedAssistantQuantization, error.AssistantOrderedEmbeddingsNotServed => "--drafter: this Gemma 4 assistant's layout is not the stock one this engine serves", + error.DrafterBitsUnsupported => "--drafter-bits: the Gemma 4 assistant runs at 4 bits on CUDA", + error.StructuresUnsupported => "structured output is not built for Gemma 4 on CUDA yet", + error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them", + error.CapturedKernels => "Gemma 4 runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: Gemma 4 on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer Gemma 4 streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.NoStreamFits => "the CUDA memory budget fits no Gemma 4 stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + else => null, + }; +} + +/// --drafter: the 12B's MTP assistant over the target's caches (a DFlash drafter is #1048's step 11). +pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []const u8, bits: u8) !struct { drafter: lanes.drafter.Drafter, ctx: *anyopaque, close: *const fn (*anyopaque) void, reserve: *const fn (*anyopaque, u32) anyerror!void } { + if (bits != 4) return error.DrafterBitsUnsupported; + const own: *Owned = @ptrCast(@alignCast(p)); + if (own.gb.lanes.count() != 0) return error.TapsAfterStreams; + if (!try asst.isAssistant(gpa, io, dir)) return error.GemmaDrafterNotServed; + const d = try asst.Drafter.init(gpa, io, &own.gb, dir); + return .{ .drafter = d.drafter(), .ctx = d, .close = closeDrafter, .reserve = reserveNone }; +} + +fn closeDrafter(p: *anyopaque) void { + const d: *asst.Drafter = @ptrCast(@alignCast(p)); + d.deinit(); +} + +fn reserveNone(_: *anyopaque, _: u32) anyerror!void {} + +fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + own.gb.deinit(); + own.m.deinit(); + own.w.deinit(); + own.s.deinit(); + own.k.deinit(); + own.gpa.destroy(own); +} + +test "every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits, error.GemmaDrafterNotServed, error.AssistantForAnotherTarget, error.DrafterBitsUnsupported }) |e| { + try std.testing.expect(explain(null, e) != null); + } + try std.testing.expect(explain(null, error.OutOfMemory) == null); + try std.testing.expectEqualStrings("gemma4_unified", Family("gemma4_unified").model_type); +} + +const Fake = struct { + stream: u64, + fixed: u64, + fn bytes(f: Fake, n: usize) !u64 { + return f.fixed + n * f.stream; + } +}; + +test "the streams allocated at open fit the room the weights leave: auto takes what fits, a fixed count refuses" { + const g: u64 = 1 << 30; + const one: Fake = .{ .stream = g, .fixed = 0 }; + try std.testing.expectEqual(@as(usize, 2), try fit(2 * g, 8, false, one)); + try std.testing.expectError(error.ParallelPastMemory, fit(2 * g, 8, true, one)); + try std.testing.expectError(error.NoStreamFits, fit(g - 1, 8, false, one)); + try std.testing.expectError(error.ParallelPastRound, fit(1024 * g, be.round_rows + 1, true, one)); +} diff --git a/zig/src/families/gemma4/cuda_snap.zig b/zig/src/families/gemma4/cuda_snap.zig new file mode 100644 index 0000000000..45375230d6 --- /dev/null +++ b/zig/src/families/gemma4/cuda_snap.zig @@ -0,0 +1,108 @@ +//! A stream's caches at a prompt position: global layers' [0, at), each sliding layer's last window - 1 positions. +const std = @import("std"); +const cuda = @import("cuda"); +const cfg = @import("cuda_config.zig"); +const fwd = @import("cuda_forward.zig"); + +/// A kept state: per layer K then V, global layers' positions [0, at), sliding layers' last positions in order. +pub const Saved = struct { at: u32, buf: cuda.DeviceBuffer }; + +/// Positions a sliding layer keeps at `at`. +fn tail(at: usize) usize { + return @min(at, cfg.window - 1); +} + +fn rowBytes(m: *const fwd.Model, i: usize) usize { + return m.geo.kvHeads(i) * cfg.headDim(i) * 2; +} + +/// Bytes a state at `at` takes. +pub fn bytes(m: *const fwd.Model, at: usize) u64 { + var n: usize = 0; + for (0..m.layers) |i| n += 2 * (if (cfg.global(i)) at else tail(at)) * rowBytes(m, i); + return n; +} + +/// Each layer's cache, K then V, positions [first, at) of `slot`: one copy a run of ring rows, in order. +fn each(m: *const fwd.Model, slot: usize, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + var off: usize = 0; + for (0..m.layers) |i| for ([_][]u64{ m.kc, m.vc }) |cache| { + const rb = rowBytes(m, i); + const global = cfg.global(i); + const len = m.cacheLen(i); + const base = cache[i] + slot * len * rb; + var p = if (global) 0 else at - tail(at); + while (p < at) { + const row = p % len; + const run = @min(at - p, len - row); + try f(ctx, base + row * rb, off, run * rb); + off += run * rb; + p += run; + } + }; +} + +const Copy = struct { d: *const cuda.Driver, s: cuda.Stream, buf: u64, out: bool }; + +fn copy(c: Copy, dev: u64, off: usize, len: usize) anyerror!void { + const src = if (c.out) dev else c.buf + off; + const dst = if (c.out) c.buf + off else dev; + try c.d.check(c.d.api.cuMemcpyDtoDAsync_v2(dst, src, len, c.s.handle), "cuMemcpyDtoDAsync"); +} + +/// Copy `slot`'s state at `at` into new storage, queued on `s` behind the chunk that ends there. +pub fn save(gpa: std.mem.Allocator, m: *const fwd.Model, s: cuda.Stream, slot: usize, at: u32) !*Saved { + if (at == 0 or at > m.context) return error.NotAtMark; + const k = try gpa.create(Saved); + errdefer gpa.destroy(k); + k.* = .{ .at = at, .buf = try cuda.DeviceBuffer.alloc(m.d, bytes(m, at)) }; + errdefer k.buf.free(); + try each(m, slot, at, Copy{ .d = m.d, .s = s, .buf = k.buf.ptr, .out = true }, copy); + return k; +} + +/// Make `slot`'s caches hold `k`'s state, queued on `s` before the next chunk. +pub fn restore(m: *const fwd.Model, s: cuda.Stream, slot: usize, k: *const Saved) !void { + if (k.at > m.context) return error.PastContext; + try each(m, slot, k.at, Copy{ .d = m.d, .s = s, .buf = k.buf.ptr, .out = false }, copy); +} + +/// Free a state once the stream's queued copies of it are done. +pub fn drop(gpa: std.mem.Allocator, s: cuda.Stream, k: *Saved) void { + s.synchronize() catch {}; + k.buf.free(); + gpa.destroy(k); +} + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +test "a state holds global layers' positions and each sliding ring's last window - 1, wrapped runs in order" { + var k = [_]u64{ 1 << 30, 2 << 30, 3 << 30, 4 << 30, 5 << 30, 6 << 30 }; + var v = [_]u64{ 7 << 30, 8 << 30, 9 << 30, 10 << 30, 11 << 30, 12 << 30 }; + var m: fwd.Model = undefined; + m.geo = cfg.unified12b; + m.layers = 6; // layer 5 is global + m.context = 4096; + m.kc = &k; + m.vc = &v; + const at = 2000; // sliding: positions 977..1999, ring rows 977..1535 then 0..463 + const sb = 8 * 256 * 2; + const gb = 1 * 512 * 2; + try testing.expectEqual(@as(u64, 5 * 2 * 1023 * sb + 2 * at * gb), bytes(&m, at)); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&m, 1, at, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 5 * 2 * 2 + 2), runs.items.len); + const r0 = runs.items[0]; + try testing.expectEqual([3]usize{ (1 << 30) + (fwd.ring + 977) * sb, 0, (fwd.ring - 977) * sb }, r0); + try testing.expectEqual([3]usize{ (1 << 30) + fwd.ring * sb, (fwd.ring - 977) * sb, 464 * sb }, runs.items[1]); + const last = runs.items[runs.items.len - 1]; + try testing.expectEqual([3]usize{ (12 << 30) + 4096 * gb, 2 * 5 * 1023 * sb + at * gb, at * gb }, last); +} diff --git a/zig/src/families/gemma4/cuda_weights.zig b/zig/src/families/gemma4/cuda_weights.zig new file mode 100644 index 0000000000..eb75261360 --- /dev/null +++ b/zig/src/families/gemma4/cuda_weights.zig @@ -0,0 +1,186 @@ +//! Gemma 4's weights on CUDA: MLX affine-4 projections packed into qlinear's tiles, norms as stored, the tied head. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const cfg = @import("cuda_config.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Affine4 = cuda.qlinear.Affine4; +const AffineN = cuda.qlinear.AffineN; + +/// A projection at its MLX width: affine-4 tiles, or the tiles of every other width. +pub const Proj = union(enum) { + a4: Affine4, + an: AffineN, + + pub fn n(p: Proj) usize { + return switch (p) { + inline else => |q| q.n, + }; + } + + pub fn k(p: Proj) usize { + return switch (p) { + inline else => |q| q.k, + }; + } +}; + +/// A layer's tensors: each projection as qlinear tiles at its width, bf16 vectors as stored (norm weights with no 1 +). +pub const Layer = struct { + q: Proj, + k: Proj, + v: ?Proj, // global layers have none: v is k_proj's output (attention_k_eq_v) + o: Proj, + gate: Proj, + up: Proj, + down: Proj, + q_norm: u64, // bf16 [head dim] + k_norm: u64, + input_norm: u64, // bf16 [hidden] + post_attention: u64, + pre_feedforward: u64, + post_feedforward: u64, + scalar: u64, // bf16 [1]: the residual's layer scalar +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: cfg.Config, + /// The embedding table as MLX stores it: words u32 [V, H bits / 32], scales and biases [V, H / 64]. + embed_w: u64 = 0, + embed_bits: u8 = 4, + embed_s: u64 = 0, + embed_b: u64 = 0, + /// The same table as the tied head's tiles. + head: Proj = undefined, + final_norm: u64 = 0, + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } +}; + +const Ctx = struct { + w: *Weights, + d: *const cuda.Driver, + s: cuda.Stream, + k: *const kern.Kernels, + ck: *core.checkpoint.Checkpoint, + stage: DeviceBuffer, + + fn tensor(x: *Ctx, name: []const u8, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + return x.w.upload(x.d, (try x.ck.expect(name, dtype, shape)).bytes); + } + + /// An n-by-k MLX affine projection of `bits` codes (words, scales, biases) staged and packed into qlinear's tiles. + fn linear(x: *Ctx, prefix: []const u8, n: usize, k: usize, bits: u8) !Proj { + var buf: [192]u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k * bits / 32 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (specs, &raw) |sp, *r| { + const t = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ prefix, sp[0] }), sp[1], &.{ n, sp[2] }); + if (at + t.bytes.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, t.bytes); + r.* = x.stage.ptr + at; + at += t.bytes.len; + } + const served = cuda.qlinear.servedBits(bits); + const out: Proj = if (served == 4) blk: { + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad }; + try x.k.affine.pack(x.s, raw[0], raw[1], raw[2], q); + break :blk .{ .a4 = q }; + } else blk: { + const l = AffineN.layout(n, k, served); + const q: AffineN = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad, .bits = served }; + try x.k.wide.packWords(x.s, raw[0], bits, n, k, q.w); + try x.k.wide.packTable(x.s, raw[1], n, k, q.npad, q.s); + try x.k.wide.packTable(x.s, raw[2], n, k, q.npad, q.b); + break :blk .{ .an = q }; + }; + try x.s.synchronize(); // the next projection's staging overwrites this one's + return out; + } +}; + +const lm = "language_model.model."; + +/// Admission as Metal's (MLX affine 2-6 and 8 bits, groups of 64), every tensor checked. `layers` caps them (checks). +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, k: *const kern.Kernels, dir: []const u8, layers: ?usize) !Weights { + const c = try cfg.Config.read(gpa, io, dir); + const geo = c.geo; + var w: Weights = .{ .gpa = gpa, .config = c }; + errdefer w.deinit(); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const h = geo.hidden; + // the largest projection as stored: the tied head's table + var stage = try DeviceBuffer.alloc(d, cfg.vocab * h * c.widths.embed.bits / 8 + 2 * cfg.vocab * (h / 64) * 2); + defer stage.free(); + var x: Ctx = .{ .w = &w, .d = d, .s = s, .k = k, .ck = &ck, .stage = stage }; + w.embed_bits = c.widths.embed.bits; + w.embed_w = try x.tensor(lm ++ "embed_tokens.weight", .u32, &.{ cfg.vocab, h * w.embed_bits / 32 }); + w.embed_s = try x.tensor(lm ++ "embed_tokens.scales", .bf16, &.{ cfg.vocab, h / 64 }); + w.embed_b = try x.tensor(lm ++ "embed_tokens.biases", .bf16, &.{ cfg.vocab, h / 64 }); + w.head = try x.linear(lm ++ "embed_tokens", cfg.vocab, h, w.embed_bits); + w.final_norm = try x.tensor(lm ++ "norm.weight", .bf16, &.{h}); + const n = @min(layers orelse geo.layers, geo.layers); + w.layers = try gpa.alloc(Layer, n); + var buf: [192]u8 = undefined; + for (w.layers, 0..) |*l, i| { + const dh = cfg.headDim(i); + const kvh = geo.kvHeads(i); + const q_out = geo.query_heads * dh; + const kv_out = kvh * dh; + const p = try std.fmt.allocPrint(gpa, lm ++ "layers.{d}.", .{i}); + defer gpa.free(p); + const name = struct { + fn of(b: []u8, pre: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "{s}{s}", .{ pre, suffix }); + } + }.of; + l.q = try x.linear(try name(&buf, p, "self_attn.q_proj"), q_out, h, c.widths.of(i, .q).bits); + l.k = try x.linear(try name(&buf, p, "self_attn.k_proj"), kv_out, h, c.widths.of(i, .k).bits); + l.v = if (cfg.global(i)) null else try x.linear(try name(&buf, p, "self_attn.v_proj"), kv_out, h, c.widths.of(i, .v).bits); + l.o = try x.linear(try name(&buf, p, "self_attn.o_proj"), h, q_out, c.widths.of(i, .o).bits); + l.gate = try x.linear(try name(&buf, p, "mlp.gate_proj"), geo.intermediate, h, c.widths.of(i, .gate).bits); + l.up = try x.linear(try name(&buf, p, "mlp.up_proj"), geo.intermediate, h, c.widths.of(i, .up).bits); + l.down = try x.linear(try name(&buf, p, "mlp.down_proj"), h, geo.intermediate, c.widths.of(i, .down).bits); + l.q_norm = try x.tensor(try name(&buf, p, "self_attn.q_norm.weight"), .bf16, &.{dh}); + l.k_norm = try x.tensor(try name(&buf, p, "self_attn.k_norm.weight"), .bf16, &.{dh}); + l.input_norm = try x.tensor(try name(&buf, p, "input_layernorm.weight"), .bf16, &.{h}); + l.post_attention = try x.tensor(try name(&buf, p, "post_attention_layernorm.weight"), .bf16, &.{h}); + l.pre_feedforward = try x.tensor(try name(&buf, p, "pre_feedforward_layernorm.weight"), .bf16, &.{h}); + l.post_feedforward = try x.tensor(try name(&buf, p, "post_feedforward_layernorm.weight"), .bf16, &.{h}); + l.scalar = try x.tensor(try name(&buf, p, "layer_scalar"), .bf16, &.{1}); + } + return w; +} diff --git a/zig/src/families/gemma4/cuda_widths.zig b/zig/src/families/gemma4/cuda_widths.zig new file mode 100644 index 0000000000..4dca1ceba6 --- /dev/null +++ b/zig/src/families/gemma4/cuda_widths.zig @@ -0,0 +1,133 @@ +//! widths.zig for CUDA builds, line for line: MLX affine widths per module (core.quant.mlx imported by name). +const std = @import("std"); +const mlx = @import("core").quant.mlx; +const cfg = @import("cuda_config.zig"); + +pub const Width = mlx.Width; + +/// The bits the lane projection and core affine_mm both read; the glue's group sums are 64 wide. +pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; +pub const group = 64; +/// The width the reg lane kernel and the fused one-row prologues take. +pub const fast: Width = .{ .bits = 4, .group = group }; + +/// A layer's linears, in Layer.mlx's order. +pub const Module = enum { q, k, v, o, gate, up, down }; +pub const module_names = [_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" }; + +pub const Widths = struct { + layers: [cfg.max_layers][module_names.len]Width = @splat(@splat(fast)), + /// The tied embedding and head. + embed: Width = fast, + /// The 31B tower's (and a quantized 12B's) embed_vision projection. + vision: Width = fast, + + pub fn of(w: *const Widths, layer: usize, m: Module) Width { + return w.layers[layer][@backingInt(m)]; + } + + /// Every linear at the fast width: the one-row windows may take the fused pipelines. + pub fn allFast(w: *const Widths, layers: usize) bool { + if (!std.meta.eql(w.embed, fast)) return false; + for (w.layers[0..layers]) |l| for (l) |x| if (!std.meta.eql(x, fast)) return false; + return true; + } +}; + +pub const Error = error{ UnsupportedGemmaQuantization, GemmaWidthNotRead, GemmaModuleUnquantized } || std.mem.Allocator.Error; + +/// Whether the family's kernels read `w`. +pub fn readable(w: Width) bool { + return w.group == group and std.mem.indexOfScalar(u8, &bit_widths, w.bits) != null; +} + +/// Each module's width from `quant` (`quantization` or `quantization_config`); global layers have no v_proj. +pub fn read(gpa: std.mem.Allocator, quant: std.json.ObjectMap, layers: usize) Error!Widths { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var it = quant.iterator(); + while (it.next()) |e| { + const k = e.key_ptr.*; + if (std.mem.eql(u8, k, "bits") or std.mem.eql(u8, k, "group_size") or std.mem.eql(u8, k, "mode")) continue; + if (e.value_ptr.* != .object or e.value_ptr.object.count() == 0) { + std.log.warn("gemma4: {s} is not a width table (false or {{}} marks a module unquantized); every linear is affine", .{k}); + return error.GemmaModuleUnquantized; + } + } + const q = (mlx.detect(a, .{ .root = quant, .quantization = quant }) catch |err| return switch (err) { + error.UnsupportedQuantization => error.UnsupportedGemmaQuantization, // two spellings of one module disagree + else => |e| e, + }) orelse return error.UnsupportedGemmaQuantization; // not MLX affine + var w: Widths = .{}; + var key: [96]u8 = undefined; + for (0..layers) |i| for (module_names, 0..) |name, mi| { + if (cfg.global(i) and mi == @backingInt(Module.v)) continue; + w.layers[i][mi] = try named(q, std.fmt.bufPrint(&key, "layers.{d}.{s}", .{ i, name }) catch unreachable); + }; + w.embed = try named(q, "embed_tokens"); + w.vision = try named(q, "embed_vision.embedding_projection"); + return w; +} + +fn named(q: mlx.Config, key: []const u8) Error!Width { + const w = q.width(key) catch { + std.log.warn("gemma4: {s}'s width is not MLX affine bits 2-6 or 8", .{key}); + return error.GemmaWidthNotRead; + }; + if (!readable(w)) { + std.log.warn("gemma4: {s} is {d}-bit in groups of {d}; the kernels read bits 2-6 and 8 in groups of 64", .{ key, w.bits, w.group }); + return error.GemmaWidthNotRead; + } + return w; +} + +fn table(text: []const u8) !Widths { + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, text, .{}); + defer parsed.deinit(); + return read(std.testing.allocator, parsed.value.object, 12); +} + +test "a uniform table gives every module the global width" { + const w4 = try table("{\"bits\": 4, \"group_size\": 64, \"mode\": \"affine\"}"); + try std.testing.expect(w4.allFast(12)); + const w8 = try table("{\"bits\": 8, \"group_size\": 64}"); + try std.testing.expect(!w8.allFast(12)); + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, w8.of(11, .down)); + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, w8.embed); +} + +test "per-module entries under any spelling reach their module only" { + const w = try table( + \\{"bits": 4, "group_size": 64, "mode": "affine", + \\ "language_model.model.layers.1.self_attn.v_proj": {"bits": 6, "group_size": 64}, + \\ "model.layers.2.mlp.down_proj": {"bits": 6, "group_size": 64, "mode": "affine"}, + \\ "language_model.model.embed_tokens": {"bits": 8, "group_size": 64}} + ); + try std.testing.expectEqual(@as(u8, 6), w.of(1, .v).bits); + try std.testing.expectEqual(@as(u8, 4), w.of(1, .k).bits); + try std.testing.expectEqual(@as(u8, 4), w.of(0, .v).bits); + try std.testing.expectEqual(@as(u8, 6), w.of(2, .down).bits); + try std.testing.expectEqual(@as(u8, 8), w.embed.bits); + try std.testing.expect(!w.allFast(12)); +} + +test "widths no kernel reads are refused by name, never guessed" { + const cases = .{ + .{ "{\"bits\": 7, \"group_size\": 64}", error.UnsupportedGemmaQuantization }, + .{ "{\"bits\": 4, \"group_size\": 32}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 128}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"mode\": \"mxfp4\"}", error.UnsupportedGemmaQuantization }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 7, \"group_size\": 64}}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 4, \"group_size\": 64, \"mode\": \"nvfp4\"}}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": false}", error.GemmaModuleUnquantized }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {}}", error.GemmaModuleUnquantized }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 6, \"group_size\": 64}, \"model.layers.3.mlp.up_proj\": {\"bits\": 8, \"group_size\": 64}}", error.UnsupportedGemmaQuantization }, + }; + inline for (cases) |c| try std.testing.expectError(c[1], table(c[0])); +} + +test "a global layer's missing v_proj takes no entry" { + const w = try table("{\"bits\": 4, \"group_size\": 64, \"layers.5.self_attn.v_proj\": {\"bits\": 4, \"group_size\": 32}}"); + try std.testing.expectEqual(@as(u8, 4), w.of(5, .k).bits); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a8463051a5..691eef4509 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -15,6 +15,7 @@ const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); const qwen35 = @import("qwen35"); const kolibri1 = @import("kolibri1"); +const gemma4 = @import("gemma4_cuda"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -25,7 +26,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31 }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; diff --git a/zig/tests/cuda/act_round_check.zig b/zig/tests/cuda/act_round_check.zig new file mode 100644 index 0000000000..1f981eec34 --- /dev/null +++ b/zig/tests/cuda/act_round_check.zig @@ -0,0 +1,75 @@ +//! MLX's activation rows against core/act_ref.zig's float64 ends, their 64-value sums against the kernel's fp32 tree. +const std = @import("std"); +const core = @import("core"); + +const aref = core.act_ref; +const bfr = core.bf16_round; +const c = core.contract.spec; + +/// fp32 tanhf and expf within 2 ulps of their value (2^-22) and the division within half one: twice that, either way. +pub const spread = 0x1p-21; + +/// One call's host rows as bf16 bits: gate and up rows at their strides, GeGLU or SwiGLU. +pub const Case = struct { geglu: bool, rows: usize, n: usize, gate: []const u16, g_ld: usize, up: []const u16, u_ld: usize }; + +/// What a check saw: outputs whose float64 range is one value, a range, non-finite in float64, and sums. +pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0, sums: usize = 0 }; + +/// act_rows.cu's group_sum8 over stored outputs: each thread's 8 in order from 0, then 8 lanes' xor tree 4, 2, 1. +pub fn groupSum(y: []const u16) f32 { + var s: [8]f32 = undefined; + for (&s, 0..) |*v, l| { + v.* = 0; + for (y[8 * l ..][0..8]) |h| v.* += bfr.value(h); + } + var m: usize = 4; + while (m > 0) : (m >>= 1) { + var t: [8]f32 = undefined; + for (&t, 0..) |*v, l| v.* = s[l] + s[l ^ m]; + s = t; + } + return s[0]; +} + +/// Element i of row `row`: its float64 ends and point. +pub fn want(k: Case, row: usize, i: usize) struct { lo: f64, hi: f64, point: f64 } { + const x: f64 = bfr.value(k.gate[row * k.g_ld + i]); + const u: f64 = bfr.value(k.up[row * k.u_ld + i]); + const r = aref.mlxRange(k.geglu, x, u, spread); + return .{ .lo = r[0], .hi = r[1], .point = if (k.geglu) aref.geglu(x, u, 0) else aref.swigluRound(x, u, 0, 0) }; +} + +/// Every output within float64's ends (NaN and infinities as float64's point), the sums bit-equal to the fp32 tree. +pub fn verify(k: Case, out: []const u16, xs: ?[]const f32, stats: *Stats, msg: []u8) !?[]const u8 { + for (0..k.rows) |row| { + for (0..k.n) |i| { + const w = want(k, row, i); + const g: f64 = bfr.value(out[row * k.n + i]); + const ok = if (std.math.isFinite(w.lo) and std.math.isFinite(w.hi)) blk: { + if (w.lo == w.hi) stats.decided += 1 else stats.ranged += 1; + break :blk std.math.isFinite(g) and w.lo <= g and g <= w.hi; + } else blk: { + stats.nonfinite += 1; + break :blk c.classOf(g) == c.classOf(w.point); + }; + if (!ok) return try std.fmt.bufPrint(msg, "row {d} out[{d}] = {e} outside float64's [{e}, {e}]", .{ row, i, g, w.lo, w.hi }); + } + if (xs) |sums| for (0..k.n / 64) |j| { + const t = groupSum(out[row * k.n + 64 * j ..][0..64]); + const got = sums[row * (k.n / 64) + j]; + stats.sums += 1; + if (@as(u32, @bitCast(got)) != @as(u32, @bitCast(t)) and !(std.math.isNan(got) and std.math.isNan(t))) + return try std.fmt.bufPrint(msg, "row {d} sum {d} = {e} where the fp32 tree of the stored out gives {e}", .{ row, j, got, t }); + }; + } + return null; +} + +test "the fp32 group tree: lanes then the xor tree, not one chain" { + var y: [64]u16 = undefined; + for (&y, 0..) |*v, i| v.* = bfr.bits(if (i == 0) 0x1p24 else 1); + var chain: f32 = 0; + for (y) |h| chain += bfr.value(h); + try std.testing.expectEqual(@as(f32, 0x1p24), chain); // one chain loses all 63 ones beside 2^24 + try std.testing.expectEqual(@as(f32, 0x1p24 + 56), groupSum(&y)); // lanes 1 to 7 keep theirs +} diff --git a/zig/tests/cuda/act_rows_tests.zig b/zig/tests/cuda/act_rows_tests.zig index 43ddaf4b7e..ccd9bf10ea 100644 --- a/zig/tests/cuda/act_rows_tests.zig +++ b/zig/tests/cuda/act_rows_tests.zig @@ -4,6 +4,7 @@ const cuda = @import("cuda"); const core = @import("core"); const check = @import("check.zig"); const rows = @import("norm_rows_tests.zig"); +const act_ = @import("act_round_check.zig"); const aref = core.act_ref; const counts = rows.counts; @@ -108,6 +109,33 @@ fn dconv(r: *Rig, g: *const cuda.act_rows.ActRows, d: usize, gs: usize, seg: usi try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * d)), std.mem.sliceAsBytes(got[0 .. m * d])); } +/// MLX's GeGLU or SwiGLU steps over [gate | up] rows (fused) or two buffers with padded rows; bits at every row count. +fn mlx(r: *Rig, g: *const cuda.act_rows.ActRows, n: usize, geglu: bool, fused: bool, stats: *act_.Stats) !void { + const ld = if (fused) 2 * n else n + 64; + const gate = r.bfs(most * ld, 4); + const up = if (fused) gate[n..] else r.bfs(most * ld, 2); + const dg = try r.dev(std.mem.sliceAsBytes(gate)); + const du = if (fused) dg + 2 * n else try r.dev(std.mem.sliceAsBytes(up)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * n * 2); + xss[c] = try r.out(m * n / 64 * 4); + if (geglu) try g.geglu(r.s, dg, ld, du, ld, os[c], xss[c], m, n) else try g.swigluRound(r.s, dg, ld, du, ld, os[c], xss[c], m, n); + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "{s} in MLX's bf16 steps, {d} wide{s}", .{ if (geglu) "geglu" else "swiglu", n, if (fused) ", [gate | up] rows" else "" }); + const o = try r.back(u16, os[counts.len - 1], most * n); + const xs = try r.back(f32, xss[counts.len - 1], most * n / 64); + var msg: [200]u8 = undefined; + const k: act_.Case = .{ .geglu = geglu, .rows = most, .n = n, .gate = gate, .g_ld = ld, .up = up, .u_ld = ld }; + if (try act_.verify(k, o, xs, stats, &msg)) |why| try check.expect(false, "{s}: {s}", .{ label, why }); + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * n)), std.mem.sliceAsBytes(o[0 .. m * n])); + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(f32, xss[c], m * n / 64)), std.mem.sliceAsBytes(xs[0 .. m * n / 64])); + } +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -128,5 +156,14 @@ pub fn run(gpu: check.Gpu) !void { try dconv(&r, &g, 5120, 16, 16, 0, false); // DFlash2's block conv (groups of 16) try dconv(&r, &g, 5120, 16, 16, 1, true); try dconv(&r, &g, 256, 64, 7, 0, true); + var stats: act_.Stats = .{}; + for ([_]usize{ 15360, 21504 }) |n| for ([_]bool{ true, false }) |fused| try mlx(&r, &g, n, true, fused, &stats); // Gemma 4 12B, 31B + try mlx(&r, &g, 8192, true, true, &stats); // the 12B's assistant + try mlx(&r, &g, 10752, false, true, &stats); // the 31B's DFlash drafter + try mlx(&r, &g, 10752, false, false, &stats); + for ([_]usize{ 64, 960, 1024, 1088, 2048, 2112 }) |n| for ([_]bool{ true, false }) |geglu| try mlx(&r, &g, n, geglu, true, &stats); + try std.testing.expectError(error.RowShape, g.geglu(r.s, 0, 128, 0, 128, 0, null, 1, 96)); + try std.testing.expectError(error.RowShape, g.swigluRound(r.s, 0, 100, 0, 128, 0, null, 1, 64)); check.pass("act rows: swiglu, output gate, dconv within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); + check.pass("act rows MLX steps: {d} outputs bit-equal to float64's rounding points, {d} within tanhf's or expf's ulps, {d} non-finite as float64; {d} group sums equal to the fp32 tree; every row count's bits equal", .{ stats.decided, stats.ranged, stats.nonfinite, stats.sums }); } diff --git a/zig/tests/cuda/contract_act_round.zig b/zig/tests/cuda/contract_act_round.zig new file mode 100644 index 0000000000..d320a2f79c --- /dev/null +++ b/zig/tests/cuda/contract_act_round.zig @@ -0,0 +1,163 @@ +//! act_rows.cu's MLX activations under the contract (GeGLU, SwiGLU's bf16 steps), and controls each dropping a point. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const act_ = @import("act_round_check.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const bfr = core.bf16_round; +const aref = core.act_ref; + +/// A registration: a served kernel (keep null), or a control keeping all rounding points but one (act_round.cuh). +pub const Entry = struct { name: []const u8, geglu: bool, keep: ?i32 = null }; + +const geglu_all = 127; +const swiglu_all = 31; + +pub const entries = [_]Entry{ + .{ .name = "cuda act tf_geglu_rows (MLX's bf16 steps)", .geglu = true }, + .{ .name = "cuda act tf_swiglu_rows_round (MLX's bf16 steps)", .geglu = false }, + .{ .name = "control: GeGLU without x^3's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 1) }, + .{ .name = "control: GeGLU without c1 x^3's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 2) }, + .{ .name = "control: GeGLU without the inner sum's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 4) }, + .{ .name = "control: GeGLU without c0 inner's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 8) }, + .{ .name = "control: GeGLU without tanh's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 16) }, + .{ .name = "control: GeGLU without 1 + tanh's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 32) }, + .{ .name = "control: GeGLU without the gelu's rounding", .geglu = true, .keep = geglu_all & ~@as(i32, 64) }, + .{ .name = "control: SwiGLU steps without exp's rounding", .geglu = false, .keep = swiglu_all & ~@as(i32, 1) }, + .{ .name = "control: SwiGLU steps without 1 + e's rounding", .geglu = false, .keep = swiglu_all & ~@as(i32, 2) }, + .{ .name = "control: SwiGLU steps without the reciprocal's rounding", .geglu = false, .keep = swiglu_all & ~@as(i32, 4) }, + .{ .name = "control: SwiGLU steps without 1 - the reciprocal's rounding", .geglu = false, .keep = swiglu_all & ~@as(i32, 8) }, + .{ .name = "control: SwiGLU steps without silu's rounding", .geglu = false, .keep = swiglu_all & ~@as(i32, 16) }, +}; + +fn depth(_: c.Shape) usize { + return 0; +} + +pub fn spec(e: Entry) c.Spec { + return .{ + .name = e.name, + .accum = .none, // no sums in an output; its 64-value sums are checked bit for bit against the fp32 tree + .depth = depth, + .mids = 16, // gate 3 is the kernel's own check: float64's rounding points, ranged over tanhf's or expf's ulps + .out = .bf16, + .tiles = .{ .{ .size = 16 }, .{ .size = 1024, .step = 64, .min = 64 }, .{} }, + // Gemma 4 12B's and 31B's MLP, the 31B DFlash drafter's and the 12B assistant's + .served = &.{ .{ .rows = 17, .n = 15360, .k = 1 }, .{ .rows = 3, .n = 21504, .k = 1 }, .{ .rows = 16, .n = 10752, .k = 1 }, .{ .rows = 6, .n = 8192, .k = 1 } }, + }; +} + +pub const Act = struct { + rig: rig_.Rig, + entry: Entry, + act: *const cuda.act_rows.ActRows, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + k: act_.Case = undefined, + gu: u64 = 0, + out: u64 = 0, + xs: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Act) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Act); + + pub fn kernel(self: *Act) g.Kernel { + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = if (self.entry.keep != null) .shapes else null }; + } + + pub fn prepare(self: *Act, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n = shape.n; + if (n % 64 != 0 or n == 0 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(); + const gu = try a.alloc(u16, shape.rows * 2 * n); // [gate | up] rows, as a fused gate_up projection writes them + for (gu, 0..) |*v, i| { + const up = i % (2 * n) >= n; + const big: f32 = if (i % 64 == 7) 20 else 1; + v.* = bfr.bits(if (up) rng.float(f32) * 4 - 2 else (rng.float(f32) * 8 - 4) * big); + } + if (fill == .nonfinite) { + const planted = [_]struct { usize, usize, f32 }{ .{ 0, n / 3, std.math.nan(f32) }, .{ 1, n / 2, std.math.inf(f32) }, .{ 2, n - 1, -std.math.inf(f32) }, .{ 3, n + 5, std.math.inf(f32) }, .{ 3, 0, -std.math.inf(f32) } }; + for (planted) |p| if (p[0] < shape.rows) { + gu[p[0] * 2 * n + p[1]] = bfr.bits(p[2]); + }; + } + self.k = .{ .geglu = self.entry.geglu, .rows = shape.rows, .n = n, .gate = gu, .g_ld = 2 * n, .up = gu[n..], .u_ld = 2 * n }; + self.gu = try self.rig.input(std.mem.sliceAsBytes(gu)); + self.out = try self.rig.output(shape.rows * n * 2); + self.xs = try self.rig.output(shape.rows * n / 64 * 4); + return true; + } + + pub fn run(self: *Act) !void { + const k = self.k; + const s = self.rig.s; + const up = self.gu + 2 * k.n; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + args.add(self.gu); + args.add(@as(c_longlong, @intCast(k.g_ld))); + args.add(up); + args.add(@as(c_longlong, @intCast(k.u_ld))); + args.add(self.out); + args.add(self.xs); + args.add(@as(c_int, @intCast(k.n))); + args.add(@as(c_int, @intFromBool(k.geglu))); + args.add(@as(c_int, self.entry.keep.?)); + const t: u32 = if (k.rows <= cuda.launch.few_rows) 128 else 256; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(k.rows), .y = @intCast((k.n / 8 + t - 1) / t), .z = 1 }, .block = .{ .x = t } }, s, &args); + } else if (k.geglu) try self.act.geglu(s, self.gu, k.g_ld, up, k.u_ld, self.out, self.xs, k.rows, k.n) else try self.act.swigluRound(s, self.gu, k.g_ld, up, k.u_ld, self.out, self.xs, k.rows, k.n); + try s.synchronize(); + } + + pub fn count(self: *Act) usize { + return self.k.rows * self.k.n; + } + + pub fn reference(self: *Act, out: []f64, mag: []f64) !void { + for (0..self.k.rows) |row| for (0..self.k.n) |i| { + out[row * self.k.n + i] = act_.want(self.k, row, i).point; + }; + @memset(mag, 0); + } + + /// Every output within float64's ends at the rounding points (a dropped point leaves them), sums as fp32's tree. + pub fn check(self: *Act, _: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const n = self.count(); + const out = try a.alloc(u16, n); + const xs = try a.alloc(f32, n / 64); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.bytes(self.out, n * 2)); + @memcpy(std.mem.sliceAsBytes(xs), try self.rig.bytes(self.xs, n / 64 * 4)); + var stats: act_.Stats = .{}; + return act_.verify(self.k, out, xs, &stats, &self.msg); + } +}; + +test { + _ = act_; +} + +test "every control keeps all but one rounding point, each point once" { + var seen = [2]i32{ 0, 0 }; + for (entries) |e| if (e.keep) |k| { + const all: i32 = if (e.geglu) geglu_all else swiglu_all; + const dropped = all & ~k; + try std.testing.expectEqual(@as(u32, 1), @popCount(dropped)); + try std.testing.expect(seen[@intFromBool(e.geglu)] & dropped == 0); + seen[@intFromBool(e.geglu)] |= dropped; + }; + try std.testing.expectEqual([2]i32{ swiglu_all, geglu_all }, seen); +} diff --git a/zig/tests/cuda/contract_embed.zig b/zig/tests/cuda/contract_embed.zig new file mode 100644 index 0000000000..b3ad45d4af --- /dev/null +++ b/zig/tests/cuda/contract_embed.zig @@ -0,0 +1,159 @@ +//! embed_rows.cu's MLX affine rows under the contract (4-bit, and 2 to 6 or 8 bits with a bf16 scale), with controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const ec = @import("embed_check.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const bfr = core.bf16_round; + +/// A registration: the code width, a scaled form (tf_embed_rows_affine) or not (tf_embed_rows_affine4), or a control. +pub const Entry = struct { name: []const u8, bits: u4, scaled: bool = true, keep: ?i32 = null }; + +const round = 1; // embed_affine.cuh's ROUND +const fold = 2; // and its FOLD + +pub const entries = [_]Entry{ + .{ .name = "cuda embed_rows tf_embed_rows_affine4 (MLX 4-bit rows)", .bits = 4, .scaled = false }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (2-bit, bf16 scale)", .bits = 2 }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (3-bit, bf16 scale)", .bits = 3 }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (4-bit, bf16 scale)", .bits = 4 }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (5-bit, bf16 scale)", .bits = 5 }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (6-bit, bf16 scale)", .bits = 6 }, + .{ .name = "cuda embed_rows tf_embed_rows_affine (8-bit, bf16 scale)", .bits = 8 }, + .{ .name = "control: scaled rows without the dequantized value's bf16 rounding", .bits = 4, .keep = 0 }, + .{ .name = "control: scaled rows with the scale folded into the group's scale and bias", .bits = 4, .keep = fold }, +}; + +fn depth(_: c.Shape) usize { + return 0; +} + +/// Shape: rows x row width (n); the group is 64. +pub fn spec(e: Entry) c.Spec { + return .{ + .name = e.name, + .accum = .none, // no sums: every output is checked bit for bit against float64 at its rounding points + .depth = depth, + .out = .bf16, + .tiles = .{ .{}, .{ .size = 2048, .step = 64, .min = 64 }, .{} }, // 256 threads of 8 values a block + // Gemma 4 12B's and 31B's hidden + .served = &.{ .{ .rows = 17, .n = 3840, .k = 1 }, .{ .rows = 3, .n = 5376, .k = 1 } }, + }; +} + +const vocab = 97; + +pub const Embed = struct { + rig: rig_.Rig, + entry: Entry, + e: *const cuda.embed_rows.EmbedRows, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + k: ec.Case = undefined, + dev: [5]u64 = undefined, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Embed) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Embed); + + pub fn kernel(self: *Embed) g.Kernel { + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = if (self.entry.keep != null) .shapes else null }; + } + + pub fn prepare(self: *Embed, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const d = shape.n; + if (d % 64 != 0 or d == 0 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(); + const bits = self.entry.bits; + const codes = try a.alloc(u32, vocab * d); + for (codes) |*q| q.* = rng.int(u32) & ((@as(u32, 1) << bits) - 1); + const s = try a.alloc(u16, vocab * d / 64); + const b = try a.alloc(u16, vocab * d / 64); + for (s, b) |*x, *y| { // the checkpoint's ranges, one group in 16 far past them + const big: f32 = if (rng.uintLessThan(u32, 16) == 0) 0x1p40 else 1; + x.* = bfr.bits((rng.float(f32) * 0.5 - 0.18) * big); + y.* = bfr.bits((rng.float(f32) * 6.8 - 4.1) * big); + } + const ids = try a.alloc(i32, shape.rows); + for (ids, 0..) |*x, i| x.* = @intCast(if (i == 0) vocab - 1 else rng.uintLessThan(usize, vocab)); + if (fill == .nonfinite) { + const planted = [_]struct { usize, u16 }{ .{ 0, bfr.bits(std.math.nan(f32)) }, .{ 1, bfr.bits(std.math.inf(f32)) }, .{ 2, bfr.bits(-std.math.inf(f32)) } }; + for (planted, 0..) |p, i| if (i < shape.rows) { + const t: usize = @intCast(ids[i]); + (if (i == 1) b else s)[t * d / 64 + p[0] + 1] = p[1]; + }; + } + if (fill == .stress) for (s) |*x| { // fp32's largest scales: q s + b and its scaled value past bf16's range + x.* = bfr.bits(0x1p126); + }; + const scale: ?f32 = if (self.entry.scaled) bfr.round32(@sqrt(@as(f32, @floatFromInt(d)))) else null; + const w = try ec.pack(a, codes, bits); + self.k = .{ .rows = shape.rows, .d = d, .bits = bits, .group = 64, .scale = scale, .ids = ids, .w = w, .s = s, .b = b }; + for ([_][]const u8{ std.mem.sliceAsBytes(ids), std.mem.sliceAsBytes(w), std.mem.sliceAsBytes(s), std.mem.sliceAsBytes(b) }, 0..) |bytes, i| + self.dev[i] = try self.rig.input(bytes); + self.dev[4] = try self.rig.output(shape.rows * d * 2); + return true; + } + + pub fn run(self: *Embed) !void { + const k = self.k; + const st = self.rig.s; + const v = self.dev; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + for (v) |p| args.add(p); + for ([_]usize{ k.d, k.bits, k.group }) |x| args.add(@as(c_int, @intCast(x))); + args.add(k.scale.?); + args.add(@as(c_int, self.entry.keep.?)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(k.rows), .y = @intCast((k.d / 8 + 255) / 256), .z = 1 }, .block = .{ .x = 256 } }, st, &args); + } else if (self.entry.scaled) { + try self.e.affine(st, v[0], v[1], v[2], v[3], v[4], k.rows, k.d, .{ .bits = k.bits, .scale = k.scale }); + } else try self.e.affine4(st, v[0], v[1], v[2], v[3], v[4], k.rows, k.d); + try st.synchronize(); + } + + pub fn count(self: *Embed) usize { + return self.k.rows * self.k.d; + } + + pub fn reference(self: *Embed, out: []f64, mag: []f64) !void { + for (0..self.k.rows) |row| for (0..self.k.d) |i| { + out[row * self.k.d + i] = ec.want(self.k, row, i); + }; + @memset(mag, 0); + } + + /// Every output bit-equal to float64's value at the rounding points (a dropped or moved point leaves it). + pub fn check(self: *Embed, _: g.Gate) !?[]const u8 { + const n = self.count(); + const out = try self.arena.?.allocator().alloc(u16, n); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.bytes(self.dev[4], n * 2)); + var stats: ec.Stats = .{}; + return ec.verify(self.k, out, &stats, &self.msg); + } +}; + +test { + _ = ec; +} + +test "the controls drop the one rounding point, or fold the scale past it" { + var seen: usize = 0; + for (entries) |e| if (e.keep) |k| { + try std.testing.expect(k & round == 0 and e.scaled); + seen += 1; + }; + try std.testing.expectEqual(@as(usize, 2), seen); +} diff --git a/zig/tests/cuda/contract_norms_tail.zig b/zig/tests/cuda/contract_norms_tail.zig new file mode 100644 index 0000000000..cbeb0aa0ff --- /dev/null +++ b/zig/tests/cuda/contract_norms_tail.zig @@ -0,0 +1,241 @@ +//! norms.cu's bf16 residual tails under the contract (MLX's rounded rows and once-rounded ones), and their controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const tail_ = @import("norms_tail_check.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const rows = core.contract.rows; +const ref = core.norms_ref; +const bfr = core.bf16_round; + +/// The kernels by the output a registration checks (y, or the residual h), and the controls dropping one point each. +pub const Which = enum { round_y, round_h, once_y, once_h, drop_branch_row, drop_branch_out, drop_sum, drop_row }; + +fn depth(_: c.Shape) usize { + return tail_.depth; +} + +const why = "an output follows its sum of squares as ss^-1/2 through bf16 rounding points: a bf16 sum's error is a few output ulps, and an fp16 sum overflows to a zero output, 51x the bound that carries the points"; + +pub fn spec(which: Which) c.Spec { + return .{ + .name = switch (which) { + .round_y => "cuda norm tail tf_add_rms_rows_bf16h_tail_round (MLX's rows), output y", + .round_h => "cuda norm tail tf_add_rms_rows_bf16h_tail_round (MLX's rows), residual h", + .once_y => "cuda norm tail tf_add_rms_rows_bf16h_tail (rounded once), output y", + .once_h => "cuda norm tail tf_add_rms_rows_bf16h_tail (rounded once), residual h", + .drop_branch_row => "control: MLX tail without the branch's normalized-row rounding", + .drop_branch_out => "control: MLX tail without the branch output's rounding", + .drop_sum => "control: MLX tail without the sum's rounding before the scalar", + .drop_row => "control: MLX tail without the output's normalized-row rounding", + }, + .accum = .f32, + .depth = depth, + .ops = 8, + .mids = 4, // a rounding point decided the other way by the fp32 inverse: two output half-ulps a point + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 2048, .step = 64, .min = 64 } }, + .served = &.{ .{ .rows = 17, .n = 3840, .k = 3840 }, .{ .rows = 3, .n = 5376, .k = 5376 } }, // Gemma 4 12B, 31B + .sums = .{ .f16 = 40, .bf16 = 0, .why = why }, + }; +} + +pub fn control(which: Which) ?g.Gate { + return switch (which) { + .round_y, .round_h, .once_y, .once_h => null, + else => .shapes, + }; +} + +/// The controls' symbols in contract_controls.cu. +pub fn symbol(which: Which) [:0]const u8 { + return switch (which) { + inline else => |w| "tf_ctl_tail_" ++ @tagName(w), + }; +} + +fn round(which: Which) bool { + return which != .once_y and which != .once_h; +} + +fn hView(which: Which) bool { + return which == .round_h or which == .once_h; +} + +/// 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]); + } +}; + +pub const Tail = struct { + rig: rig_.Rig, + which: Which, + norms: *const cuda.norms.Norms, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + k: tail_.Case = undefined, + dev: [5]u64 = undefined, // x, r, wb, sc, w + out: [3]u64 = undefined, // h, y, xs + msg: [200]u8 = undefined, + + pub fn deinit(self: *Tail) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Tail); + + pub fn kernel(self: *Tail) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = &vt, .control = control(self.which) }; + } + + fn bfs(rng: std.Random, v: []u16, scale: f32) void { + for (v, 0..) |*x, i| x.* = bfr.bits((rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); + } + + /// Norm weights as Gemma 4's: signed, mostly in [-2, 6), one in 97 thirty times larger (12B: -143 to 193). + fn weights(rng: std.Random, v: []u16) void { + for (v, 0..) |*x, i| x.* = bfr.bits((rng.float(f32) * 8 - 2) * (if (i % 97 == 3) @as(f32, 30) else 1)); + } + + /// Gate 4's rows: integer squares summing to about 2^20, exact in fp32 in any order. + fn integers(rng: std.Random, v: []u16) void { + for (v) |*x| x.* = bfr.bits(@as(f32, @floatFromInt(rng.uintLessThan(u32, 65))) - 32); + } + + pub fn prepare(self: *Tail, 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(); + const n = shape.rows * width; + const x = try a.alloc(u16, n); + const r = try a.alloc(u16, n); + const wb = try a.alloc(u16, width); + const w = try a.alloc(u16, width); + bfs(rng, x, 1); + bfs(rng, r, 4); + weights(rng, wb); + weights(rng, w); + var sc = bfr.bits(0.004 + 0.9 * rng.float(f32)); // Gemma 4 12B's layer scalars run 0.0045 to 0.89 + if (fill == .nonfinite) { + const planted = [_]struct { usize, bool, usize, f32 }{ .{ 0, true, width / 3, std.math.nan(f32) }, .{ 1, false, width / 2, std.math.inf(f32) }, .{ 2, true, width - 1, -std.math.inf(f32) }, .{ 3, false, 0, std.math.inf(f32) }, .{ 3, false, width - 1, -std.math.inf(f32) } }; + for (planted) |p| if (p[0] < shape.rows) { + (if (p[1]) x else r)[p[0] * width + p[2]] = bfr.bits(p[3]); + }; + } + if (fill == .stress) { // only the checked output's sum of squares counts, so make it the exact one + sc = bfr.bits(1); + integers(rng, x); + if (hView(self.which)) @memset(r, 0) else { + @memset(wb, 0); // the branch adds zeros: h is the residual's integers + integers(rng, r); + } + } + self.k = .{ .rows = shape.rows, .width = width, .round = round(self.which), .x = x, .w = w, .r = r, .wb = wb, .sc = sc }; + for ([_][]const u8{ std.mem.sliceAsBytes(x), std.mem.sliceAsBytes(r), std.mem.sliceAsBytes(wb), std.mem.asBytes(&sc), std.mem.sliceAsBytes(w) }, &self.dev) |bytes, *d| d.* = try self.rig.input(bytes); + if (hView(self.which)) { + self.out[0] = try self.rig.output(n * 2); + self.out[1] = try self.rig.output(n * 2); + } else { + self.out[1] = try self.rig.output(n * 2); + self.out[0] = try self.rig.output(n * 2); + } + self.out[2] = try self.rig.output(n / 64 * 4); + return true; + } + + pub fn run(self: *Tail) !void { + const s = self.rig.s; + const k = self.k; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + for (self.dev[0..4]) |d| args.add(d); + args.add(self.dev[4]); + for (self.out) |d| args.add(d); + args.add(@as(c_int, @intCast(k.width))); + args.add(@as(f32, 1e-6)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(k.rows), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &args); + } else try self.norms.addRmsTail(s, self.dev[0], self.dev[1], self.dev[4], self.out[0], self.out[1], self.out[2], k.rows, k.width, 1e-6, .{ .round = k.round, .branch = self.dev[2], .scalar = self.dev[3] }); + try s.synchronize(); + } + + pub fn count(self: *Tail) usize { + return self.k.rows * self.k.width; + } + + fn view(self: *Tail, a: std.mem.Allocator, out: []f64, sums: anytype) !void { + const width = self.k.width; + const h = try a.alloc(f64, 4 * width); + defer a.free(h); + for (0..self.k.rows) |row| { + const at = row * width; + const t = try tail_.tailOf(a, self.k, row, 0); + ref.tail(t.t, t.x, t.w, sums, h[0..width], h[width .. 2 * width], h[2 * width .. 3 * width], h[3 * width ..]); + @memcpy(out[at..][0..width], if (hView(self.which)) h[0..width] else h[2 * width .. 3 * width]); + } + } + + pub fn reference(self: *Tail, want: []f64, mag: []f64) !void { + var arena: std.heap.ArenaAllocator = .init(self.rig.gpu.gpa); + defer arena.deinit(); + try self.view(arena.allocator(), want, ref.Exact{}); + for (mag, want) |*m, v| m.* = @abs(v) / 2; // v ~ ss^-1/2 + } + + pub fn simulate(self: *Tail, acc: c.Accum, order: g.Order, out: []f64) !void { + var arena: std.heap.ArenaAllocator = .init(self.rig.gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + try self.view(a, out, Sim{ .acc = acc, .order = order, .buf = try a.alloc(f64, self.k.width) }); + } + + /// Every output in float64's range over the fp32 inverse's spread (a dropped point leaves it), sums as fp32's tree. + pub fn check(self: *Tail, _: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const n = self.count(); + const h = try a.alloc(u16, n); + const y = try a.alloc(u16, n); + const xs = try a.alloc(f32, n / 64); + @memcpy(std.mem.sliceAsBytes(h), try self.rig.bytes(self.out[0], n * 2)); + @memcpy(std.mem.sliceAsBytes(y), try self.rig.bytes(self.out[1], n * 2)); + @memcpy(std.mem.sliceAsBytes(xs), try self.rig.bytes(self.out[2], n / 64 * 4)); + var stats: tail_.Stats = .{}; + return tail_.verify(self.rig.gpu.gpa, self.k, h, y, xs, &stats, &self.msg); + } +}; + +test { + _ = tail_; +} + +test "gate 4's rows: the checked output's squares are exact in fp32, fp16 overflows and bf16 misses" { + var terms: [5376]f64 = undefined; + var prng = std.Random.DefaultPrng.init(0x6b63); + 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 > 0x1p20); + 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); + } +} diff --git a/zig/tests/cuda/contract_qk_round.zig b/zig/tests/cuda/contract_qk_round.zig new file mode 100644 index 0000000000..73d1c79652 --- /dev/null +++ b/zig/tests/cuda/contract_qk_round.zig @@ -0,0 +1,226 @@ +//! qk_rows' tf_qk_rows_once and tf_qk_rows_round (dims to 512, v's unit norm, MLX's rows) and their controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const qr = @import("qk_round_check.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const bfr = core.bf16_round; +const ref = core.qk_ref; +const qk = cuda.qk_rows; + +pub const Rope = enum { table, inv, proportional }; + +/// A registration: head dim, query heads a key head, the options; a control keeps the bits in `keep` (qk_round.cuh). +pub const Form = struct { name: []const u8, dim: usize, group: usize, rope: Rope, round: bool, v_norm: bool, v_from_k: bool = false, keep: ?i32 = null }; + +const row_pt = 1; +const out_pt = 2; +const vnorm_pt = 4; + +/// Gemma 4 31B's global (k's columns as v) and sliding layers, the 12B's global rounded once; controls drop one each. +pub const forms = [_]Form{ + .{ .name = "cuda qk_rows tf_qk_rows_round (MLX's rows, dim 512, k's columns as v by a unit norm, 64 of 256 pairs rotated)", .dim = 512, .group = 8, .rope = .proportional, .round = true, .v_norm = true, .v_from_k = true }, + .{ .name = "cuda qk_rows tf_qk_rows_round (MLX's rows, dim 256, unit v norm, rotation from inv_freq)", .dim = 256, .group = 2, .rope = .inv, .round = true, .v_norm = true }, + .{ .name = "cuda qk_rows tf_qk_rows_once (dim 512, norm, rotation from a table, v as is)", .dim = 512, .group = 16, .rope = .table, .round = false, .v_norm = false }, + .{ .name = "control: MLX's qk rows without the normalized row's rounding", .dim = 512, .group = 8, .rope = .proportional, .round = true, .v_norm = true, .v_from_k = true, .keep = out_pt | vnorm_pt }, + .{ .name = "control: MLX's qk rows without the weighted row's rounding before RoPE", .dim = 512, .group = 8, .rope = .proportional, .round = true, .v_norm = true, .v_from_k = true, .keep = row_pt | vnorm_pt }, + .{ .name = "control: MLX's qk rows without v's unit norm", .dim = 512, .group = 8, .rope = .proportional, .round = true, .v_norm = true, .v_from_k = true, .keep = row_pt | out_pt }, +}; + +fn depth(shape: c.Shape) usize { + return qr.depth(shape.k); +} + +const why = "an output follows its sum of squares as ss^-1/2 through bf16 rounding points: a bf16 sum's error is a few output ulps, and an fp16 sum overflows to a zero output, 51x the bound that carries the points"; + +/// Shape: rows x query heads (n) x head dim (k). +pub fn spec(f: Form) c.Spec { + const served = switch (f.group) { + 8 => &[_]c.Shape{ .{ .rows = 1, .n = 32, .k = 512 }, .{ .rows = 300, .n = 32, .k = 512 } }, + 2 => &[_]c.Shape{ .{ .rows = 1, .n = 32, .k = 256 }, .{ .rows = 300, .n = 32, .k = 256 } }, + else => &[_]c.Shape{ .{ .rows = 1, .n = 16, .k = 512 }, .{ .rows = 300, .n = 16, .k = 512 } }, + }; + return .{ + .name = f.name, + .accum = .f32, + .depth = depth, + .ops = 12, // the division, eps, the root, the reciprocal, the products and RoPE's + .mids = 4, // a rounding point decided the other way by the fp32 inverse: two output half-ulps a point + .out = .bf16, + .tiles = .{ .{ .size = 16 }, .{}, .{} }, + .served = served, + .sums = .{ .f16 = 40, .bf16 = 0, .why = why }, + }; +} + +const positions = 4096; +const len = 600; + +/// Sums of squares in an accumulator and order (gate 4's simulation), or float64's. +const Sim = struct { acc: c.Accum, order: rows.Order }; + +pub const Qk = struct { + rig: rig_.Rig, + form: Form, + k: *const qk.QkRows, + ctl: ?cuda.Function = null, + arena: std.heap.ArenaAllocator, + case: qr.Case = undefined, + 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); + + pub fn prepare(self: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const f = self.form; + if (shape.rows == 0 or shape.k != f.dim or shape.n % f.group != 0 or shape.n == 0) return false; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const hk = shape.n / f.group; + const ld = (shape.n + (if (f.v_from_k) hk else 2 * hk)) * f.dim; + const x = try a.alloc(u16, r * ld); + for (x, 0..) |*h, i| h.* = bfr.bits(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + // Gemma 4 12B's norm weights: q 0.95 to 1.05, k 0.06 to 0.13 + const wb = try a.alloc(u16, 2 * f.dim); + for (wb, 0..) |*v, i| v.* = bfr.bits(if (i < f.dim) 0.9 + 0.2 * rng.float(f32) else 0.05 + 0.1 * rng.float(f32)); + const rot = f.dim; + const inv = try a.alloc(f32, rot / 2); + for (inv, 0..) |*v, j| v.* = if (f.rope == .proportional and j >= 64) 0 else @floatCast(1 / std.math.pow(f64, if (f.rope == .proportional) 1e6 else 1e4, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(rot)))); + const pos = try a.alloc(i32, r); + const slot = try a.alloc(i32, r); + for (pos, slot, 0..) |*p, *sl, i| { // distinct (slot, pos % len) a row: no row's cache entry overwrites another's + p.* = @intCast((i / 2) % len + len * (i % 5)); + 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, f.dim - 1 } }; + for (at, put) |p, v| if (p[0] < r) { + x[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 (x) |*h| h.* = bfr.bits(@as(f32, @floatFromInt(32 + rng.uintLessThan(u32, 33))) * (if (rng.boolean()) @as(f32, 1) else -1)); + @memset(inv, 0); + } + const cosv = try a.alloc(f32, positions * rot); + const sinv = try a.alloc(f32, positions * rot); + if (f.rope == .table) for (0..positions) |p| for (0..rot) |i| { + const ang = @as(f64, @floatFromInt(p)) * inv[i % (rot / 2)]; + cosv[p * rot + i] = @floatCast(@cos(ang)); + sinv[p * rot + i] = @floatCast(@sin(ang)); + }; + // the host's cos and sin of each row: the table's, or of fp32 angle pos * inv as the kernel forms it + const cs = try a.alloc(f64, r * rot); + const sn = try a.alloc(f64, r * rot); + for (0..r) |row| for (0..rot) |i| { + const p: usize = @intCast(pos[row]); + if (f.rope == .table) { + cs[row * rot + i] = cosv[p * rot + i]; + sn[row * rot + i] = sinv[p * rot + i]; + } else { + const ang: f64 = @as(f32, @floatFromInt(p)) * inv[i % (rot / 2)]; + cs[row * rot + i] = @cos(ang); + sn[row * rot + i] = @sin(ang); + } + }; + const wv = try a.alloc(f64, 2 * f.dim); + for (wv, wb) |*v, h| v.* = bfr.value(h); + const pts: ref.Points = if (f.round) .{ .row = true, .out = true } else .{}; + self.case = .{ .rows = r, .h = shape.n, .hk = hk, .dim = f.dim, .rot = rot, .x = x, .ld = ld, .v_from_k = f.v_from_k, .wq = wv[0..f.dim], .wk = wv[f.dim..], .cos = cs, .sin = sn, .computed = f.rope != .table, .points = pts, .v_norm = f.v_norm }; + self.dx = try self.rig.input(std.mem.sliceAsBytes(x)); + const nw = try self.rig.input(std.mem.sliceAsBytes(wb)); + if (f.rope == .table) { + self.sh = .{ .norms = nw, .norms_bf16 = true, .cos = try self.rig.input(std.mem.sliceAsBytes(cosv)), .sin = try self.rig.input(std.mem.sliceAsBytes(sinv)), .rot = rot, .round = f.round, .v_norm = f.v_norm }; + } else self.sh = .{ .norms = nw, .norms_bf16 = true, .inv = try self.rig.input(std.mem.sliceAsBytes(inv)), .rot = rot, .round = f.round, .v_norm = f.v_norm }; + self.dpos = try self.rig.input(std.mem.sliceAsBytes(pos)); + self.dslot = try self.rig.input(std.mem.sliceAsBytes(slot)); + const kv = r * hk * f.dim * 2; + const cache = 2 * len * hk * f.dim * 2; + self.out = .{ .q = try self.rig.output(r * shape.n * f.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 k = self.case; + const ld: i64 = @intCast(k.ld); + const dim: i64 = @intCast(k.dim); + const kp = self.dx + k.h * k.dim * 2; + const src: qk.Src = .{ .q = self.dx, .q_ld = ld, .q_hs = dim, .k = kp, .k_ld = ld, .k_hs = dim, .v = if (k.v_from_k) kp else kp + k.hk * k.dim * 2, .v_ld = ld, .v_hs = dim }; + try self.k.rowsVia(self.ctl, self.form.keep orelse 0, self.rig.s, src, .{ .h = k.h, .hk = k.hk, .dim = k.dim }, self.sh, self.dpos, self.dslot, self.out, k.rows); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Qk) usize { + return self.case.rows * self.case.heads() * self.case.dim; + } + + /// Every head's float64 point at its rounding points, its sums of squares in `sim` (null: float64). + fn outputs(self: *Qk, sim: ?Sim, want: []f64, mag: []f64) !void { + const k = self.case; + const a = self.arena.allocator(); + const x = try a.alloc(f64, k.dim); + const sq = try a.alloc(f64, k.dim); + const hi = try a.alloc(f64, k.dim); + for (0..k.rows) |row| for (0..k.heads()) |hd| { + for (x, sq, 0..) |*v, *q, i| { + v.* = bfr.value(k.x[k.src(row, hd) + i]); + q.* = v.* * v.*; + } + const ss: ?f64 = if (sim) |s| rows.sumIn(s.acc, s.order, sq) else null; + const at = qr.outAt(k, row, hd); + qr.head(k, row, hd, .{}, ss, x, want[at..][0..k.dim], hi); + for (mag[at..][0..k.dim], want[at..][0..k.dim]) |*m, v| m.* = @abs(v) / 2; // v ~ ss^-1/2 + }; + } + + pub fn reference(self: *Qk, want: []f64, mag: []f64) !void { + return self.outputs(null, want, mag); + } + + pub fn simulate(self: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.arena.allocator().alloc(f64, out.len); + return self.outputs(.{ .acc = acc, .order = order }, out, mag); + } + + /// Every output in float64's range at the rounding points (a dropped point leaves it). + pub fn check(self: *Qk, _: g.Gate) !?[]const u8 { + const n = self.count(); + const out = try self.arena.allocator().alloc(u16, n); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.outputBytes()); + var stats: qr.Stats = .{}; + return qr.verify(self.rig.gpu.gpa, self.case, out, &stats, &self.msg); + } +}; + +test { + _ = qr; +} + +test "every control keeps all but one option of the form it controls" { + for (forms) |f| if (f.keep) |k| { + try std.testing.expectEqual(@as(u32, 1), @popCount((row_pt | out_pt | vnorm_pt) & ~k)); + }; +} diff --git a/zig/tests/cuda/contract_softcap.zig b/zig/tests/cuda/contract_softcap.zig new file mode 100644 index 0000000000..edb2fa23da --- /dev/null +++ b/zig/tests/cuda/contract_softcap.zig @@ -0,0 +1,143 @@ +//! act_rows.cu's tf_softcap_rows under the contract (a soft-capped head's logits), and controls each dropping a point. +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 bfr = core.bf16_round; +const aref = core.act_ref; + +/// tanhf within 2 ulps of its value (2^-22), twice that either way, as act_round_check's spread. +const spread = 0x1p-21; +const cap = 30; +const all = 3; + +/// A registration: the served kernel (keep null), or a control keeping all rounding points but one (act_round.cuh). +pub const Entry = struct { name: []const u8, keep: ?i32 = null }; + +pub const entries = [_]Entry{ + .{ .name = "cuda act tf_softcap_rows (MLX's bf16 steps, cap 30)" }, + .{ .name = "control: soft cap without x / cap's rounding", .keep = all & ~@as(i32, 1) }, + .{ .name = "control: soft cap without tanh's rounding", .keep = all & ~@as(i32, 2) }, +}; + +fn depth(_: c.Shape) usize { + return 0; +} + +pub fn spec(e: Entry) c.Spec { + return .{ + .name = e.name, + .accum = .none, + .depth = depth, + .mids = 16, // gate 3 is the kernel's own check: float64's rounding points, ranged over tanhf's ulps + .out = .bf16, + .tiles = .{ .{ .size = 16 }, .{ .size = 2048, .step = 8, .min = 8 }, .{} }, + // Gemma 4's head over its 262,144-token vocabulary: one row, a verify window, a prompt's last rows + .served = &.{ .{ .rows = 1, .n = 262144, .k = 1 }, .{ .rows = 16, .n = 262144, .k = 1 }, .{ .rows = 3, .n = 4096, .k = 1 } }, + }; +} + +pub const SoftCap = struct { + rig: rig_.Rig, + entry: Entry, + act: *const cuda.act_rows.ActRows, + ctl: ?cuda.Function = null, + arena: ?std.heap.ArenaAllocator = null, + rows: usize = 0, + n: usize = 0, + x: []const u16 = &.{}, + in: u64 = 0, + out: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *SoftCap) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(SoftCap); + + pub fn kernel(self: *SoftCap) g.Kernel { + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = if (self.entry.keep != null) .shapes else null }; + } + + /// Logits spread over [-80, 80) (some past the cap, many near 0); non-finite: NaN and both infinities planted. + pub fn prepare(self: *SoftCap, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n % 8 != 0 or shape.n == 0 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(); + const x = try a.alloc(u16, shape.rows * shape.n); + for (x, 0..) |*v, i| { + const wide: f32 = if (i % 16 == 3) 80 else if (i % 16 == 9) 0.5 else 30; + v.* = bfr.bits((rng.float(f32) * 2 - 1) * wide); + } + if (fill == .nonfinite) { + const planted = [_]struct { usize, f32 }{ .{ 0, std.math.nan(f32) }, .{ shape.n / 2, std.math.inf(f32) }, .{ shape.n - 1, -std.math.inf(f32) } }; + for (planted) |p| x[p[0]] = bfr.bits(p[1]); + } + self.rows = shape.rows; + self.n = shape.n; + self.x = x; + self.in = try self.rig.input(std.mem.sliceAsBytes(x)); + self.out = try self.rig.output(x.len * 2); + return true; + } + + pub fn run(self: *SoftCap) !void { + const s = self.rig.s; + if (self.ctl) |f| { + var args: cuda.Args = .{}; + args.add(self.in); + args.add(self.out); + args.add(@as(c_int, @intCast(self.n))); + args.add(@as(f32, cap)); + args.add(@as(c_int, self.entry.keep.?)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(self.rows), .y = @intCast((self.n / 8 + 255) / 256), .z = 1 }, .block = .{ .x = 256 } }, s, &args); + } else try self.act.softcap(s, self.in, self.out, self.rows, self.n, cap); + try s.synchronize(); + } + + pub fn count(self: *SoftCap) usize { + return self.rows * self.n; + } + + pub fn reference(self: *SoftCap, out: []f64, mag: []f64) !void { + for (self.x, out) |v, *o| o.* = aref.softcap(bfr.value(v), cap, 0); + @memset(mag, 0); + } + + /// Every output within float64's ends at the rounding points (a dropped point leaves them); non-finite as float64. + pub fn check(self: *SoftCap, _: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const out = try a.alloc(u16, self.count()); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.bytes(self.out, out.len * 2)); + for (self.x, out, 0..) |v, got_bits, i| { + const x: f64 = bfr.value(v); + const r = aref.softcapRange(x, cap, spread); + const got: f64 = bfr.value(got_bits); + const ok = if (std.math.isFinite(r[0]) and std.math.isFinite(r[1])) + std.math.isFinite(got) and r[0] <= got and got <= r[1] + else + c.classOf(got) == c.classOf(aref.softcap(x, cap, 0)); + if (!ok) return try std.fmt.bufPrint(&self.msg, "out[{d}] = {e} for x {e} outside float64's [{e}, {e}]", .{ i, got, x, r[0], r[1] }); + } + return null; + } +}; + +test "every control keeps all but one rounding point, each point once" { + var seen: i32 = 0; + for (entries) |e| if (e.keep) |k| { + const dropped = all & ~k; + try std.testing.expectEqual(@as(u32, 1), @popCount(dropped)); + seen |= dropped; + }; + try std.testing.expectEqual(@as(i32, all), seen); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..a4b3bf0d79 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -6,6 +6,7 @@ const nemotron = @import("nemotron"); 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 router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const mx_k = @import("contract_mxfp8.zig"); @@ -23,12 +24,17 @@ const ck_k = @import("contract_ck.zig"); const scan_k = @import("contract_scan.zig"); const lane8_k = @import("contract_ck_lane.zig"); const latent_k = @import("contract_latent.zig"); +const act_round_k = @import("contract_act_round.zig"); const round_k = @import("contract_tree_round.zig"); const act_k = @import("contract_act.zig"); +const softcap_k = @import("contract_softcap.zig"); const gdn_k = @import("contract_gdn.zig"); const qk_k = @import("contract_qk.zig"); +const wrows_k = @import("contract_window_rows.zig"); +const embed_k = @import("contract_embed.zig"); +const qkr_k = @import("contract_qk_round.zig"); const contract = core.contract; const g = contract.gates; @@ -47,6 +53,7 @@ const h3_image = struct { /// The named embedded image; only contract_controls may select the planted controls. pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (@import("contract_deepseek_images.zig").image(name)) |bytes| return bytes; + if (std.mem.eql(u8, name, "window_rows")) return kk.window_rows; if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "nemotron_mamba", "sample", "qwen35", "qwen35_keep", "qwen35_kv8", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; @@ -97,6 +104,11 @@ pub fn run(gpu: check.Gpu) !void { nrows += 1; } for (row_list[0..nrows]) |*r| try kernels.append(a, r.kernel()); + const tail_which = comptime std.enums.values(tail_k.Which); + var tails: [tail_which.len]tail_k.Tail = undefined; + 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 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") }; @@ -256,6 +268,14 @@ pub fn run(gpu: check.Gpu) !void { for (std.enums.values(act_k.Which), &acts) |w, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .k = &act_rows, .arena = .init(a) }; defer for (&acts) |*x| x.deinit(); for (&acts) |*x| try kernels.append(a, x.kernel()); + var act_rounds: [act_round_k.entries.len]act_round_k.Act = undefined; + for (act_round_k.entries, &act_rounds) |e, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .act = &act_rows, .ctl = if (e.keep != null) try mods[5].function("tf_ctl_act_round") else null }; + defer for (&act_rounds) |*x| x.deinit(); + for (&act_rounds) |*x| try kernels.append(a, x.kernel()); + var softcaps: [softcap_k.entries.len]softcap_k.SoftCap = undefined; + for (softcap_k.entries, &softcaps) |e, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .act = &act_rows, .ctl = if (e.keep != null) try mods[5].function("tf_ctl_softcap") else null }; + defer for (&softcaps) |*x| x.deinit(); + for (&softcaps) |*x| try kernels.append(a, x.kernel()); var gdn_prep = try cuda.gdn_prep.GdnPrep.load(gpu.d); defer gdn_prep.unload(); var gdns: [9]gdn_k.Gdn = undefined; @@ -268,6 +288,22 @@ pub fn run(gpu: check.Gpu) !void { for (&qks, qk_k.forms) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .form = f, .k = &qk_rows, .arena = .init(a) }; defer for (&qks) |*x| x.deinit(); for (&qks) |*x| try kernels.append(a, x.kernel()); + var window_fns = try cuda.window_attention.WindowAttention.load(gpu.d); + defer window_fns.unload(); + var wrows: [wrows_k.cfgs.len]wrows_k.Rows = undefined; + for (&wrows, wrows_k.cfgs) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .g = &window_fns, .ctl = mods[5], .arena = .init(a) }; + defer for (&wrows) |*x| x.deinit(); + for (&wrows) |*x| try kernels.append(a, x.kernel()); + var embed_rows = try cuda.embed_rows.EmbedRows.load(gpu.d); + defer embed_rows.unload(); + var embeds: [embed_k.entries.len]embed_k.Embed = undefined; + for (embed_k.entries, &embeds) |e, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .e = &embed_rows, .ctl = if (e.keep != null) try mods[5].function("tf_ctl_embed_affine") else null }; + defer for (&embeds) |*x| x.deinit(); + for (&embeds) |*x| try kernels.append(a, x.kernel()); + var qkrs: [qkr_k.forms.len]qkr_k.Qk = undefined; + for (&qkrs, qkr_k.forms) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .form = f, .k = &qk_rows, .ctl = if (f.keep != null) try mods[5].function("tf_ctl_qk_round") else null, .arena = .init(a) }; + defer for (&qkrs) |*x| x.deinit(); + for (&qkrs) |*x| try kernels.append(a, x.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); diff --git a/zig/tests/cuda/contract_window_rows.zig b/zig/tests/cuda/contract_window_rows.zig new file mode 100644 index 0000000000..53520bc5b1 --- /dev/null +++ b/zig/tests/cuda/contract_window_rows.zig @@ -0,0 +1,307 @@ +//! window_rows.cu under the contract: Gemma 4's sliding (256) and global (512) layers, outputs in float64's range. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const wc = @import("window_rows_check.zig"); +const wt = @import("window_rows_tests.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const ar = core.attention_range; +const bfr = core.bf16_round; +const wa = cuda.window_attention; +const WA = wa.WindowAttention; + +/// A planted fault, from contract_controls.cu (window_rows.cuh's Fault), or the served bf16-P kernel (p16). +pub const Ctl = enum { none, edge, rescale, trunc, prompt_edge, last, p16, rescale512 }; + +/// A call's fixed heads, window and ring, and which kernels it runs. +pub const Cfg = struct { + name: []const u8, + h: usize, + hk: usize, + d: usize = 256, + win: usize = 1024, + ring: usize = 1536, + streams: usize = 3, + prompt: bool = false, + p32: bool = false, + image: bool = false, + ctl: Ctl = .none, + served: []const c.Shape, +}; + +const decode_served = &[_]c.Shape{ .{ .rows = 16, .n = 16, .k = 40000 }, .{ .rows = 64, .n = 16, .k = 3000 }, .{ .rows = 1, .n = 16, .k = 1100 } }; +const prompt_served = &[_]c.Shape{ .{ .rows = 512, .n = 16, .k = 0 }, .{ .rows = 512, .n = 16, .k = 5000 }, .{ .rows = 96, .n = 16, .k = 1300 } }; +const global_decode = &[_]c.Shape{ .{ .rows = 4, .n = 16, .k = 32768 }, .{ .rows = 16, .n = 16, .k = 3000 }, .{ .rows = 1, .n = 16, .k = 300 } }; +const global_prompt = &[_]c.Shape{ .{ .rows = 256, .n = 16, .k = 0 }, .{ .rows = 128, .n = 16, .k = 4096 }, .{ .rows = 32, .n = 16, .k = 1000 } }; +const global31_prompt = &[_]c.Shape{ .{ .rows = 128, .n = 32, .k = 0 }, .{ .rows = 64, .n = 32, .k = 3000 } }; + +/// Gemma 4 12B's sliding layers (16 query heads over 8 of 256, window 1024, ring 1536); controls each plant one fault. +pub const cfgs = [_]Cfg{ + .{ .name = "cuda window_rows decode tf_wrows_part_256 + tf_wrows_merge_256, 16/8 heads, window 1024, ring 1536", .h = 16, .hk = 8, .served = decode_served }, + .{ .name = "cuda window_rows prompt tf_wrows_prompt_256_p32, 16/8 heads, window 1024, ring 1536, image rows", .h = 16, .hk = 8, .prompt = true, .p32 = true, .image = true, .served = prompt_served }, + .{ .name = "cuda window_rows prompt tf_wrows_prompt_256 (P as bf16), 16/8 heads, window 1024, ring 1536", .h = 16, .hk = 8, .prompt = true, .served = prompt_served }, + .{ .name = "control: decode partials whose window starts a key early", .h = 16, .hk = 8, .ctl = .edge, .served = decode_served }, + .{ .name = "control: decode merge without each chunk's exp(m_c - M)", .h = 16, .hk = 8, .ctl = .rescale, .served = decode_served }, + .{ .name = "control: decode merge rounding its output toward zero", .h = 16, .hk = 8, .ctl = .trunc, .served = decode_served }, + .{ .name = "control: P32 prompt rows whose window starts a key early", .h = 16, .hk = 8, .prompt = true, .p32 = true, .ctl = .prompt_edge, .served = prompt_served }, + .{ .name = "control: P32 prompt rows ignoring an image's last key", .h = 16, .hk = 8, .prompt = true, .p32 = true, .image = true, .ctl = .last, .served = prompt_served }, + .{ .name = "control: P rounded to bf16 under the P32 range", .h = 16, .hk = 8, .prompt = true, .p32 = true, .ctl = .p16, .served = prompt_served }, + // Gemma 4's global layers (head dim 512, no window over a linear cache): 12B 16 query heads over 1, 31B 32 over 4 + .{ .name = "cuda window_rows decode tf_wrows_part_512 + tf_wrows_merge_512, 16/1 heads, no window, linear 32768", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 32768, .streams = 1, .served = global_decode }, + .{ .name = "cuda window_rows prompt tf_wrows_prompt_512_p32, 16/1 heads, no window, linear cache", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 0, .prompt = true, .p32 = true, .served = global_prompt }, + .{ .name = "cuda window_rows prompt tf_wrows_prompt_512 (P as bf16), 32/4 heads, no window, linear cache", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = 0, .prompt = true, .served = global31_prompt }, + .{ .name = "control: head dim 512 merge over up to 128 chunks without exp(m_c - M)", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 32768, .streams = 1, .ctl = .rescale512, .served = global_decode }, +}; + +/// The control kernel's symbol in contract_controls.cu, if it has one. +pub fn symbol(ctl: Ctl) ?[:0]const u8 { + return switch (ctl) { + .none, .p16 => null, + .edge => "tf_ctl_wrows_part_edge", + .rescale => "tf_ctl_wrows_merge_rescale", + .rescale512 => "tf_ctl_wrows_merge_rescale_512", + .trunc => "tf_ctl_wrows_merge_trunc", + .prompt_edge => "tf_ctl_wrows_prompt_edge", + .last => "tf_ctl_wrows_prompt_last", + }; +} + +/// The longest chain: a chunk's keys and the chunks (decode), or the tiles' folds (prompt), as the range counts them. +fn depth(s: c.Shape) usize { + return 256 + s.k / 32 + 16; +} + +const why = "all keys alike make every weight exactly 1: o is a plain sum of V's integers, exact in fp32 (any order); fp16 sums overflow, bf16 ones lose its low bits"; + +pub fn spec(f: Cfg) c.Spec { + const tiles: [3]c.Tile = if (f.prompt) .{ .{ .size = 16 }, .{}, .{ .size = 32, .min = 0 } } else .{ .{}, .{}, .{ .size = 256 } }; + return .{ .name = f.name, .accum = .f32, .depth = depth, .ops = 4, .out = .bf16, .tiles = tiles, .served = f.served, .sums = .{ .f16 = 100, .bf16 = 0, .why = why } }; +} + +pub const Rows = struct { + rig: rig_.Rig, + f: Cfg, + g: *const WA, + /// The controls' image (contract_controls.cu), when this registration runs one. + ctl: ?cuda.Module = null, + arena: std.heap.ArenaAllocator, + call: wc.Call = undefined, + fns: WA.Wide = undefined, + dev: [6]u64 = @splat(0), + scratch: u64 = 0, + out: u64 = 0, + half: []f64 = &.{}, + msg: [240]u8 = undefined, + + pub fn deinit(self: *Rows) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Rows) g.Kernel { + return .{ .spec = spec(self.f), .ctx = self, .vt = &vt, .control = if (self.f.ctl != .none) .shapes else null }; + } + + const vt = g.hooks(Rows); + + fn win(self: *const Rows) wa.Window { + const f = self.f; + return .{ .h = f.h, .hk = f.hk, .d = f.d, .win = f.win, .ring = f.ring, .scale = 1, .p32 = f.p32 and f.ctl != .p16 }; + } + + /// The kernels a run launches: the served ones, one of them swapped for its control. + fn kernels(self: *Rows) !WA.Wide { + var k = self.g.wideFns(self.f.d).*; + const sym = symbol(self.f.ctl) orelse return k; + const fun = try self.ctl.?.function(sym); + switch (self.f.ctl) { + .edge => k.part = fun, + .rescale, .trunc, .rescale512 => k.merge = fun, + else => { + try fun.allowDynamicShared(WA.promptShared(self.f.d)); + k.prompt[1] = fun; + }, + } + return k; + } + + pub fn prepare(self: *Rows, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const f = self.f; + if (shape.n != f.h or shape.rows == 0 or (!f.prompt and shape.k == 0)) return false; + if (f.prompt and f.ring > 0 and shape.rows + f.win > f.ring) return false; // chunk + window fit the ring + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const n = shape.rows; + // gate 4 without a window: at most 1,000 keys, so the mirrored integers' partial sums stay exact in fp32 + const at = if (fill == .stress and f.win == 0) @min(shape.k, 1000) else shape.k; + const crow = if (f.prompt) (if (f.ring > 0) f.ring else at + n) else f.streams * f.ring; + const q = try a.alloc(u16, n * f.h * f.d); + const kc = try a.alloc(u16, crow * f.hk * f.d); + const vc = try a.alloc(u16, kc.len); + wt.fill(rng, q, kc, vc, f.d); + const pos = try a.alloc(i32, n); + const slot = try a.alloc(i32, n); + const hi = try a.alloc(i32, n); + for (pos, slot, hi, 0..) |*p, *sl, *x, i| { + // short histories, wrapped rings, the last position; gate 4: one stream's row at the bound, every row alike + p.* = @intCast(if (fill == .stress or i == n - 1) at -| 1 else if (i % 4 == 0) i % @min(40, shape.k) else (i * 7919) % shape.k); + sl.* = @intCast(if (fill == .stress) 0 else i % f.streams); + x.* = @intCast(at + i); + } + if (f.image and fill != .stress and n > 8) { // an image: rows n/8 .. n/2 see each other + for (n / 8..n / 2) |i| hi[i] = @intCast(shape.k + n / 2 - 1); + } + self.call = .{ .h = f.h, .hk = f.hk, .d = f.d, .win = f.win, .ring = f.ring, .scale = 1, .prompt = f.prompt, .p32 = f.p32, .rows = n, .q = q, .kc = kc, .vc = vc, .pos = pos, .slot = slot, .p0 = at, .hi = if (f.image) hi else null }; + if (fill == .random) { + var r: usize = 0; + while (r < n) : (r += 3) { + const w = self.call.at(r); + const lo = self.call.first(w.p); + for (0..f.hk) |kh| wt.sharpen(self.call, kc, r, lo + (r * 31 + kh * 7) % (w.last + 1 - lo), kh); + } + } + switch (fill) { + .random => {}, + .nonfinite => self.plant(q, kc, vc), + .stress => try self.stress(rng, kc, vc), + } + for ([_][]const u8{ std.mem.sliceAsBytes(q), std.mem.sliceAsBytes(kc), std.mem.sliceAsBytes(vc), std.mem.sliceAsBytes(pos), std.mem.sliceAsBytes(slot), std.mem.sliceAsBytes(hi) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + self.scratch = if (f.prompt) 0 else try self.rig.scratch(WA.scratchBytes(self.win(), n)); + self.out = try self.rig.output(n * f.h * f.d * 2); + self.fns = try self.kernels(); + self.half = try a.alloc(f64, self.count()); + return true; + } + + fn cell(self: *Rows, x: []u16, r: usize, key: usize, kh: usize) []u16 { + const w = self.call.at(r); + return x[(self.call.cacheRow(w.base, key) * self.f.hk + kh) * self.f.d ..][0..self.f.d]; + } + + /// NaN in q and a key; +inf in a later row's value; +-inf in a column; +inf at float64 weight e^-200 (row 4). + fn plant(self: *Rows, q: []u16, kc: []u16, vc: []u16) void { + const f = self.f; + const n = self.call.rows; + const nan = bfr.bits(std.math.nan(f32)); + const inf = bfr.bits(std.math.inf(f32)); + q[3] = nan; // row 0, head 0 + self.cell(vc, 1, self.call.at(1).p, 1 % f.hk)[7] = inf; + if (n > 2) { + const p = self.call.at(2).p; + self.cell(vc, 2, p, 2 % f.hk)[11] = inf; + if (p > self.call.first(p)) self.cell(vc, 2, p - 1, 2 % f.hk)[11] = bfr.bits(-std.math.inf(f32)); + } + if (n > 3) self.cell(kc, 3, self.call.at(3).p, 3 % f.hk)[0] = nan; + if (n > 4) { + const r = 4; + const kh = 4 % f.hk; + const w = self.call.at(r); + const grp = f.h / f.hk; + for (0..grp) |hh| @memset(q[(r * f.h + kh * grp + hh) * f.d ..][0..f.d], bfr.bits(0.625)); + for (self.call.first(w.p)..w.last + 1) |key| { + const k = self.cell(kc, r, key, kh); + @memset(k, bfr.bits(if (key == w.p) 0.625 else -0.625)); + if (key != w.p) self.cell(vc, r, key, kh)[21] = inf; + } + } + } + + /// Gate 4: every key alike (each weight 1); values mirrored over the keys every row sees, 0 elsewhere. + fn stress(self: *Rows, rng: std.Random, kc: []u16, vc: []u16) !void { + const a = self.arena.allocator(); + const f = self.f; + const kv = f.hk * f.d; + for (0..kc.len / kv) |i| @memcpy(kc[i * kv ..][0..kv], kc[0..kv]); + @memset(vc, 0); + const n = self.call.rows; + const lo = self.call.first(self.call.at(n - 1).p); + const hi = self.call.at(0).last; + const col = try a.alloc(f32, hi + 1 - lo); + for (0..f.hk) |kh| for (0..f.d) |e| { + mirrored(rng, col); + for (col, lo..) |y, k| self.cell(vc, 0, k, kh)[e] = bfr.bits(y); + }; + } + + pub fn run(self: *Rows) !void { + const n = self.call.rows; + const w = self.win(); + if (self.f.prompt) { + try WA.promptWith(&self.fns, self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.out, self.call.p0, n, w, if (self.f.image) self.dev[5] else null); + } else try WA.decodeWith(&self.fns, self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.out, self.dev[3], self.dev[4], w, self.scratch, n); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Rows) usize { + return self.call.rows * self.f.h * self.f.d; + } + + /// Every output in its float64 range (gates 1 and 3); at gate 1 some, not all, of float64's outputs non-finite. + pub fn check(self: *Rows, at: g.Gate) !?[]const u8 { + const a = self.arena.allocator(); + const out = try a.alloc(u16, self.count()); + @memcpy(std.mem.sliceAsBytes(out), try self.rig.outputBytes()); + var st: ar.Stats = .{}; + if (try wc.verify(self.rig.gpu.gpa, self.call, out, &st, &self.msg)) |why_| return why_; + if (at == .nonfinite and (st.nonfinite == 0 or st.nonfinite == out.len)) return try std.fmt.bufPrint(&self.msg, "the planted values make {d} of {d} reference outputs non-finite: no proof", .{ st.nonfinite, out.len }); + return null; + } + + pub fn reference(self: *Rows, want: []f64, mag: []f64) !void { + try wc.reference(self.rig.gpu.gpa, self.call, want, self.half); + for (mag, want) |*m, w| m.* = @abs(w); + } + + /// The range's half width, then the output's own rounding at its far end. + pub fn bounds(self: *Rows, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, self.half) |*l, w, h| l.* = h + 0x1p-8 * (@abs(w) + h) + 0x1p-133; + } + + /// Gate 4's sums in an accumulator and order: every weight is 1, so o is V's column sum over l's count of ones. + pub fn simulate(self: *Rows, acc: c.Accum, order: g.Order, out: []f64) !void { + const a = self.arena.allocator(); + const call = self.call; + const d = call.d; + const grp = call.h / call.hk; + for (0..call.rows) |r| { + const w = call.at(r); + const lo = call.first(w.p); + const n = w.last + 1 - lo; + const terms = try a.alloc(f64, n); + const ones = try a.alloc(f64, n); + @memset(ones, 1); + const l = rows.sumIn(acc, order, ones); + for (0..call.hk) |kh| for (0..d) |e| { + for (terms, lo..) |*t, key| t.* = bfr.value(call.vc[(call.cacheRow(w.base, key) * call.hk + kh) * d + e]); + const o = rows.sumIn(acc, order, terms) / l; + for (0..grp) |hh| out[(r * call.h + kh * grp + hh) * d + e] = o; + }; + } + } +}; + +/// Gate 4's values: odd integers, every other one times 64 (bf16-exact), the second half cancelling the first. +fn mirrored(rng: std.Random, out: []f32) void { + const h = out.len / 2; + for (0..h) |i| { + const v: f32 = @floatFromInt((1 + 2 * rng.uintLessThan(u32, 128)) * @as(u32, if (i % 2 == 1) 64 else 1)); + out[i] = v; + out[h + i] = (if (i < 8) @as(f32, 1) else 0) - v; + } + if (out.len % 2 == 1) out[out.len - 1] = 0; +} + +test "every control names a kernel of contract_controls.cu, or runs the served bf16-P kernel" { + for (cfgs) |f| if (f.ctl != .none) { + try std.testing.expect(symbol(f.ctl) != null or f.ctl == .p16); + try std.testing.expect(std.mem.startsWith(u8, f.name, "control: ")); + }; +} diff --git a/zig/tests/cuda/embed_check.zig b/zig/tests/cuda/embed_check.zig new file mode 100644 index 0000000000..22eeaf7876 --- /dev/null +++ b/zig/tests/cuda/embed_check.zig @@ -0,0 +1,66 @@ +//! MLX affine token rows against core/embed_ref.zig's float64 value at their rounding points, bit for bit. +const std = @import("std"); +const core = @import("core"); + +const eref = core.embed_ref; +const bfr = core.bf16_round; + +/// One call's host inputs: the table's words, scales and biases as bf16 bits, the ids, and the form. +pub const Case = struct { + rows: usize, + d: usize, + bits: u4, + group: usize, + scale: ?f32, + ids: []const i32, + w: []const u32, + s: []const u16, + b: []const u16, +}; + +/// What a check saw: outputs bit-equal to float64's value, and NaN where float64's is NaN. +pub const Stats = struct { equal: usize = 0, nan: usize = 0 }; + +/// Element i of row `row`: float64's value at the rounding points. +pub fn want(k: Case, row: usize, i: usize) f64 { + const t: usize = @intCast(k.ids[row]); + const words = k.d / 32 * k.bits; + const q = eref.code(k.w[t * words ..][0..words], i, k.bits); + const g = t * (k.d / k.group) + i / k.group; + return eref.affine(q, bfr.value(k.s[g]), bfr.value(k.b[g]), if (k.scale) |x| x else null); +} + +/// Every output's bits equal float64's value rounded at the points (NaN where it is NaN, any payload). +pub fn verify(k: Case, out: []const u16, stats: *Stats, msg: []u8) !?[]const u8 { + for (0..k.rows) |row| for (0..k.d) |i| { + const w = want(k, row, i); + const got = out[row * k.d + i]; + const ok = if (std.math.isNan(w)) std.math.isNan(bfr.value(got)) else got == bfr.bits(@floatCast(w)); + if (!ok) return try std.fmt.bufPrint(msg, "row {d} out[{d}] = {e} where float64 at the rounding points gives {e}", .{ row, i, bfr.value(got), w }); + if (std.math.isNan(w)) stats.nan += 1 else stats.equal += 1; + }; + return null; +} + +/// MLX's words of `codes` `bits` wide, low code first. +pub fn pack(a: std.mem.Allocator, codes: []const u32, bits: u4) ![]u32 { + const w = try a.alloc(u32, codes.len * bits / 32); + @memset(w, 0); + for (codes, 0..) |q, i| { + const bit = i * bits; + w[bit / 32] |= q << @intCast(bit % 32); + if (bit % 32 + bits > 32) w[bit / 32 + 1] |= q >> @intCast(32 - bit % 32); + } + return w; +} + +test "pack is embed_ref.code's inverse" { + var prng = std.Random.DefaultPrng.init(4); + var codes: [192]u32 = undefined; + for ([_]u4{ 2, 3, 5, 6, 8 }) |bits| { + for (&codes) |*q| q.* = prng.random().int(u32) & ((@as(u32, 1) << bits) - 1); + const w = try pack(std.testing.allocator, &codes, bits); + defer std.testing.allocator.free(w); + for (codes, 0..) |q, i| try std.testing.expectEqual(q, eref.code(w, i, bits)); + } +} diff --git a/zig/tests/cuda/embed_tests.zig b/zig/tests/cuda/embed_tests.zig new file mode 100644 index 0000000000..300554f80b --- /dev/null +++ b/zig/tests/cuda/embed_tests.zig @@ -0,0 +1,88 @@ +//! embed_rows.cu's MLX affine rows (2 to 6 or 8 bits, a bf16 scale) against float64's rounding points, bit for bit. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rows = @import("norm_rows_tests.zig"); +const ec = @import("embed_check.zig"); + +const bfr = core.bf16_round; +const counts = rows.counts; +const most = rows.most; +const Rig = rows.Rig; +const vocab = 300; + +/// sqrt(d) as MLX multiplies by it: its bf16 value. +pub fn scaleOf(d: usize) f32 { + return bfr.round32(@sqrt(@as(f32, @floatFromInt(d)))); +} + +/// A table of `bits` wide codes: scales and biases in the checkpoint's ranges (-0.18 to 0.32, -4.1 to 2.7). +fn table(r: *Rig, d: usize, bits: u4, group: usize) !struct { w: []u32, s: []u16, b: []u16 } { + const codes = try r.a.alloc(u32, vocab * d); + for (codes) |*q| q.* = r.rng.int(u32) & ((@as(u32, 1) << bits) - 1); + const s = try r.a.alloc(u16, vocab * d / group); + const b = try r.a.alloc(u16, vocab * d / group); + for (s, b) |*x, *y| { + x.* = bfr.bits(r.rng.float(f32) * 0.5 - 0.18); + y.* = bfr.bits(r.rng.float(f32) * 6.8 - 4.1); + } + return .{ .w = try ec.pack(r.a, codes, bits), .s = s, .b = b }; +} + +/// One form at every row count: the 64-row call bit-equal to float64, the smaller calls its leading rows. +fn form(e: *const cuda.embed_rows.EmbedRows, r: *Rig, d: usize, bits: u4, group: usize, scale: ?f32, stats: *ec.Stats) !void { + const t = try table(r, d, bits, group); + const ids = try r.a.alloc(i32, most); + for (ids, 0..) |*x, i| x.* = @intCast(if (i == 0) vocab - 1 else r.rng.uintLessThan(usize, vocab)); + const did = try r.dev(std.mem.sliceAsBytes(ids)); + const dw = try r.dev(std.mem.sliceAsBytes(t.w)); + const ds = try r.dev(std.mem.sliceAsBytes(t.s)); + const db = try r.dev(std.mem.sliceAsBytes(t.b)); + var os: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * d * 2); + try e.affine(r.s, did, dw, ds, db, os[c], m, d, .{ .bits = bits, .group = group, .scale = scale }); + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "embed affine {d}-bit, groups of {d}, {d} wide, scale {?d}", .{ bits, group, d, scale }); + const o = try r.back(u16, os[counts.len - 1], most * d); + var msg: [200]u8 = undefined; + const k: ec.Case = .{ .rows = most, .d = d, .bits = bits, .group = group, .scale = scale, .ids = ids, .w = t.w, .s = t.s, .b = t.b }; + if (try ec.verify(k, o, stats, &msg)) |why| try check.expect(false, "{s}: {s}", .{ label, why }); + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * d)), std.mem.sliceAsBytes(o[0 .. m * d])); + if (bits == 4 and group != 128 and scale == null) { // the new kernel at scale 1 gives tf_embed_rows_affine4's bytes + const one = try r.out(most * d * 2); + try e.affine(r.s, did, dw, ds, db, one, most, d, .{ .bits = 4, .group = group, .scale = 1 }); + try check.sameBytes("embed affine at scale 1 and affine4", std.mem.sliceAsBytes(try r.back(u16, one, most * d)), std.mem.sliceAsBytes(o)); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(31); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var e = try cuda.embed_rows.EmbedRows.load(gpu.d); + defer e.unload(); + var stats: ec.Stats = .{}; + for ([_]u4{ 2, 3, 4, 5, 6, 8 }) |bits| { + for ([_]usize{ 3840, 5376 }) |d| try form(&e, &r, d, bits, 64, scaleOf(d), &stats); // Gemma 4 12B, 31B + try form(&e, &r, 3840, bits, 64, null, &stats); + for ([_]usize{ 64, 1984, 2048, 2112 }) |d| try form(&e, &r, d, bits, 64, scaleOf(3840), &stats); // a block's edges + try form(&e, &r, 1024, bits, 32, scaleOf(1024), &stats); + try form(&e, &r, 1024, bits, 128, null, &stats); + } + try form(&e, &r, 2560, 4, 32, null, &stats); // affine4's group of 32 + const z: u64 = 0; + for ([_]cuda.embed_rows.EmbedRows.Affine{ .{ .bits = 7 }, .{ .bits = 1 }, .{ .group = 48 }, .{ .scale = @sqrt(@as(f32, 3840)) }, .{ .scale = std.math.inf(f32) } }) |o| + try std.testing.expectError(error.RowShape, e.affine(r.s, z, z, z, z, z, 1, 3840, o)); + try std.testing.expectError(error.RowShape, e.affine(r.s, z, z, z, z, 8, 1, 3840, .{ .bits = 8, .scale = 62 })); + check.pass("embed rows MLX affine: {d} outputs bit-equal to float64's rounding points (2 to 6 and 8 bits, groups of 32, 64 and 128, scaled and not); every row count's bits equal; scale 1 gives affine4's bytes", .{stats.equal}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..85a4c7de53 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -46,6 +46,8 @@ const exlr8_tests = @import("exlr8_tests.zig"); const exlr8_dense_tests = @import("exlr8_dense_tests.zig"); const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); +const embed_tests = @import("embed_tests.zig"); +const norm_tail_tests = @import("norm_tail_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const contract_tests = @import("contract_tests.zig"); @@ -62,6 +64,8 @@ const act_rows_tests = @import("act_rows_tests.zig"); const gdn_prep_tests = @import("gdn_prep_tests.zig"); const qk_rows_tests = @import("qk_rows_tests.zig"); const window_tests = @import("window_tests.zig"); +const window_rows_tests = @import("window_rows_tests.zig"); +const window_bench = @import("window_bench.zig"); const moe_glue_tests = @import("moe_glue_tests.zig"); const dense_rows_tests = @import("dense_rows_tests.zig"); @@ -126,6 +130,8 @@ const usage = \\ triton a Triton cubin's bits against the Python oracle's fixture \\ norms RMS norm and the residual step against float64 references (synthetic) \\ norm-rows bf16-residual and gated norms, token rows against float64 (synthetic) + \\ embed-rows MLX affine token rows (2 to 8 bits, bf16 scale) against float64 (synthetic) + \\ norm-tail bf16 residual tails (branch norm, scalar, MLX rounding) against float64 (synthetic) \\ tree-round tree_round.cu's picks, batched verified trees and commits against core/tree_round \\ act-rows SwiGLU, output gates, group sums, DFlash2's dynamic conv against float64 (synthetic) \\ gdn-prep Gated DeltaNet conv, norms, g and beta against float64 (synthetic) @@ -181,6 +187,8 @@ const usage = \\ tree tree attention (committed keys and paths) against float64 references (synthetic) \\ block DFlash block attention over windowed contexts against float64 (synthetic) \\ window sliding-window attention against float64 references (synthetic) + \\ window-rows window attention at head dims 256 and 512 against float64 ranges (synthetic) + \\ window-bench WHICH window attention's GPU time: callers (Kolibri-1's shapes) or gemma (head dims 256, 512) \\ moe-glue sigmoid top-k routing and the weighted combine against float64 references (synthetic) \\ dense-rows bf16 rows times weight rows (routers, heads) against float64 references (synthetic) \\ exl3 [dir] EXL3 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) @@ -307,6 +315,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "q35-argmax")) return sample_tests.q35Argmax(gpu); if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "embed-rows")) return embed_tests.run(gpu); + if (std.mem.eql(u8, cmd, "norm-tail")) return norm_tail_tests.run(gpu); if (std.mem.eql(u8, cmd, "tree-round")) return tree_round_tests.run(gpu); if (std.mem.eql(u8, cmd, "act-rows")) return act_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "gdn-prep")) return gdn_prep_tests.run(gpu); @@ -361,6 +371,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu); if (std.mem.eql(u8, cmd, "block")) return block_tests.run(gpu); if (std.mem.eql(u8, cmd, "window")) return window_tests.run(gpu); + if (std.mem.eql(u8, cmd, "window-rows")) return window_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "window-bench")) return window_bench.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "moe-glue")) return moe_glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "dense-rows")) return dense_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "train-pack")) return train_pack.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24); diff --git a/zig/tests/cuda/norm_tail_tests.zig b/zig/tests/cuda/norm_tail_tests.zig new file mode 100644 index 0000000000..deed776a51 --- /dev/null +++ b/zig/tests/cuda/norm_tail_tests.zig @@ -0,0 +1,112 @@ +//! norms.cu's bf16 residual tails at Gemma 4's widths and the 16-byte tiles' edges: float64 ranges, exact sums, bits. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rows = @import("norm_rows_tests.zig"); +const tail_ = @import("norms_tail_check.zig"); + +const bfr = core.bf16_round; +const counts = rows.counts; +const most = rows.most; +const eps: f32 = 1e-6; + +/// A tail's options as a family calls them. +const Form = struct { name: []const u8, round: bool, residual: bool, branch: bool, scalar: bool }; + +const input_norm: Form = .{ .name = "input norm (MLX)", .round = true, .residual = false, .branch = false, .scalar = false }; +const attn_tail: Form = .{ .name = "attention tail (MLX)", .round = true, .residual = true, .branch = true, .scalar = false }; +const dense_tail: Form = .{ .name = "dense tail with layer scalar (MLX)", .round = true, .residual = true, .branch = true, .scalar = true }; +const forms = [_]Form{ + input_norm, + attn_tail, + dense_tail, + .{ .name = "scalar alone (MLX)", .round = true, .residual = false, .branch = false, .scalar = true }, + .{ .name = "dense tail, rounded once", .round = false, .residual = true, .branch = true, .scalar = true }, + .{ .name = "branch norm alone, rounded once", .round = false, .residual = false, .branch = true, .scalar = false }, +}; + +fn case(r: *rows.Rig, n: *const cuda.norms.Norms, width: usize, f: Form, stats: *tail_.Stats) !void { + const x = r.bfs(most * width, 1); + const res = r.bfs(most * width, 4); + const wb = try r.a.alloc(u16, width); + const w = try r.a.alloc(u16, width); + for (wb, w, 0..) |*p, *q, i| { // signed, as Gemma 4's norm weights (12B: -143 to 193), one in 97 large + const big: f32 = if (i % 97 == 3) 30 else 1; + p.* = rows.bf16((r.rng.float(f32) * 8 - 2) * big); + q.* = rows.bf16((r.rng.float(f32) * 8 - 2) * big); + } + const sc = [8]u16{ bfr.bits(0.004 + 0.9 * r.rng.float(f32)), 0, 0, 0, 0, 0, 0, 0 }; // one bf16, 16 bytes + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dr = try r.dev(std.mem.sliceAsBytes(res)); + const dwb = try r.dev(std.mem.sliceAsBytes(wb)); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const dsc = try r.dev(std.mem.sliceAsBytes(&sc)); + var hs: [counts.len]u64 = undefined; + var ys: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + hs[c] = try r.out(m * width * 2); + ys[c] = try r.out(m * width * 2); + xss[c] = try r.out(m * width / 64 * 4); + const t: cuda.norms.Tail = .{ .round = f.round, .branch = if (f.branch) dwb else null, .scalar = if (f.scalar) dsc else null }; + try n.addRmsTail(r.s, dx, if (f.residual) dr else null, dw, hs[c], ys[c], xss[c], m, width, eps, t); + } + var name: [128]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "norm tail {s}, width {d}", .{ f.name, width }); + const h = try r.back(u16, hs[counts.len - 1], most * width); + const y = try r.back(u16, ys[counts.len - 1], most * width); + const xs = try r.back(f32, xss[counts.len - 1], most * width / 64); + const k: tail_.Case = .{ .rows = most, .width = width, .round = f.round, .x = x, .w = w, .r = if (f.residual) res else null, .wb = if (f.branch) wb else null, .sc = if (f.scalar) sc[0] else null }; + var msg: [200]u8 = undefined; + if (try tail_.verify(r.gpu.gpa, k, h, y, xs, stats, &msg)) |why| try check.expect(false, "{s}: {s}", .{ label, why }); + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, ys[c], m * width)), std.mem.sliceAsBytes(y[0 .. m * width])); + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(f32, xss[c], m * width / 64)), std.mem.sliceAsBytes(xs[0 .. m * width / 64])); + if (k.writesH()) try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, hs[c], m * width)), std.mem.sliceAsBytes(h[0 .. m * width])); + } +} + +/// No option set is addRmsH itself: the same launch, the same bytes. +fn unchanged(r: *rows.Rig, n: *const cuda.norms.Norms, width: usize) !void { + const x = try r.dev(std.mem.sliceAsBytes(r.bfs(most * width, 1))); + const res = try r.dev(std.mem.sliceAsBytes(r.bfs(most * width, 4))); + const wv = try r.a.alloc(u16, width); + for (wv, r.weights(width)) |*o, v| o.* = rows.bf16(v); + const w = try r.dev(std.mem.sliceAsBytes(wv)); + var outs: [2][3]u64 = undefined; + for (&outs) |*o| for (o, [_]usize{ 2, 2, 4 }) |*p, b| { + p.* = try r.out(most * width * b / @as(usize, if (b == 4) 64 else 1)); + }; + try n.addRmsH(r.s, x, res, w, .bf16, outs[0][0], outs[0][1], outs[0][2], most, width, eps); + try n.addRmsTail(r.s, x, res, w, outs[1][0], outs[1][1], outs[1][2], most, width, eps, .{}); + for (outs[0], outs[1], [_]usize{ 2, 2, 4 }) |p, q, b| { + const len = most * width / @as(usize, if (b == 4) 64 else 1); + try check.sameBytes("norm tail without options", std.mem.sliceAsBytes(try r.back(u16, q, len * b / 2)), std.mem.sliceAsBytes(try r.back(u16, p, len * b / 2))); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(4); + var r: rows.Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var n = try cuda.norms.Norms.load(gpu.d); + defer n.unload(); + var stats: tail_.Stats = .{}; + // 12B, 31B, the 12B's assistant + for ([_]usize{ 3840, 5376, 1024 }) |width| for (forms) |f| try case(&r, &n, width, f, &stats); + for ([_]usize{ 64, 1984, 2048, 2112, 8192 }) |width| for ([_]Form{ input_norm, dense_tail }) |f| try case(&r, &n, width, f, &stats); + try unchanged(&r, &n, 3840); + const sc = try r.out(16); + const buf = try r.out(2 * 3904); + try std.testing.expectError(error.RowShape, n.addRmsTail(r.s, buf, null, buf, buf, buf, null, 1, 3848, eps, .{ .round = true })); + try std.testing.expectError(error.RowShape, n.addRmsTail(r.s, buf, null, buf, buf, buf, null, 1, 3840, eps, .{ .scalar = sc + 1 })); + try std.testing.expectError(error.RowShape, n.addRmsTail(r.s, buf, null, buf, buf, buf, null, 1, 8256, eps, .{ .round = true })); + check.pass("norm tail: {d} outputs bit-equal to float64's rounding points, {d} within an fp32 inverse's range, {d} non-finite as float64; {d} group sums equal to the fp32 tree; every row count's bits equal; no options = addRmsH's bytes; bad rows refused", .{ stats.decided, stats.ranged, stats.nonfinite, stats.sums }); +} diff --git a/zig/tests/cuda/norms_tail_check.zig b/zig/tests/cuda/norms_tail_check.zig new file mode 100644 index 0000000000..61021ff303 --- /dev/null +++ b/zig/tests/cuda/norms_tail_check.zig @@ -0,0 +1,192 @@ +//! A norm tail's outputs against core/norms_ref.zig's float64 ranges, its 64-value sums against the kernel's fp32 tree. +const std = @import("std"); +const core = @import("core"); + +const ref = core.norms_ref; +const bfr = core.bf16_round; +const c = core.contract.spec; + +/// The vector norm's sum chain: 32 fma a thread over 16-byte loads, then 256 threads' tree (contract_rows' vecDepth). +pub const depth = 42; + +/// An fp32 inverse RMS's relative spread from float64's: the sum's gamma halved by the root, fp32 steps, the products. +pub fn spread() f64 { + return c.gamma(depth, 0x1p-24) / 2 + 8 * 0x1p-24; +} + +/// One call's host inputs as bf16 bits: rows of `width`, optional residual, branch weight and scalar. +pub const Case = struct { + rows: usize, + width: usize, + round: bool, + x: []const u16, + w: []const u16, + r: ?[]const u16 = null, + wb: ?[]const u16 = null, + sc: ?u16 = null, + + /// The kernel writes h when the tail has a residual, a branch norm or a scalar. + pub fn writesH(k: Case) bool { + return k.r != null or k.wb != null or k.sc != null; + } +}; + +/// What a check saw: outputs whose float64 range is one value, a range, and outputs non-finite in float64. +pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0, sums: usize = 0 }; + +fn vals(a: std.mem.Allocator, b: []const u16) ![]f64 { + const o = try a.alloc(f64, b.len); + for (o, b) |*v, h| v.* = bfr.value(h); + return o; +} + +/// The float64 point reference of a row: h and y (spread 0). +pub fn point(a: std.mem.Allocator, k: Case, row: usize, h: []f64, y: []f64) !void { + return rowRef(a, k, row, 0, h, try a.alloc(f64, k.width), y, try a.alloc(f64, k.width)); +} + +/// Row `row` of `k` as core/norms_ref.zig's tail and its float64 x and w, the inverse's spread `spread_`. +pub fn tailOf(a: std.mem.Allocator, k: Case, row: usize, spread_: f64) !struct { t: ref.Tail, x: []f64, w: []f64 } { + const at = row * k.width; + return .{ + .t = .{ + .round = k.round, + .residual = if (k.r) |r| try vals(a, r[at..][0..k.width]) else null, + .branch = if (k.wb) |wb| try vals(a, wb) else null, + .scalar = if (k.sc) |s| bfr.value(s) else null, + .spread = spread_, + }, + .x = try vals(a, k.x[at..][0..k.width]), + .w = try vals(a, k.w), + }; +} + +fn rowRef(a: std.mem.Allocator, k: Case, row: usize, spread_: f64, h_lo: []f64, h_hi: []f64, y_lo: []f64, y_hi: []f64) !void { + const t = try tailOf(a, k, row, spread_); + ref.tail(t.t, t.x, t.w, ref.Exact{}, h_lo, h_hi, y_lo, y_hi); +} + +/// got in [lo, hi]: by its bits where the ends are one value, else between ordered ends (-0 below +0); NaN for NaN. +fn inRange(got: u16, lo: f64, hi: f64, stats: *Stats) bool { + const g: f64 = bfr.value(got); + if (std.math.isFinite(lo) and std.math.isFinite(hi)) { + if (sameBits(lo, hi)) stats.decided += 1 else stats.ranged += 1; + } else stats.nonfinite += 1; + if (std.math.isNan(lo) or std.math.isNan(hi)) return std.math.isNan(g); + if (sameBits(lo, hi)) return bfr.value(bfr.bits(@floatCast(lo))) == lo and got == bfr.bits(@floatCast(lo)); + const a, const b = if (atMost(lo, hi)) .{ lo, hi } else .{ hi, lo }; + return atMost(a, g) and atMost(g, b); +} + +fn sameBits(a: f64, b: f64) bool { + return @as(u64, @bitCast(a)) == @as(u64, @bitCast(b)); +} + +/// a <= b with -0 below +0 (false for NaN). +fn atMost(a: f64, b: f64) bool { + return a < b or (a == b and (std.math.signbit(a) or !std.math.signbit(b))); +} + +/// norms_tail.cuh's group_sum8 over stored outputs: each thread's 8 in order from 0, then 8 lanes' xor tree 4, 2, 1. +pub fn groupSum(y: []const u16) f32 { + var s: [8]f32 = undefined; + for (&s, 0..) |*v, l| { + v.* = 0; + for (y[8 * l ..][0..8]) |h| v.* += bfr.value(h); + } + var m: usize = 4; + while (m > 0) : (m >>= 1) { + var t: [8]f32 = undefined; + for (&t, 0..) |*v, l| v.* = s[l] + s[l ^ m]; + s = t; + } + return s[0]; +} + +/// Every h (when written), y and xs of `k`: the outputs in their float64 ranges, the sums bit-equal to the fp32 tree. +pub fn verify(a: std.mem.Allocator, k: Case, h: []const u16, y: []const u16, xs: ?[]const f32, stats: *Stats, msg: []u8) !?[]const u8 { + const n = k.width; + const b = try a.alloc(f64, 4 * n); + defer a.free(b); + const h_lo, const h_hi, const y_lo, const y_hi = .{ b[0..n], b[n .. 2 * n], b[2 * n .. 3 * n], b[3 * n ..] }; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + for (0..k.rows) |row| { + _ = arena.reset(.retain_capacity); + const ra = arena.allocator(); + try rowRef(ra, k, row, spread(), h_lo, h_hi, y_lo, y_hi); + const at = row * n; + for (0..n) |i| { + if (k.writesH() and !inRange(h[at + i], h_lo[i], h_hi[i], stats)) + return try std.fmt.bufPrint(msg, "row {d} h[{d}] = {e} outside float64's [{e}, {e}]", .{ row, i, bfr.value(h[at + i]), h_lo[i], h_hi[i] }); + if (!inRange(y[at + i], y_lo[i], y_hi[i], stats)) + return try std.fmt.bufPrint(msg, "row {d} y[{d}] = {e} outside float64's [{e}, {e}]", .{ row, i, bfr.value(y[at + i]), y_lo[i], y_hi[i] }); + } + if (xs) |sums| for (0..n / 64) |g| { + const want = groupSum(y[at + 64 * g ..][0..64]); + const got = sums[row * (n / 64) + g]; + stats.sums += 1; + const same = @as(u32, @bitCast(got)) == @as(u32, @bitCast(want)) or (std.math.isNan(got) and std.math.isNan(want)); + if (!same) return try std.fmt.bufPrint(msg, "row {d} sum {d} = {e} where the fp32 tree of the stored y gives {e}", .{ row, g, got, want }); + }; + } + return null; +} + +test "the fp32 group tree: lanes then the xor tree, not one chain" { + var y: [64]u16 = undefined; + for (&y, 0..) |*v, i| v.* = bfr.bits(if (i == 0) 0x1p24 else 1); + var chain: f32 = 0; + for (y) |h| chain += bfr.value(h); + try std.testing.expectEqual(@as(f32, 0x1p24), chain); // one chain loses all 63 ones beside 2^24 + try std.testing.expectEqual(@as(f32, 0x1p24 + 56), groupSum(&y)); // lanes 1 to 7 keep theirs +} + +test "a range with one infinite end: a zero below its finite end is refused, the point and the end pass" { + var x: [64]u16 = undefined; + var w: [64]u16 = undefined; + for (&x, &w) |*v, *u| { + v.* = bfr.bits(1); + u.* = bfr.bits(1); + } + x[62] = bfr.bits(0.625); + x[63] = bfr.bits(1.0546875); + w[0] = 0x7f7f; // the largest finite bf16: bf16(x0 inv) of 1 or 1 + 2^-7 within the spread, so y0 in [w0, +Inf] + const k: Case = .{ .rows = 1, .width = 64, .round = true, .x = &x, .w = &w }; + const a = std.testing.allocator; + var lo: [64]f64 = undefined; + var hi: [64]f64 = undefined; + var scratch: [2][64]f64 = undefined; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + try rowRef(arena.allocator(), k, 0, spread(), &scratch[0], &scratch[1], &lo, &hi); + try std.testing.expectEqual(@as(f64, bfr.value(0x7f7f)), lo[0]); + try std.testing.expect(std.math.isPositiveInf(hi[0])); + var yp: [64]f64 = undefined; + try point(arena.allocator(), k, 0, &scratch[0], &yp); + try std.testing.expect(std.math.isFinite(yp[0])); + var y: [64]u16 = undefined; + for (&y, yp) |*o, v| o.* = bfr.bits(@floatCast(v)); + var msg: [200]u8 = undefined; + var stats: Stats = .{}; + try std.testing.expectEqual(@as(?[]const u8, null), try verify(a, k, &y, &y, null, &stats, &msg)); + y[0] = 0x7f80; // +Inf, the range's upper end + try std.testing.expectEqual(@as(?[]const u8, null), try verify(a, k, &y, &y, null, &stats, &msg)); + y[0] = 0; // the planted zero + try std.testing.expect(try verify(a, k, &y, &y, null, &stats, &msg) != null); +} + +test "range ends: the negative infinite side, signed zeros by their bits, NaN only for NaN" { + var s: Stats = .{}; + const inf = std.math.inf(f64); + const big = bfr.value(0x7f7f); + try std.testing.expect(!inRange(0, -inf, -big, &s)); // a zero above the finite upper end + try std.testing.expect(inRange(0xff7f, -inf, -big, &s) and inRange(0xff80, -inf, -big, &s)); + try std.testing.expect(!inRange(0x8000, 0, 0, &s)); // -0 where float64 decides +0 + try std.testing.expect(!inRange(0, -0.0, -0.0, &s)); + try std.testing.expect(!inRange(0x8000, 0, 1, &s) and inRange(0, 0, 1, &s)); // -0 below a +0 end + try std.testing.expect(inRange(0x8000, -0.0, 0, &s) and inRange(0, 0, -0.0, &s)); + try std.testing.expect(inRange(0x7f80, inf, inf, &s) and !inRange(0x7f7f, inf, inf, &s)); + try std.testing.expect(inRange(0x3f80, -inf, inf, &s) and !inRange(0x7fc0, -inf, inf, &s)); + try std.testing.expect(!inRange(0x3f80, std.math.nan(f64), std.math.nan(f64), &s) and inRange(0x7fc0, std.math.nan(f64), std.math.nan(f64), &s)); +} diff --git a/zig/tests/cuda/qk_round_check.zig b/zig/tests/cuda/qk_round_check.zig new file mode 100644 index 0000000000..52ecc1c4b5 --- /dev/null +++ b/zig/tests/cuda/qk_round_check.zig @@ -0,0 +1,115 @@ +//! qk_rows' tf_qk_rows_once and tf_qk_rows_round against core/qk_ref.zig's float64 ranges at their rounding points. +const std = @import("std"); +const core = @import("core"); + +const ref = core.qk_ref; +const bfr = core.bf16_round; +const c = core.contract.spec; + +/// cosf and sinf within 2 ulps of their value either way. +pub const trig = 0x1p-22; + +/// The sum of squares' chain: a warp's xor tree (5), then the head's warps in order. +pub fn depth(dim: usize) usize { + return 5 + dim / 32; +} + +/// An fp32 inverse RMS's relative spread from float64's: the sum's gamma halved by the root, four fp32 steps. +pub fn spread(dim: usize) f64 { + return c.gamma(@floatFromInt(depth(dim)), 0x1p-24) / 2 + 4 * 0x1p-24; +} + +/// One call's host inputs: rows of h query, hk key and (v_from_k: none, k's columns read again) hk value heads. +pub const Case = struct { + rows: usize, + h: usize, + hk: usize, + dim: usize, + rot: usize, + /// bf16 bits [rows, ld]: q heads, then k heads, then (unless v_from_k) v heads. + x: []const u16, + ld: usize, + v_from_k: bool = false, + /// Norm weights' values (null: no norm), cos and sin [rows, rot] as served (null: no RoPE). + wq: ?[]const f64 = null, + wk: ?[]const f64 = null, + cos: ?[]const f64 = null, + sin: ?[]const f64 = null, + /// cos and sin from cosf and sinf (inv_freq), not tables. + computed: bool = false, + points: ref.Points = .{}, + v_norm: bool = false, + eps: f64 = 1e-6, + + pub fn heads(k: Case) usize { + return k.h + 2 * k.hk; + } + + /// Head hd's (all q, then k, then v) first input of row `row`. + pub fn src(k: Case, row: usize, hd: usize) usize { + const at = if (k.v_from_k and hd >= k.h + k.hk) hd - k.hk else hd; + return row * k.ld + at * k.dim; + } +}; + +/// What a check saw: outputs whose float64 range is one value, a range, and outputs non-finite in float64. +pub const Stats = struct { decided: usize = 0, ranged: usize = 0, nonfinite: usize = 0 }; + +/// Head hd of row `row` as float64 ends under `st` from its sum of squares `ss` (null: float64's). +pub fn head(k: Case, row: usize, hd: usize, st: ref.Stray, ss: ?f64, x: []f64, lo: []f64, hi: []f64) void { + for (x, 0..) |*v, i| v.* = bfr.value(k.x[k.src(row, hd) + i]); + const sq = ss orelse ref.squares(x); + if (hd >= k.h + k.hk) { + if (k.v_norm) ref.unitRange(x, sq, k.eps, st, lo, hi) else { + @memcpy(lo, x); + @memcpy(hi, x); + } + return; + } + const w = if (hd < k.h) k.wq else k.wk; + const cs = if (k.cos) |t| t[row * k.rot ..][0..k.rot] else null; + const sn = if (k.sin) |t| t[row * k.rot ..][0..k.rot] else null; + ref.headRange(x, w, sq, k.eps, cs, sn, k.rot, k.points, st, lo, hi); +} + +/// The output index of head hd's first value: q [rows, h, dim], then k and v [rows, hk, dim]. +pub fn outAt(k: Case, row: usize, hd: usize) usize { + if (hd < k.h) return (row * k.h + hd) * k.dim; + const base = k.rows * k.h * k.dim; + if (hd < k.h + k.hk) return base + (row * k.hk + hd - k.h) * k.dim; + return base + k.rows * k.hk * k.dim + (row * k.hk + hd - k.h - k.hk) * k.dim; +} + +/// The fp32 run's stray for the case's kernel. +pub fn stray(k: Case) ref.Stray { + return .{ .spread = spread(k.dim), .u = 0x1p-24, .trig = if (k.computed) trig else 0 }; +} + +/// Every output (q, k, v as outAt lays them) within float64's ends, or of the float64 point's class where not finite. +pub fn verify(a: std.mem.Allocator, k: Case, out: []const u16, stats: *Stats, msg: []u8) !?[]const u8 { + const buf = try a.alloc(f64, 5 * k.dim); + defer a.free(buf); + const x, const lo, const hi, const plo, const phi = .{ buf[0..k.dim], buf[k.dim .. 2 * k.dim], buf[2 * k.dim .. 3 * k.dim], buf[3 * k.dim .. 4 * k.dim], buf[4 * k.dim ..] }; + for (0..k.rows) |row| for (0..k.heads()) |hd| { + head(k, row, hd, stray(k), null, x, lo, hi); + head(k, row, hd, .{}, null, x, plo, phi); + const at = outAt(k, row, hd); + for (0..k.dim) |i| { + const g: f64 = bfr.value(out[at + i]); + const ok = if (std.math.isFinite(lo[i]) and std.math.isFinite(hi[i])) blk: { + if (lo[i] == hi[i]) stats.decided += 1 else stats.ranged += 1; + break :blk std.math.isFinite(g) and lo[i] <= g and g <= hi[i]; + } else blk: { + stats.nonfinite += 1; + break :blk c.classOf(g) == c.classOf(plo[i]); + }; + if (!ok) return try std.fmt.bufPrint(msg, "row {d} head {d} out[{d}] = {e} outside float64's [{e}, {e}]", .{ row, hd, i, g, lo[i], hi[i] }); + } + }; + return null; +} + +test "the inverse's spread grows with the head's warps and stays far below a bf16 ulp" { + try std.testing.expect(spread(512) > spread(256)); + try std.testing.expect(spread(512) < 0x1p-18); +} diff --git a/zig/tests/cuda/qk_rows_tests.zig b/zig/tests/cuda/qk_rows_tests.zig index cdb5105ef3..a6b6a8b178 100644 --- a/zig/tests/cuda/qk_rows_tests.zig +++ b/zig/tests/cuda/qk_rows_tests.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); const check = @import("check.zig"); +const qr = @import("qk_round_check.zig"); const counts = [_]usize{ 1, 3, 17, 300 }; const most = counts[counts.len - 1]; @@ -207,6 +208,139 @@ fn run1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.qk_ } } +const Rope = enum { none, sliding, proportional, table }; + +/// A Gemma-shaped call: dim 256 or 512, the rotation, MLX's rows or one rounding, v's unit norm, v from k's columns. +const Wide = struct { h: usize, hk: usize, dim: usize, rope: Rope, round: bool, v_norm: bool, v_from_k: bool = false }; + +const Dev = struct { + gpu: check.Gpu, + s: cuda.Stream, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn deinit(d: *Dev) void { + for (d.bufs.items) |*b| b.free(); + d.bufs.deinit(d.gpu.gpa); + } + + fn up(d: *Dev, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(d.gpu.d, bytes); + try d.bufs.append(d.gpu.gpa, b); + return b.ptr; + } + + fn poisoned(d: *Dev, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(d.gpu.d, @max(n, 2)); + try d.bufs.append(d.gpu.gpa, b); + try b.fill8(0xff, d.s.handle); + return b.ptr; + } + + fn back(d: *Dev, a: std.mem.Allocator, ptr: u64, n: usize) ![]u16 { + try d.s.synchronize(); + const raw = try a.alloc(u16, n); + const b: cuda.DeviceBuffer = .{ .d = d.gpu.d, .ptr = ptr, .len = n * 2 }; + try b.download(0, std.mem.sliceAsBytes(raw)); + return raw; + } +}; + +/// One Wide call at every row count: outputs in float64's ranges, caches as the outputs, bits equal at every count. +fn wide(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.qk_rows.QkRows, rng: std.Random, c: Wide, stats: *qr.Stats) !void { + var d: Dev = .{ .gpu = gpu, .s = s }; + defer d.deinit(); + const ld = (c.h + (if (c.v_from_k) c.hk else 2 * c.hk)) * c.dim; + const xs = try a.alloc(u16, most * ld); + for (xs, 0..) |*h, i| h.* = bf16(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const wb = try a.alloc(u16, 2 * c.dim); // Gemma 4 12B's: q 0.95 to 1.05, k 0.06 to 0.13 + for (wb, 0..) |*v, i| v.* = bf16(if (i < c.dim) 0.9 + 0.2 * rng.float(f32) else 0.05 + 0.1 * rng.float(f32)); + const rot: usize = if (c.rope == .none) 0 else c.dim; + const half = @max(rot / 2, 1); + const inv = try a.alloc(f32, half); // families/gemma4/config.zig's inverseFrequencies + for (inv, 0..) |*v, j| v.* = if (c.rope == .proportional and j >= 64) 0 else @floatCast(1 / std.math.pow(f64, if (c.rope == .proportional) 1e6 else 1e4, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(c.dim)))); + const tc = try a.alloc(f32, positions * @max(rot, 1)); + const ts = try a.alloc(f32, positions * @max(rot, 1)); + if (c.rope == .table) for (0..positions) |p| for (0..rot) |i| { + const ang = @as(f64, @floatFromInt(p)) * inv[i % half]; + tc[p * rot + i] = @floatCast(@cos(ang)); + ts[p * rot + i] = @floatCast(@sin(ang)); + }; + const len = 600; + const pos = try a.alloc(i32, most); + const slot = try a.alloc(i32, most); + for (pos, slot, 0..) |*p, *sl, i| { // distinct (slot, pos % len) a row + p.* = @intCast((i / 2) % len + len * (i % 5)); + sl.* = @intCast(i % 2); + } + const cs = try a.alloc(f64, most * @max(rot, 1)); + const sn = try a.alloc(f64, most * @max(rot, 1)); + for (0..most) |r| for (0..rot) |i| { + const p: usize = @intCast(pos[r]); + const ang: f64 = @as(f32, @floatFromInt(p)) * inv[i % half]; + cs[r * rot + i] = if (c.rope == .table) tc[p * rot + i] else @cos(ang); + sn[r * rot + i] = if (c.rope == .table) ts[p * rot + i] else @sin(ang); + }; + const dx = try d.up(std.mem.sliceAsBytes(xs)); + const nw = try d.up(std.mem.sliceAsBytes(wb)); + var shape: cuda.qk_rows.Shape = .{ .norms = nw, .norms_bf16 = true, .rot = rot, .round = c.round, .v_norm = c.v_norm }; + switch (c.rope) { + .none => {}, + .table => { + shape.cos = try d.up(std.mem.sliceAsBytes(tc)); + shape.sin = try d.up(std.mem.sliceAsBytes(ts)); + }, + .sliding, .proportional => shape.inv = try d.up(std.mem.sliceAsBytes(inv)), + } + const dp = try d.up(std.mem.sliceAsBytes(pos)); + const dsl = try d.up(std.mem.sliceAsBytes(slot)); + const kp = dx + c.h * c.dim * 2; + const src: cuda.qk_rows.Src = .{ .q = dx, .q_ld = @intCast(ld), .q_hs = @intCast(c.dim), .k = kp, .k_ld = @intCast(ld), .k_hs = @intCast(c.dim), .v = if (c.hk == 0) 0 else if (c.v_from_k) kp else kp + c.hk * c.dim * 2, .v_ld = @intCast(ld), .v_hs = @intCast(c.dim) }; + const cache = 2 * len * c.hk * c.dim; + var outs: [counts.len]cuda.qk_rows.Out = undefined; + for (counts, &outs) |m, *o| { + o.* = .{ .q = try d.poisoned(m * c.h * c.dim * 2), .k = try d.poisoned(m * c.hk * c.dim * 2), .v = try d.poisoned(m * c.hk * c.dim * 2), .kc = if (c.hk > 0) try d.poisoned(cache * 2) else null, .vc = if (c.hk > 0) try d.poisoned(cache * 2) else null, .len = len }; + try g.rowsFrom(s, src, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o.*, m); + } + const o = outs[counts.len - 1]; + const nq = most * c.h * c.dim; + const nk = most * c.hk * c.dim; + const out = try a.alloc(u16, nq + 2 * nk); + @memcpy(out[0..nq], try d.back(a, o.q, nq)); + @memcpy(out[nq..][0..nk], try d.back(a, o.k, nk)); + @memcpy(out[nq + nk ..], try d.back(a, o.v, nk)); + const wv = try a.alloc(f64, 2 * c.dim); + for (wv, wb) |*v, h| v.* = bf(h); + const k: qr.Case = .{ .rows = most, .h = c.h, .hk = c.hk, .dim = c.dim, .rot = rot, .x = xs, .ld = ld, .v_from_k = c.v_from_k, .wq = wv[0..c.dim], .wk = wv[c.dim..], .cos = if (rot > 0) cs else null, .sin = if (rot > 0) sn else null, .computed = c.rope == .sliding or c.rope == .proportional, .points = if (c.round) .{ .row = true, .out = true } else .{}, .v_norm = c.v_norm }; + var msg: [200]u8 = undefined; + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "qk rows {d}/{d}x{d} {t}{s}{s}{s}", .{ c.h, c.hk, c.dim, c.rope, if (c.round) ", MLX's rows" else "", if (c.v_norm) ", unit v norm" else "", if (c.v_from_k) ", v from k" else "" }); + if (try qr.verify(a, k, out, stats, &msg)) |why| try check.expect(false, "{s}: {s}", .{ label, why }); + if (c.hk > 0) { // the caches hold what the outputs hold, at (slot * len + pos % len) + const kc = try d.back(a, o.kc.?, cache); + const vc = try d.back(a, o.vc.?, cache); + for (0..most) |r| { + const at = @as(usize, @intCast(slot[r])) * len + @as(usize, @intCast(pos[r])) % len; + for (0..c.hk) |j| { + try check.sameBytes(label, std.mem.sliceAsBytes(kc[(at * c.hk + j) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(out[nq + (r * c.hk + j) * c.dim ..][0..c.dim])); + try check.sameBytes(label, std.mem.sliceAsBytes(vc[(at * c.hk + j) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(out[nq + nk + (r * c.hk + j) * c.dim ..][0..c.dim])); + } + } + } + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, oo| { + try check.sameBytes(label, std.mem.sliceAsBytes(try d.back(a, oo.q, m * c.h * c.dim)), std.mem.sliceAsBytes(out[0 .. m * c.h * c.dim])); + try check.sameBytes(label, std.mem.sliceAsBytes(try d.back(a, oo.k, m * c.hk * c.dim)), std.mem.sliceAsBytes(out[nq..][0 .. m * c.hk * c.dim])); + try check.sameBytes(label, std.mem.sliceAsBytes(try d.back(a, oo.v, m * c.hk * c.dim)), std.mem.sliceAsBytes(out[nq + nk ..][0 .. m * c.hk * c.dim])); + } + if (c.dim <= 256 and !c.round and c.v_norm) { // q and k as tf_qk_rows writes them: v's norm alone moves the call + const o2: cuda.qk_rows.Out = .{ .q = try d.poisoned(nq * 2), .k = try d.poisoned(nk * 2), .v = try d.poisoned(nk * 2), .len = len }; + var plain = shape; + plain.v_norm = false; + try g.rowsFrom(s, src, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, plain, dp, dsl, o2, most); + try check.sameBytes("qk rows once: q as tf_qk_rows'", std.mem.sliceAsBytes(try d.back(a, o2.q, nq)), std.mem.sliceAsBytes(out[0..nq])); + try check.sameBytes("qk rows once: k as tf_qk_rows'", std.mem.sliceAsBytes(try d.back(a, o2.k, nk)), std.mem.sliceAsBytes(out[nq..][0..nk])); + } +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -223,4 +357,16 @@ pub fn run(gpu: check.Gpu) !void { try run1(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 2, .dim = 128, .rot = 128, .norm = false }, &worst); // a warp a head, no norm try run1(gpu, a, s, &g, prng.random(), .{ .h = 24, .hk = 4, .dim = 256, .rot = 64, .norm = true }, &worst); // Qwen3.8-27B check.pass("qk rows: q and k within float64 tolerance (worst {d:.3}), v and the caches exact, bits at any row count", .{worst}); + var stats: qr.Stats = .{}; + for ([_]usize{ 16, 32 }) |h| { // Gemma 4 12B and 31B: sliding 2:1 x 256, global (k's columns as v) 16:1 and 8:1 x 512 + try wide(gpu, a, s, &g, prng.random(), .{ .h = h, .hk = h / 2, .dim = 256, .rope = .sliding, .round = true, .v_norm = true }, &stats); + try wide(gpu, a, s, &g, prng.random(), .{ .h = h, .hk = if (h == 16) 1 else 4, .dim = 512, .rope = .proportional, .round = true, .v_norm = true, .v_from_k = true }, &stats); + } + for ([_]usize{ 256, 512 }) |dim| try wide(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 0, .dim = dim, .rope = if (dim == 256) .sliding else .proportional, .round = true, .v_norm = false }, &stats); // the 12B's assistant: q only + try wide(gpu, a, s, &g, prng.random(), .{ .h = 32, .hk = 4, .dim = 512, .rope = .table, .round = false, .v_norm = false }, &stats); // one rounding at 512 + try wide(gpu, a, s, &g, prng.random(), .{ .h = 32, .hk = 4, .dim = 512, .rope = .none, .round = true, .v_norm = true }, &stats); + try wide(gpu, a, s, &g, prng.random(), .{ .h = 24, .hk = 4, .dim = 256, .rope = .table, .round = false, .v_norm = true }, &stats); // the 27B's shape, v normed + try wide(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .dim = 128, .rope = .sliding, .round = false, .v_norm = true }, &stats); + try std.testing.expectError(error.HeadShape, g.rows(s, 0, .{ .h = 1, .hk = 1, .dim = 544 }, .{}, 0, null, .{ .q = 0, .k = 0, .v = 0 }, 1)); + check.pass("qk rows dims to 512, MLX's rows, unit v norm: {d} outputs bit-equal to float64's at the rounding points, {d} within the fp32 inverse's or cosf's range, {d} non-finite as float64; caches as the outputs; every row count's bits equal", .{ stats.decided, stats.ranged, stats.nonfinite }); } diff --git a/zig/tests/cuda/window_bench.zig b/zig/tests/cuda/window_bench.zig new file mode 100644 index 0000000000..f4b7fa6441 --- /dev/null +++ b/zig/tests/cuda/window_bench.zig @@ -0,0 +1,134 @@ +//! window_attention's GPU time at Kolibri-1's served shapes (its caller) and Gemma 4's, 1k to 128k keys (synthetic). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); + +const wa = cuda.window_attention; + +const warm = 3; +const reps = 25; + +/// One timed call: decode rows at a position, or a prompt chunk at p0. +pub const Shape = struct { name: []const u8, h: usize, hk: usize, d: usize, win: usize, ring: usize, rows: usize, at: usize, prompt: bool, p32: bool = false }; + +/// Kolibri-1's calls: 48 query heads over 4 of 128, window 513, ring 8768 (an 8192-row chunk); full layers unwindowed. +pub const callers = [_]Shape{ + .{ .name = "kolibri decode tf_window_decode_tc, 1 row", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 1, .at = 6000, .prompt = false }, + .{ .name = "kolibri decode tf_window_decode_tc, 4 rows", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 4, .at = 6000, .prompt = false }, + .{ .name = "kolibri decode tf_window_decode_tc, 16 rows", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 16, .at = 6000, .prompt = false }, + .{ .name = "kolibri decode tf_window_decode_tc, 64 rows", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 64, .at = 6000, .prompt = false }, + .{ .name = "kolibri sliding prompt wattn_kernel<128,8,4,8>, 2048 rows at 0", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 2048, .at = 0, .prompt = true }, + .{ .name = "kolibri sliding prompt wattn_kernel<128,8,4,8>, 8192 rows at 8192", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 8768, .rows = 8192, .at = 8192, .prompt = true }, + .{ .name = "kolibri full prompt wattn_kernel<128,8,4,8>, 2048 rows at 8192", .h = 48, .hk = 4, .d = 128, .win = 0, .ring = 32768, .rows = 2048, .at = 8192, .prompt = true }, + .{ .name = "kolibri full prompt wattn_kernel<128,8,4,8>, 2048 rows at 30720", .h = 48, .hk = 4, .d = 128, .win = 0, .ring = 32768, .rows = 2048, .at = 30720, .prompt = true }, +}; + +const contexts = [_]usize{ 1024, 8192, 32768, 131072 }; + +/// Gemma 4's sliding (256, window 1024, ring 1536) and global layers (512, linear): decode rows, 512-row prompt chunks. +pub fn sliding(out: []Shape) []Shape { + var n: usize = 0; + for (contexts) |ctx| { + for ([_]usize{ 1, 4 }) |rows| { + out[n] = .{ .name = "12B sliding decode (16/8 x 256)", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .rows = rows, .at = ctx - 1, .prompt = false }; + n += 1; + } + out[n] = .{ .name = "31B sliding decode (32/16 x 256)", .h = 32, .hk = 16, .d = 256, .win = 1024, .ring = 1536, .rows = 1, .at = ctx - 1, .prompt = false }; + n += 1; + for ([_]bool{ true, false }) |p32| { + out[n] = .{ .name = if (p32) "12B sliding prompt (16/8 x 256), P32" else "12B sliding prompt (16/8 x 256), P bf16", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .rows = 512, .at = ctx - 512, .prompt = true, .p32 = p32 }; + n += 1; + } + if (!wa.wide(512)) continue; + for ([_]usize{ 1, 4 }) |rows| { + out[n] = .{ .name = "12B global decode (16/1 x 512)", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = ctx, .rows = rows, .at = ctx - 1, .prompt = false }; + n += 1; + } + out[n] = .{ .name = "31B global decode (32/4 x 512)", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = ctx, .rows = 1, .at = ctx - 1, .prompt = false }; + n += 1; + for ([_]bool{ true, false }) |p32| { + out[n] = .{ .name = if (p32) "12B global prompt (16/1 x 512), P32" else "12B global prompt (16/1 x 512), P bf16", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = ctx, .rows = 512, .at = ctx - 512, .prompt = true, .p32 = p32 }; + n += 1; + } + out[n] = .{ .name = "31B global prompt (32/4 x 512), P32", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = ctx, .rows = 512, .at = ctx - 512, .prompt = true, .p32 = true }; + n += 1; + } + return out[0..n]; +} + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// Microseconds a call takes on the GPU (median of `reps`); inputs are normal bf16 values, outputs discarded. +pub fn time(gpu: check.Gpu, s: cuda.Stream, g: *const wa.WindowAttention, x: Shape) !f64 { + const crow = if (x.prompt) x.ring else x.ring * x.rows; // decode: a stream a row + const q_bytes = x.rows * x.h * x.d * 2; + const c_bytes = crow * x.hk * x.d * 2; + var prng = std.Random.DefaultPrng.init(x.rows *% 31 +% x.at); + const host = try gpu.gpa.alloc(u16, @max(q_bytes, c_bytes) / 2); + defer gpu.gpa.free(host); + for (host) |*v| { + const f: f32 = prng.random().floatNorm(f32) * 0.5; + v.* = @truncate(@as(u32, @bitCast(f)) >> 16); + } + var bufs: [7]cuda.DeviceBuffer = undefined; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + const sizes = [_]usize{ q_bytes, c_bytes, c_bytes, q_bytes, x.rows * 4, x.rows * 4, if (x.prompt) 16 else wa.WindowAttention.scratchBytes(.{ .h = x.h, .hk = x.hk, .d = x.d, .win = x.win, .ring = x.ring, .scale = 1 }, x.rows) }; + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, @max(n, 16)); + made += 1; + } + try bufs[0].upload(0, std.mem.sliceAsBytes(host[0 .. q_bytes / 2])); + try bufs[1].upload(0, std.mem.sliceAsBytes(host[0 .. c_bytes / 2])); + try bufs[2].upload(0, std.mem.sliceAsBytes(host[0 .. c_bytes / 2])); + const pos = try gpu.gpa.alloc(i32, 2 * x.rows); + defer gpu.gpa.free(pos); + for (pos[0..x.rows], pos[x.rows..], 0..) |*p, *sl, i| { + p.* = @intCast(x.at); + sl.* = @intCast(i); + } + try bufs[4].upload(0, std.mem.sliceAsBytes(pos[0..x.rows])); + try bufs[5].upload(0, std.mem.sliceAsBytes(pos[x.rows..])); + const w: wa.Window = .{ .h = x.h, .hk = x.hk, .d = x.d, .win = x.win, .ring = x.ring, .scale = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(x.d)))) }; + var start = try cuda.Event.init(gpu.d, true); + defer start.deinit(); + var end = try cuda.Event.init(gpu.d, true); + defer end.deinit(); + var us: [reps]f64 = undefined; + for (0..warm + reps) |i| { + try start.record(s); + if (x.prompt) { + try call(g, s, bufs[0].ptr, bufs[1].ptr, bufs[2].ptr, bufs[3].ptr, x, w); + } else try g.decode(s, bufs[0].ptr, bufs[1].ptr, bufs[2].ptr, bufs[3].ptr, bufs[4].ptr, bufs[5].ptr, w, bufs[6].ptr, x.rows); + try end.record(s); + try s.synchronize(); + if (i >= warm) us[i - warm] = @as(f64, try start.elapsedMs(end)) * 1000; + } + return median(&us); +} + +fn call(g: *const wa.WindowAttention, s: cuda.Stream, q: u64, kc: u64, vc: u64, out: u64, x: Shape, w0: wa.Window) !void { + if (@hasDecl(wa, "wide") and wa.wide(x.d)) { + var w = w0; + w.p32 = x.p32; + return g.promptRows(s, q, kc, vc, out, x.at, x.rows, w, null); + } + return g.prompt(s, q, kc, vc, out, x.at, x.rows, w0); +} + +/// `which`: callers (Kolibri-1's shapes) or gemma (the new head dims). +pub fn run(gpu: check.Gpu, which: []const u8) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try wa.WindowAttention.load(gpu.d); + defer g.unload(); + var buf: [64]Shape = undefined; + const list: []const Shape = if (std.mem.eql(u8, which, "callers")) &callers else if (@hasDecl(wa, "wide")) sliding(&buf) else return error.UnknownBench; + for (list) |x| { + const us = try time(gpu, s, &g, x); + std.debug.print("RESULT window-bench {s}, rows {d}, {s} {d}: {d:.1} us\n", .{ x.name, x.rows, if (x.prompt) "p0" else "position", x.at, us }); + } +} diff --git a/zig/tests/cuda/window_rows_check.zig b/zig/tests/cuda/window_rows_check.zig new file mode 100644 index 0000000000..823d0debf3 --- /dev/null +++ b/zig/tests/cuda/window_rows_check.zig @@ -0,0 +1,167 @@ +//! window_rows.cu's calls on the host: each output against core/attention_range.zig's float64 point and range. +const std = @import("std"); +const core = @import("core"); + +const ar = core.attention_range; +const bfr = core.bf16_round; + +/// One decode or prompt call as the host holds it: q [rows, h, d], caches [cache rows, hk, d] (bf16 bits). +pub const Call = struct { + h: usize, + hk: usize, + d: usize, + win: usize, + ring: usize, + scale: f32, + prompt: bool, + p32: bool = false, + rows: usize, + q: []const u16, + kc: []const u16, + vc: []const u16, + /// Decode: each row's position and stream slot (cache rows slot * ring + key % ring). + pos: []const i32 = &.{}, + slot: []const i32 = &.{}, + /// Prompt: the first row's position and each row's last key (null: its own position); ring 0 is a linear cache. + p0: usize = 0, + hi: ?[]const i32 = null, + + /// Row r's position, its last key and its stream's first cache row. + pub fn at(c: Call, r: usize) struct { p: usize, last: usize, base: usize } { + if (c.prompt) return .{ .p = c.p0 + r, .last = if (c.hi) |h| @intCast(h[r]) else c.p0 + r, .base = 0 }; + const p: usize = @intCast(c.pos[r]); + return .{ .p = p, .last = p, .base = @as(usize, @intCast(c.slot[r])) * c.ring }; + } + + pub fn first(c: Call, p: usize) usize { + return if (c.win > 0 and p + 1 > c.win) p + 1 - c.win else 0; + } + + pub fn cacheRow(c: Call, base: usize, key: usize) usize { + return base + if (c.ring > 0) key % c.ring else key; + } + + pub fn form(c: Call, keys: usize) ar.Form { + return if (c.prompt) ar.prompt(c.d, keys, c.p32) else ar.decode(c.d, keys); + } +}; + +/// Scratch for a call's heads: one (row, key head)'s keys and values as float64, and a head's weights and ends. +pub const Work = struct { + a: std.mem.Allocator, + k: []f64, + v: []f64, + q: []f64, + w: []f64, + t: []f64, + point: []f64, + half: []f64, + + pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Work { + var self: Work = undefined; + self.a = a; + inline for (.{ "k", "v" }) |n| @field(self, n) = try a.alloc(f64, most * d); + inline for (.{ "w", "t" }) |n| @field(self, n) = try a.alloc(f64, most); + inline for (.{ "q", "point", "half" }) |n| @field(self, n) = try a.alloc(f64, d); + return self; + } + + pub fn deinit(self: *Work) void { + inline for (.{ "k", "v", "w", "t", "q", "point", "half" }) |n| self.a.free(@field(self, n)); + } +}; + +/// The most keys any row of the call sees. +pub fn mostKeys(c: Call) usize { + var n: usize = 1; + for (0..c.rows) |r| { + const w = c.at(r); + n = @max(n, w.last + 1 - c.first(w.p)); + } + return n; +} + +/// Each output head's float64 point and half range, in out's order ([rows, h, d]); `each` sees every head's. +fn heads(c: Call, wk: *Work, ctx: anytype, comptime each: fn (@TypeOf(ctx), usize, []const f64, []const f64) anyerror!void) !void { + const d = c.d; + const g = c.h / c.hk; + for (0..c.rows) |r| { + const w = c.at(r); + const lo = c.first(w.p); + const n = w.last + 1 - lo; + for (0..c.hk) |kh| { + for (lo..w.last + 1, 0..) |key, j| for (0..d) |e| { + const at = (c.cacheRow(w.base, key) * c.hk + kh) * d + e; + wk.k[j * d + e] = bfr.value(c.kc[at]); + wk.v[j * d + e] = bfr.value(c.vc[at]); + }; + for (0..g) |hh| { + const head = kh * g + hh; + for (wk.q, 0..) |*x, e| x.* = bfr.value(c.q[(r * c.h + head) * d + e]); + ar.head(wk.q, wk.k, wk.v, n, d, c.scale, c.form(n), wk.w, wk.t, wk.point, wk.half); + try each(ctx, (r * c.h + head) * d, wk.point, wk.half); + } + } + } +} + +const Verify = struct { out: []const u16, stats: *ar.Stats, msg: []u8, d: usize, h: usize, bad: ?[]const u8 = null }; + +fn verifyOne(v: *Verify, at: usize, point: []const f64, half: []const f64) !void { + if (v.bad != null) return; + for (point, half, 0..) |p, h, e| { + const got = v.out[at + e]; + if (ar.within(p, h, got, v.stats)) continue; + v.bad = try std.fmt.bufPrint(v.msg, "row {d} head {d} out[{d}] = {e} (0x{x:0>4}) outside float64's {e} +- {e}", .{ at / v.d / v.h, at / v.d % v.h, e, bfr.value(got), got, p, h }); + return; + } +} + +/// Every output within its range (NaN where float64 gives NaN, infinities by their bits); null, or what failed. +pub fn verify(a: std.mem.Allocator, c: Call, out: []const u16, stats: *ar.Stats, msg: []u8) !?[]const u8 { + var wk = try Work.init(a, c.d, mostKeys(c)); + defer wk.deinit(); + var v: Verify = .{ .out = out, .stats = stats, .msg = msg, .d = c.d, .h = c.h }; + try heads(c, &wk, &v, verifyOne); + return v.bad; +} + +const Ref = struct { want: []f64, half: []f64 }; + +fn refOne(r: *Ref, at: usize, point: []const f64, half: []const f64) !void { + @memcpy(r.want[at..][0..point.len], point); + @memcpy(r.half[at..][0..half.len], half); +} + +/// Every output's float64 point and half range, in out's order. +pub fn reference(a: std.mem.Allocator, c: Call, want: []f64, half: []f64) !void { + var wk = try Work.init(a, c.d, mostKeys(c)); + defer wk.deinit(); + var r: Ref = .{ .want = want, .half = half }; + try heads(c, &wk, &r, refOne); +} + +/// A row of d bf16 values of RMS `rms` (a unit-RMS normal row times `rms`, as a norm with a constant weight leaves it). +pub fn row(rng: std.Random, x: []u16, rms: f32) void { + var ss: f64 = 0; + var buf: [512]f32 = undefined; + for (buf[0..x.len]) |*v| { + v.* = rng.floatNorm(f32); + ss += @as(f64, v.*) * v.*; + } + const inv: f32 = @floatCast(@sqrt(@as(f64, @floatFromInt(x.len)) / ss)); + for (x, buf[0..x.len]) |*h, v| h.* = bfr.bits(bfr.round32(v * inv) * rms); +} + +/// Gemma 4's q and k weights (all equal within a layer: q 0.95 to 1.05, k 0.06 to 0.13; mlx-community 12B 4-bit). +pub const q_rms = [2]f32{ 0.95, 1.05 }; +pub const k_rms = [2]f32{ 0.06, 0.13 }; + +test "a row's RMS is as asked" { + var prng = std.Random.DefaultPrng.init(1); + var x: [256]u16 = undefined; + row(prng.random(), &x, 0.125); + var ss: f64 = 0; + for (x) |h| ss += bfr.value(h) * bfr.value(h); + try std.testing.expectApproxEqRel(@as(f64, 0.125), @sqrt(ss / 256), 0.01); +} diff --git a/zig/tests/cuda/window_rows_tests.zig b/zig/tests/cuda/window_rows_tests.zig new file mode 100644 index 0000000000..d66591ad77 --- /dev/null +++ b/zig/tests/cuda/window_rows_tests.zig @@ -0,0 +1,231 @@ +//! window_rows.cu at Gemma 4's shapes and values: outputs in float64's range, a row's bits its own (batch, chunk). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const wc = @import("window_rows_check.zig"); + +const wa = cuda.window_attention; +const ar = core.attention_range; +const bfr = core.bf16_round; + +/// Device buffers freed together. +pub const Bufs = struct { + gpu: check.Gpu, + list: std.ArrayList(cuda.DeviceBuffer) = .empty, + + pub fn deinit(b: *Bufs) void { + for (b.list.items) |*x| x.free(); + b.list.deinit(b.gpu.gpa); + } + + pub fn up(b: *Bufs, bytes: []const u8) !u64 { + const x = try cuda.DeviceBuffer.fromHost(b.gpu.d, bytes); + try b.list.append(b.gpu.gpa, x); + return x.ptr; + } + + /// `n` bytes of 0xff (an output no launch writes reads as NaN). + pub fn out(b: *Bufs, s: cuda.Stream, n: usize) !u64 { + const x = try cuda.DeviceBuffer.alloc(b.gpu.d, @max(n, 16)); + try b.list.append(b.gpu.gpa, x); + try x.fill8(0xff, s.handle); + return x.ptr; + } + + pub fn get(b: *Bufs, ptr: u64, out_: []u16) !void { + const x: cuda.DeviceBuffer = .{ .d = b.gpu.d, .ptr = ptr, .len = out_.len * 2 }; + try x.download(0, std.mem.sliceAsBytes(out_)); + } +}; + +/// Gemma 4's attention inputs: q rows of RMS 0.95-1.05, k 0.06-0.13 (the norms' constant weights), v unit RMS. +pub fn fill(rng: std.Random, q: []u16, kc: []u16, vc: []u16, d: usize) void { + const cq = wc.q_rms[0] + (wc.q_rms[1] - wc.q_rms[0]) * rng.float(f32); + const ck = wc.k_rms[0] + (wc.k_rms[1] - wc.k_rms[0]) * rng.float(f32); + var i: usize = 0; + while (i < q.len) : (i += d) wc.row(rng, q[i..][0..d], cq); + i = 0; + while (i < kc.len) : (i += d) wc.row(rng, kc[i..][0..d], ck); + i = 0; + while (i < vc.len) : (i += d) wc.row(rng, vc[i..][0..d], 1); +} + +/// Key `key` of (row r, key head kh) set to r's first query head of kh at k's RMS: one sharp weight for that head. +pub fn sharpen(c: wc.Call, kc: []u16, r: usize, key: usize, kh: usize) void { + const w = c.at(r); + const src = c.q[(r * c.h + kh * (c.h / c.hk)) * c.d ..][0..c.d]; + const dst = kc[(c.cacheRow(w.base, key) * c.hk + kh) * c.d ..][0..c.d]; + for (dst, src) |*x, y| x.* = bfr.bits(bfr.round32(bfr.value(y) * 0.1)); +} + +fn report(what: []const u8, st: ar.Stats) void { + check.pass("{s}: {d} outputs bit-equal to float64's value, {d} within their range, {d} non-finite as float64", .{ what, st.decided, st.ranged, st.nonfinite }); +} + +const DCase = struct { name: []const u8, h: usize, hk: usize, d: usize, win: usize, ring: usize, streams: usize, top: usize, rows: usize = 64 }; + +/// Decode at 1, 3 and c.rows rows: every output in range, a row's bits at any row count and (linear) cache length. +fn decode1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const wa.WindowAttention, rng: std.Random, c: DCase, total: *ar.Stats) !void { + const counts = [_]usize{ 1, 3, c.rows }; + const most = c.rows; + var b: Bufs = .{ .gpu = gpu }; + defer b.deinit(); + const q = try a.alloc(u16, most * c.h * c.d); + const kc = try a.alloc(u16, c.streams * c.ring * c.hk * c.d); + const vc = try a.alloc(u16, kc.len); + fill(rng, q, kc, vc, c.d); + const pos = try a.alloc(i32, most); + const slot = try a.alloc(i32, most); + for (pos, slot, 0..) |*p, *sl, i| { + p.* = @intCast(if (i == most - 1) c.top - 1 else if (i % 4 == 0) i % 40 else (i * 7919) % c.top); + sl.* = @intCast(i % c.streams); + } + var call: wc.Call = .{ .h = c.h, .hk = c.hk, .d = c.d, .win = c.win, .ring = c.ring, .scale = 1, .prompt = false, .rows = most, .q = q, .kc = kc, .vc = vc, .pos = pos, .slot = slot }; + for (0..most) |r| if (r % 3 == 0) { // a sharp key in every third row's window, for each key head + const w = call.at(r); + const lo = call.first(w.p); + for (0..c.hk) |kh| sharpen(call, kc, r, lo + (r * 31 + kh * 7) % (w.p + 1 - lo), kh); + }; + const dq = try b.up(std.mem.sliceAsBytes(q)); + const dk = try b.up(std.mem.sliceAsBytes(kc)); + const dv = try b.up(std.mem.sliceAsBytes(vc)); + const dp = try b.up(std.mem.sliceAsBytes(pos)); + const ds = try b.up(std.mem.sliceAsBytes(slot)); + const w: wa.Window = .{ .h = c.h, .hk = c.hk, .d = c.d, .win = c.win, .ring = c.ring, .scale = 1 }; + const scratch = try b.out(s, wa.WindowAttention.scratchBytes(w, most)); + var outs: [counts.len]u64 = undefined; + for (counts, &outs) |m, *o| { + o.* = try b.out(s, m * c.h * c.d * 2); + try g.decode(s, dq, dk, dv, o.*, dp, ds, w, scratch, m); + } + try s.synchronize(); + const wide = try a.alloc(u16, most * c.h * c.d); + try b.get(outs[counts.len - 1], wide); + var st: ar.Stats = .{}; + var msg: [240]u8 = undefined; + if (try wc.verify(a, call, wide, &st, &msg)) |why| return check.expect(false, "window rows decode {s}: {s}", .{ c.name, why }); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, o| { + const got = try a.alloc(u16, m * c.h * c.d); + try b.get(o, got); + try check.sameBytes("window rows decode: a row count's bits equal the widest call's", std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(wide[0 .. m * c.h * c.d])); + } + if (c.win == 0 and c.streams == 1 and c.ring <= 32768) { // a cache declared 16x longer: the same bits + var w2 = w; + w2.ring = 16 * c.ring; + const o2 = try b.out(s, most * c.h * c.d * 2); + try g.decode(s, dq, dk, dv, o2, dp, ds, w2, try b.out(s, wa.WindowAttention.scratchBytes(w2, most)), most); + try s.synchronize(); + const got = try a.alloc(u16, most * c.h * c.d); + try b.get(o2, got); + try check.sameBytes("window rows decode: empty chunks past a row's keys leave its bits", std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(wide)); + } + report(c.name, st); + total.decided += st.decided; + total.ranged += st.ranged; + total.nonfinite += st.nonfinite; +} + +const PCase = struct { name: []const u8, h: usize, hk: usize, d: usize, win: usize, ring: usize, p0: usize, n: usize, p32: bool, image: ?[2]usize = null }; + +/// Prompt rows p0 .. p0 + n, an image span's rows seeing its end: every output in range, the same bits in two chunks. +fn prompt1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const wa.WindowAttention, rng: std.Random, c: PCase, total: *ar.Stats) !void { + var b: Bufs = .{ .gpu = gpu }; + defer b.deinit(); + const q = try a.alloc(u16, c.n * c.h * c.d); + const crow = if (c.ring > 0) c.ring else c.p0 + c.n; + const kc = try a.alloc(u16, crow * c.hk * c.d); + const vc = try a.alloc(u16, kc.len); + fill(rng, q, kc, vc, c.d); + const hi = try a.alloc(i32, c.n); + for (hi, 0..) |*x, i| x.* = @intCast(c.p0 + i); + if (c.image) |sp| for (sp[0]..sp[1]) |i| { + hi[i] = @intCast(c.p0 + sp[1] - 1); + }; + var call: wc.Call = .{ .h = c.h, .hk = c.hk, .d = c.d, .win = c.win, .ring = c.ring, .scale = 1, .prompt = true, .p32 = c.p32, .rows = c.n, .q = q, .kc = kc, .vc = vc, .p0 = c.p0, .hi = if (c.image != null) hi else null }; + var r: usize = 0; + while (r < c.n) : (r += 5) { + const w = call.at(r); + const lo = call.first(w.p); + for (0..c.hk) |kh| sharpen(call, kc, r, lo + (r * 31 + kh * 7) % (w.last + 1 - lo), kh); + } + const dq = try b.up(std.mem.sliceAsBytes(q)); + const dk = try b.up(std.mem.sliceAsBytes(kc)); + const dv = try b.up(std.mem.sliceAsBytes(vc)); + const dh = try b.up(std.mem.sliceAsBytes(hi)); + const w: wa.Window = .{ .h = c.h, .hk = c.hk, .d = c.d, .win = c.win, .ring = c.ring, .scale = 1, .p32 = c.p32 }; + const hp: ?u64 = if (c.image != null) dh else null; + const full = try b.out(s, c.n * c.h * c.d * 2); + try g.promptRows(s, dq, dk, dv, full, c.p0, c.n, w, hp); + // the chunk in two at a cut no image span crosses + var cut = c.n / 2; + if (c.image) |sp| if (cut > sp[0] and cut < sp[1]) { + cut = sp[1]; + }; + const row_bytes = c.h * c.d * 2; + const two = try b.out(s, c.n * row_bytes); + try g.promptRows(s, dq, dk, dv, two, c.p0, cut, w, hp); + if (cut < c.n) try g.promptRows(s, dq + cut * row_bytes, dk, dv, two + cut * row_bytes, c.p0 + cut, c.n - cut, w, if (hp) |p| p + cut * 4 else null); + try s.synchronize(); + const got = try a.alloc(u16, c.n * c.h * c.d); + try b.get(full, got); + var st: ar.Stats = .{}; + var msg: [240]u8 = undefined; + if (try wc.verify(a, call, got, &st, &msg)) |why| return check.expect(false, "window rows prompt {s}: {s}", .{ c.name, why }); + const split = try a.alloc(u16, got.len); + try b.get(two, split); + try check.sameBytes("window rows prompt: a chunk split in two gives each row's bits", std.mem.sliceAsBytes(split), std.mem.sliceAsBytes(got)); + report(c.name, st); + total.decided += st.decided; + total.ranged += st.ranged; + total.nonfinite += st.nonfinite; +} + +/// Gemma 4's sliding layers (12B: 16 query heads over 8, 31B: 32 over 16; head dim 256, window 1024, ring 1536). +pub const sliding_decode = [_]DCase{ + .{ .name = "12B sliding decode (16/8 heads of 256, window 1024, ring 1536, 3 streams)", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .streams = 3, .top = 40000 }, + .{ .name = "31B sliding decode (32/16 heads of 256, window 1024, ring 1536)", .h = 32, .hk = 16, .d = 256, .win = 1024, .ring = 1536, .streams = 2, .top = 3000 }, + .{ .name = "no window over a linear cache (16/4 heads of 256, 3000 positions)", .h = 16, .hk = 4, .d = 256, .win = 0, .ring = 3000, .streams = 1, .top = 3000 }, +}; + +/// Gemma 4's global layers: head dim 512, no window over a linear cache (12B 16 query heads over 1, 31B 32 over 4). +pub const global_decode = [_]DCase{ + .{ .name = "12B global decode (16/1 heads of 512, linear cache of 262144, a row at its end)", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 262144, .streams = 1, .top = 262144, .rows = 4 }, + .{ .name = "12B global decode (16/1 heads of 512, linear cache of 20000)", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 20000, .streams = 1, .top = 20000, .rows = 17 }, + .{ .name = "31B global decode (32/4 heads of 512, linear caches of 40000, 2 streams)", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = 40000, .streams = 2, .top = 40000, .rows = 17 }, +}; + +pub const global_prompt = [_]PCase{ + .{ .name = "12B global prompt, 128 rows at 8000, P32", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 0, .p0 = 8000, .n = 128, .p32 = true }, + .{ .name = "12B global prompt, 16 rows at 130000, P32", .h = 16, .hk = 1, .d = 512, .win = 0, .ring = 0, .p0 = 130000, .n = 16, .p32 = true }, + .{ .name = "31B global prompt, 300 rows at 0, P32", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = 0, .p0 = 0, .n = 300, .p32 = true }, + .{ .name = "31B global prompt, 96 rows at 3000, P as bf16", .h = 32, .hk = 4, .d = 512, .win = 0, .ring = 0, .p0 = 3000, .n = 96, .p32 = false }, +}; + +pub const sliding_prompt = [_]PCase{ + .{ .name = "12B sliding prompt, first chunk, P32", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .p0 = 0, .n = 512, .p32 = true }, + .{ .name = "12B sliding prompt, the ring wraps, P32", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .p0 = 5000, .n = 512, .p32 = true }, + .{ .name = "12B sliding prompt, an image's rows 40-300 see each other, P32", .h = 16, .hk = 8, .d = 256, .win = 1024, .ring = 1536, .p0 = 1300, .n = 512, .p32 = true, .image = .{ 40, 300 } }, + .{ .name = "31B sliding prompt, 333 rows, P as bf16", .h = 32, .hk = 16, .d = 256, .win = 1024, .ring = 1536, .p0 = 700, .n = 333, .p32 = false }, + .{ .name = "no window over a linear cache (16/4 heads of 256), P32", .h = 16, .hk = 4, .d = 256, .win = 0, .ring = 0, .p0 = 900, .n = 300, .p32 = true }, +}; + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try wa.WindowAttention.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(0x6a77); + var total: ar.Stats = .{}; + for (sliding_decode ++ global_decode) |c| { + _ = arena.reset(.retain_capacity); + try decode1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); + } + for (sliding_prompt ++ global_prompt) |c| { + _ = arena.reset(.retain_capacity); + try prompt1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); + } + report("window rows (all)", total); +}