From a2bed92ab1344ac0ae8f088254956515b54383e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:01:33 +0200 Subject: [PATCH 1/2] cuda window_attention: head dim 256, any GQA, a row's last key and fp32 probabilities window_rows.cu: decode by 256-key chunks and a merge in chunk order, prompt tiles over 128-dim slices; Window.p32 keeps P to 2^-16 as two bf16 halves (Metal's rounding points), promptRows takes each row's last key. window_attention.cu and its callers are unchanged. core/attention_range.zig: float64 at the rounding points and the range fp32 allows. --- docs/kernel-contract.md | 6 +- zig/build/cuda.zig | 8 + zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/contract_controls.cu | 40 ++ zig/kernels/cuda/window_rows.cu | 38 ++ zig/kernels/cuda/window_rows.cuh | 560 ++++++++++++++++++ zig/src/core/attention_range.zig | 170 ++++++ zig/src/core/bf16_round.zig | 79 +++ zig/src/core/contract/registry.zig | 2 +- zig/src/core/contract/window_rows_records.zig | 10 + zig/src/core/root.zig | 4 + zig/src/cuda/kernels.zig | 1 + zig/src/cuda/window_attention.zig | 104 +++- zig/tests/cuda/contract_tests.zig | 8 + zig/tests/cuda/contract_window_rows.zig | 296 +++++++++ zig/tests/cuda/main.zig | 6 + zig/tests/cuda/window_bench.zig | 121 ++++ zig/tests/cuda/window_rows_check.zig | 167 ++++++ zig/tests/cuda/window_rows_tests.zig | 217 +++++++ 19 files changed, 1832 insertions(+), 6 deletions(-) create mode 100644 zig/kernels/cuda/window_rows.cu create mode 100644 zig/kernels/cuda/window_rows.cuh create mode 100644 zig/src/core/attention_range.zig create mode 100644 zig/src/core/bf16_round.zig create mode 100644 zig/src/core/contract/window_rows_records.zig create mode 100644 zig/tests/cuda/contract_window_rows.zig create mode 100644 zig/tests/cuda/window_bench.zig create mode 100644 zig/tests/cuda/window_rows_check.zig create mode 100644 zig/tests/cuda/window_rows_tests.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index bc1e3f6768..71bfdecc89 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -111,8 +111,10 @@ Mamba scan of a verify window (its host twin's bytes), qmm_group's cluster tile, 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 dim 256 (`window_rows.cu`: decode chunks and their +merge, prompt tiles with P as bf16 or in two halves, a row's last key), each output within float64's range over the +spread its fp32 steps and tensor-core sums allow (`core/attention_range.zig`; six controls, each one planted fault). ## CUDA records per nvcc release diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6dfec448ca..688f7832ac 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -321,6 +321,14 @@ 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); + const wrows_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_window_rows.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&wrows_contract.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| { diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 25315ca223..5801b0a7d8 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.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/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index e0d3bb2d51..39fc742438 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -85,3 +85,43 @@ 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 dim 256, 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); +} diff --git a/zig/kernels/cuda/window_rows.cu b/zig/kernels/cuda/window_rows.cu new file mode 100644 index 0000000000..aa31b1bf4a --- /dev/null +++ b/zig/kernels/cuda/window_rows.cu @@ -0,0 +1,38 @@ +// window_attention's rows at head dim 256 (window_rows.cuh): decode partials and merge, prompt tiles. + +#include "window_rows.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); +} 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/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/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..6186781bea 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -221,7 +221,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; /// 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..6ce2a8aebd --- /dev/null +++ b/zig/src/core/contract/window_rows_records.zig @@ -0,0 +1,10 @@ +//! window_rows.cu's launches on CUDA (window attention at head dim 256), 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 } } }, +}; diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 721310a5eb..ed4d86ba01 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -27,6 +27,7 @@ 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"); /// 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 +41,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 +66,11 @@ test { _ = tuning; _ = norms_ref; _ = act_ref; + _ = bf16_round; _ = contract; _ = tree_round; _ = attention_ref; + _ = attention_range; _ = gdn_prep_ref; _ = qk_ref; _ = moe_ref; 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/window_attention.zig b/zig/src/cuda/window_attention.zig index 40569b7ef9..c3e180194b 100644 --- a/zig/src/cuda/window_attention.zig +++ b/zig/src/cuda/window_attention.zig @@ -10,8 +10,18 @@ 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; +} + +/// window_rows.cu's decode chunk: keys [c * chunk_keys, (c + 1) * chunk_keys), merged in chunk order. +pub const chunk_keys = 256; +/// Prompt rows' dynamic shared bytes at 256: 4 slots of 32 keys, then the warps' partial scores. +const rows_smem = [_]u32{4 * 32 * 32 * 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 +35,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 +50,41 @@ 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}, &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); + } 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 { + _ = d; + return rows_smem[0]; + } + + pub fn wideFns(g: *const WindowAttention, d: usize) *const Wide { + _ = d; + return &g.wide_fns[0]; + } + + /// 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 +92,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 +103,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 +140,38 @@ 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); + try launch_.launch(f.part, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = @intCast(t) }, .block = .{ .x = 256 } }, 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 +186,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 = rows_smem[0] }; + try launch_.launch(f.prompt[@intFromBool(w.p32)], cfg, s, &a); + } }; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..febd12c1c2 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -29,6 +29,7 @@ const round_k = @import("contract_tree_round.zig"); const act_k = @import("contract_act.zig"); const gdn_k = @import("contract_gdn.zig"); const qk_k = @import("contract_qk.zig"); +const wrows_k = @import("contract_window_rows.zig"); const contract = core.contract; const g = contract.gates; @@ -47,6 +48,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 &.{}; @@ -268,6 +270,12 @@ 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 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..16831af820 --- /dev/null +++ b/zig/tests/cuda/contract_window_rows.zig @@ -0,0 +1,296 @@ +//! window_rows.cu under the contract on CUDA: Gemma 4's sliding layers (head dim 256), every output 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 }; + +/// 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 } }; + +/// 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 }, +}; + +/// 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", + .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 => 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 shape.rows + f.win > f.ring) return false; // a chunk and its window fit the ring + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const n = shape.rows; + const crow = if (f.prompt) f.ring 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) shape.k -| 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(shape.k + 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 = shape.k, .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/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..bc6ae9da31 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -62,6 +62,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"); @@ -181,6 +183,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) @@ -361,6 +365,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/window_bench.zig b/zig/tests/cuda/window_bench.zig new file mode 100644 index 0000000000..3c6b22bcae --- /dev/null +++ b/zig/tests/cuda/window_bench.zig @@ -0,0 +1,121 @@ +//! 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 layers (head dim 256, window 1024, ring 1536): decode rows and 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; + } + } + 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..7d6044b0ac --- /dev/null +++ b/zig/tests/cuda/window_rows_tests.zig @@ -0,0 +1,217 @@ +//! 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 }; + +/// Decode at 1, 3, 17, 64 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, 17, 64 }; + const most = counts[counts.len - 1]; + 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) { // the linear cache declared 16x longer: 16x the chunks, 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 }, +}; + +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) |c| { + _ = arena.reset(.retain_capacity); + try decode1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); + } + for (sliding_prompt) |c| { + _ = arena.reset(.retain_capacity); + try prompt1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); + } + report("window rows (all)", total); +} From 742480004f82838962f1d284811c01e5b6ef0a51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:21:12 +0200 Subject: [PATCH 2/2] cuda window_attention: head dim 512 over a linear cache, for global attention window_rows.cu at head dim 512: decode partials with K and V tiles in shared memory, a block for 4 query heads and a warp for each key run, merged in chunk order as at 256; prompt tiles as at 256. No window over a linear cache to 262144 keys, GQA 8 and 16. --- docs/kernel-contract.md | 7 +- zig/kernels/cuda/contract_controls.cu | 10 +- zig/kernels/cuda/window_rows.cu | 36 +++- zig/kernels/cuda/window_rows512.cuh | 181 ++++++++++++++++++ zig/src/core/contract/window_rows_records.zig | 6 +- zig/src/cuda/window_attention.zig | 22 ++- zig/tests/cuda/contract_window_rows.zig | 27 ++- zig/tests/cuda/window_bench.zig | 15 +- zig/tests/cuda/window_rows_tests.zig | 28 ++- 9 files changed, 300 insertions(+), 32 deletions(-) create mode 100644 zig/kernels/cuda/window_rows512.cuh diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 71bfdecc89..aaf018e491 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -112,9 +112,10 @@ bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles w 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), the -reduce-scatter's `tf_rank_sum`, and `window_attention`'s rows at head dim 256 (`window_rows.cu`: decode chunks and their -merge, prompt tiles with P as bf16 or in two halves, a row's last key), each output within float64's range over the -spread its fp32 steps and tensor-core sums allow (`core/attention_range.zig`; six controls, each one planted fault). +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). ## CUDA records per nvcc release diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 39fc742438..4582af4d3b 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -86,7 +86,7 @@ extern "C" __global__ void tf_ctl_heavy(const float* X, float* Y) { Y[t] = s; } -// Gate 3: window_rows.cu's kernels at head dim 256, each with one planted fault (window_rows.cuh's Fault). +// 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, @@ -125,3 +125,11 @@ extern "C" __global__ void __launch_bounds__(tf_wrows::NT, 1) tf_ctl_wrows_promp 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); +} diff --git a/zig/kernels/cuda/window_rows.cu b/zig/kernels/cuda/window_rows.cu index aa31b1bf4a..46887ad334 100644 --- a/zig/kernels/cuda/window_rows.cu +++ b/zig/kernels/cuda/window_rows.cu @@ -1,6 +1,7 @@ -// window_attention's rows at head dim 256 (window_rows.cuh): decode partials and merge, prompt tiles. +// 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; @@ -36,3 +37,36 @@ extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256_p32(cons 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_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/contract/window_rows_records.zig b/zig/src/core/contract/window_rows_records.zig index 6ce2a8aebd..c4c9e0db06 100644 --- a/zig/src/core/contract/window_rows_records.zig +++ b/zig/src/core/contract/window_rows_records.zig @@ -1,4 +1,4 @@ -//! window_rows.cu's launches on CUDA (window attention at head dim 256), read by kernel-contract-limits from every SM. +//! 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. @@ -7,4 +7,8 @@ pub const cuda = [_]registry.CudaRecord{ .{ .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/cuda/window_attention.zig b/zig/src/cuda/window_attention.zig index c3e180194b..7a17eb5acf 100644 --- a/zig/src/cuda/window_attention.zig +++ b/zig/src/cuda/window_attention.zig @@ -15,13 +15,15 @@ pub const Window = struct { h: usize, hk: usize, d: usize, win: usize, ring: usi /// 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; + 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; -/// Prompt rows' dynamic shared bytes at 256: 4 slots of 32 keys, then the warps' partial scores. -const rows_smem = [_]u32{4 * 32 * 32 * 16 + 8 * 16 * 32 * 4}; +/// 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"; @@ -53,9 +55,10 @@ pub const WindowAttention = struct { var rows = try Module.load(d, kernels.window_rows); errdefer rows.unload(); var wf: [rows_smem.len]Wide = undefined; - inline for (.{256}, &wf, rows_smem) |dim, *f, smem| { + 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"); @@ -71,13 +74,11 @@ pub const WindowAttention = struct { /// The dynamic shared bytes window_rows.cu's prompt rows take at head dim d. pub fn promptShared(d: usize) u32 { - _ = d; - return rows_smem[0]; + return rows_smem[@intFromBool(d == 512)]; } pub fn wideFns(g: *const WindowAttention, d: usize) *const Wide { - _ = d; - return &g.wide_fns[0]; + 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. @@ -156,7 +157,8 @@ pub const WindowAttention = struct { a.add(@as(i32, @intCast(w.win))); a.add(@as(i64, @intCast(w.ring))); a.add(w.scale); - try launch_.launch(f.part, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = @intCast(t) }, .block = .{ .x = 256 } }, s, &a); + 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))); @@ -207,7 +209,7 @@ pub const WindowAttention = struct { 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 = rows_smem[0] }; + 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/tests/cuda/contract_window_rows.zig b/zig/tests/cuda/contract_window_rows.zig index 16831af820..53520bc5b1 100644 --- a/zig/tests/cuda/contract_window_rows.zig +++ b/zig/tests/cuda/contract_window_rows.zig @@ -1,4 +1,4 @@ -//! window_rows.cu under the contract on CUDA: Gemma 4's sliding layers (head dim 256), every output in float64's range. +//! 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"); @@ -16,7 +16,7 @@ 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 }; +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 { @@ -36,6 +36,9 @@ pub const Cfg = struct { 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{ @@ -48,6 +51,11 @@ pub const cfgs = [_]Cfg{ .{ .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. @@ -56,6 +64,7 @@ pub fn symbol(ctl: Ctl) ?[:0]const u8 { .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", @@ -112,7 +121,7 @@ pub const Rows = struct { const fun = try self.ctl.?.function(sym); switch (self.f.ctl) { .edge => k.part = fun, - .rescale, .trunc => k.merge = fun, + .rescale, .trunc, .rescale512 => k.merge = fun, else => { try fun.allowDynamicShared(WA.promptShared(self.f.d)); k.prompt[1] = fun; @@ -126,12 +135,14 @@ pub const Rows = struct { _ = 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 shape.rows + f.win > f.ring) return false; // a chunk and its window fit the ring + 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; - const crow = if (f.prompt) f.ring else f.streams * f.ring; + // 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); @@ -141,14 +152,14 @@ pub const Rows = struct { 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) shape.k -| 1 else if (i % 4 == 0) i % @min(40, shape.k) else (i * 7919) % shape.k); + 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(shape.k + i); + 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 = shape.k, .hi = if (f.image) hi else null }; + 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) { diff --git a/zig/tests/cuda/window_bench.zig b/zig/tests/cuda/window_bench.zig index 3c6b22bcae..f4b7fa6441 100644 --- a/zig/tests/cuda/window_bench.zig +++ b/zig/tests/cuda/window_bench.zig @@ -25,7 +25,7 @@ pub const callers = [_]Shape{ const contexts = [_]usize{ 1024, 8192, 32768, 131072 }; -/// Gemma 4's sliding layers (head dim 256, window 1024, ring 1536): decode rows and 512-row prompt chunks. +/// 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| { @@ -39,6 +39,19 @@ pub fn sliding(out: []Shape) []Shape { 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]; } diff --git a/zig/tests/cuda/window_rows_tests.zig b/zig/tests/cuda/window_rows_tests.zig index 7d6044b0ac..d66591ad77 100644 --- a/zig/tests/cuda/window_rows_tests.zig +++ b/zig/tests/cuda/window_rows_tests.zig @@ -63,12 +63,12 @@ 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 }; +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, 17, 64 rows: every output in range, a row's bits at any row count and (linear) cache length. +/// 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, 17, 64 }; - const most = counts[counts.len - 1]; + 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); @@ -110,7 +110,7 @@ fn decode1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const wa.Wi 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) { // the linear cache declared 16x longer: 16x the chunks, the same bits + 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); @@ -188,6 +188,20 @@ pub const sliding_decode = [_]DCase{ .{ .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 }, @@ -205,11 +219,11 @@ pub fn run(gpu: check.Gpu) !void { defer g.unload(); var prng = std.Random.DefaultPrng.init(0x6a77); var total: ar.Stats = .{}; - for (sliding_decode) |c| { + for (sliding_decode ++ global_decode) |c| { _ = arena.reset(.retain_capacity); try decode1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); } - for (sliding_prompt) |c| { + for (sliding_prompt ++ global_prompt) |c| { _ = arena.reset(.retain_capacity); try prompt1(gpu, arena.allocator(), s, &g, prng.random(), c, &total); }