From 751e7f82a870238067c0f233d61ff5c091b65e6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:57:56 +0200 Subject: [PATCH 01/41] cuda: tree attention and DFlash block attention for every family (tree_attention.cu, block_attention.cu), checked against core/attention_ref.zig --- zig/build/cuda.zig | 2 + zig/kernels/cuda/block_attention.cu | 203 ++++++++++++++++ zig/kernels/cuda/tree_attention.cu | 344 ++++++++++++++++++++++++++++ zig/src/core/attention_ref.zig | 38 +++ zig/src/core/root.zig | 2 + zig/src/cuda/block_attention.zig | 48 ++++ zig/src/cuda/kernels.zig | 2 + zig/src/cuda/root.zig | 4 + zig/src/cuda/tree_attention.zig | 100 ++++++++ zig/tests/cuda/block_tests.zig | 152 ++++++++++++ zig/tests/cuda/main.zig | 6 + zig/tests/cuda/tree_tests.zig | 201 ++++++++++++++++ 12 files changed, 1102 insertions(+) create mode 100644 zig/kernels/cuda/block_attention.cu create mode 100644 zig/kernels/cuda/tree_attention.cu create mode 100644 zig/src/core/attention_ref.zig create mode 100644 zig/src/cuda/block_attention.zig create mode 100644 zig/src/cuda/tree_attention.zig create mode 100644 zig/tests/cuda/block_tests.zig create mode 100644 zig/tests/cuda/tree_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..4ca82f73c0 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -47,6 +47,8 @@ const kernels = [_]Kernel{ .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 + .{ .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 = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu new file mode 100644 index 0000000000..e3cd06b121 --- /dev/null +++ b/zig/kernels/cuda/block_attention.cu @@ -0,0 +1,203 @@ +// DFlash-style block attention for every CUDA family (float64 reference: core/attention_ref.zig): context, then block. + +// Block (stream, key head): a warp a query head, its 16 mma rows the block's rows; key tiles of 32 shared by the warps. + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_block { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +constexpr int D = 128, C = D / 8, K16 = D / 16, NB = D / 8, KT = 32; +constexpr float LOG2E = 1.4426950408889634f; + +__device__ __forceinline__ void wait_all() { asm volatile("cp.async.wait_group 0;\n" ::); } + +__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))); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +__device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); } + +// One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V. +struct Fold { + float o[NB][4]; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f; + + __device__ __forceinline__ Fold() { +#pragma unroll + for (int i = 0; i < NB; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f; + } + + template + __device__ __forceinline__ void tile(const uint32_t (&qa)[K16][4], const uint4* ks, const uint4* vs, float scale, + int lane, int k0, Vis vis) { + const int g8 = lane / 4, qd = lane % 4, mt = lane / 8; + float sc[4][4]; +#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, ks + sw(8 * (n + mt / 2) + lane % 8, 2 * k + mt % 2)); + if (k == 0) { + mma0(sc[n], qa[k], b[0], b[1]); + mma0(sc[n + 1], qa[k], b[2], b[3]); + } else { + mma(sc[n], qa[k], b[0], b[1]); + mma(sc[n + 1], qa[k], b[2], b[3]); + } + } + } + float tm0 = -INFINITY, tm1 = -INFINITY; +#pragma unroll + for (int n = 0; n < 4; ++n) { + const int c = k0 + 8 * n + 2 * qd; + sc[n][0] = vis(g8, c) ? __fmul_rn(sc[n][0], scale) : -INFINITY; + sc[n][1] = vis(g8, c + 1) ? __fmul_rn(sc[n][1], scale) : -INFINITY; + sc[n][2] = vis(g8 + 8, c) ? __fmul_rn(sc[n][2], scale) : -INFINITY; + sc[n][3] = vis(g8 + 8, c + 1) ? __fmul_rn(sc[n][3], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1])); + tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); + } + const float n0 = fmaxf(m0, tm0), n1 = fmaxf(m1, tm1); + const float a0 = tm0 == -INFINITY ? 1.f : (m0 == -INFINITY ? 0.f : texp(__fsub_rn(m0, n0))); + const float a1 = tm1 == -INFINITY ? 1.f : (m1 == -INFINITY ? 0.f : texp(__fsub_rn(m1, n1))); + float c0 = 0.f, c1 = 0.f; +#pragma unroll + for (int n = 0; n < 4; ++n) { + sc[n][0] = sc[n][0] != -INFINITY ? texp(__fsub_rn(sc[n][0], n0)) : 0.f; + sc[n][1] = sc[n][1] != -INFINITY ? texp(__fsub_rn(sc[n][1], n0)) : 0.f; + sc[n][2] = sc[n][2] != -INFINITY ? texp(__fsub_rn(sc[n][2], n1)) : 0.f; + sc[n][3] = sc[n][3] != -INFINITY ? texp(__fsub_rn(sc[n][3], n1)) : 0.f; + c0 = __fadd_rn(c0, __fadd_rn(sc[n][0], sc[n][1])); + c1 = __fadd_rn(c1, __fadd_rn(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + c0 = __fadd_rn(c0, __shfl_xor_sync(0xffffffffu, c0, x)); + c1 = __fadd_rn(c1, __shfl_xor_sync(0xffffffffu, c1, x)); + } + l0 = __fmaf_rn(l0, a0, c0); + l1 = __fmaf_rn(l1, a1, c1); + m0 = n0; + m1 = n1; +#pragma unroll + for (int i = 0; i < NB; ++i) { + o[i][0] = __fmul_rn(o[i][0], a0); + o[i][1] = __fmul_rn(o[i][1], a0); + o[i][2] = __fmul_rn(o[i][2], a1); + o[i][3] = __fmul_rn(o[i][3], a1); + } + uint32_t pa[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) { + pa[j][0] = pack(sc[2 * j][0], sc[2 * j][1]); + pa[j][1] = pack(sc[2 * j][2], sc[2 * j][3]); + pa[j][2] = pack(sc[2 * j + 1][0], sc[2 * j + 1][1]); + pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); + } +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, vs + sw(16 * j + 8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa[j], b[0], b[1]); + mma(o[i + 1], pa[j], b[2], b[3]); + } + } +}; + +} // namespace tf_block + +using namespace tf_block; + +// Block rows r see context key k when lens[j] + r - k <= window, then block keys c < L (c <= r if causal). +extern "C" __global__ void __launch_bounds__(256) tf_block_attention( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, const __nv_bfloat16* __restrict__ vb, + const long long* __restrict__ table, const int* __restrict__ lens, __nv_bfloat16* __restrict__ out, float scale, + int window, int rows, int h, int hk, int L, int causal) { + extern __shared__ uint4 smem[]; + uint4* ks = smem; + uint4* vs = ks + KT * C; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g8 = lane / 4, qd = lane % 4; + const int j = blockIdx.x, g = blockIdx.y, G = h / hk; + const int head = g * G + warp; + const int s = lens[j]; + const __nv_bfloat16* kc = reinterpret_cast(table[2 * j]) + (int64_t)g * s * D; + const __nv_bfloat16* vc = reinterpret_cast(table[2 * j + 1]) + (int64_t)g * s * D; + uint32_t qa[K16][4]; + { + const int r0 = g8, r1 = g8 + 8; + const uint32_t* q0 = reinterpret_cast(q + ((int64_t)head * rows + j * L + min(r0, L - 1)) * D); + const uint32_t* q1 = reinterpret_cast(q + ((int64_t)head * rows + j * L + min(r1, L - 1)) * D); +#pragma unroll + for (int k = 0; k < K16; ++k) { + qa[k][0] = r0 < L ? q0[8 * k + qd] : 0u; + qa[k][1] = r1 < L ? q1[8 * k + qd] : 0u; + qa[k][2] = r0 < L ? q0[8 * k + qd + 4] : 0u; + qa[k][3] = r1 < L ? q1[8 * k + qd + 4] : 0u; + } + } + Fold f; + const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) + for (int k0 = lo / KT * KT; k0 < s; k0 += KT) { + for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { + const int key = k0 + i / C; + const bool in = key < s; + cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + } + commit(); + wait_all(); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, k0, [&](int r, int key) { return r < L && key < s && s + r - key <= window; }); + __syncthreads(); + } + for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { // the block's own keys + const int c = i / C; + const bool in = c < L; + const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; + cp16z(ks + sw(c, i % C), kb + at, in); + cp16z(vs + sw(c, i % C), vb + at, in); + } + commit(); + wait_all(); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, 0, [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }); + if (g8 < L) +#pragma unroll + for (int i = 0; i < NB; ++i) + *reinterpret_cast(out + ((int64_t)(j * L + g8) * h + head) * D + 8 * i + 2 * qd) = + pack(__fdiv_rn(f.o[i][0], f.l0), __fdiv_rn(f.o[i][1], f.l0)); + if (g8 + 8 < L) +#pragma unroll + for (int i = 0; i < NB; ++i) + *reinterpret_cast(out + ((int64_t)(j * L + g8 + 8) * h + head) * D + 8 * i + 2 * qd) = + pack(__fdiv_rn(f.o[i][2], f.l1), __fdiv_rn(f.o[i][3], f.l1)); +} diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu new file mode 100644 index 0000000000..035b1d642f --- /dev/null +++ b/zig/kernels/cuda/tree_attention.cu @@ -0,0 +1,344 @@ +// Tree attention (float64 references: core/attention_ref.zig): committed keys, then the row's path, groups in order. + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_tree { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +constexpr int SPAN = 2048; // keys a group: one partial +constexpr int MAXD = 128; // a path's most rows +constexpr float LOG2E = 1.4426950408889634f; + +template +__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__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))); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +template +__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } + +} // namespace tf_tree + +using namespace tf_tree; + +// Each row's root-to-row path (global window rows, root first) and its depth, from parents (-1 at a root). +extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* __restrict__ paths, + int* __restrict__ depths, int rows) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= rows) return; + int depth = 0; + for (int cur = node; cur >= 0 && depth < MAXD; cur = parents[cur]) ++depth; + depths[node] = depth; + int cur = node; + for (int slot = depth - 1; slot >= 0; --slot, cur = parents[cur]) paths[node * MAXD + slot] = cur; +} + +// Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. +template +__device__ __forceinline__ void tree_group( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, + const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, + const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, + float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int gi, int ng, + int64_t rq, int pf, int npairs, int64_t row) { + constexpr int C = D / 8, K16 = D / 16, NB = D / 8, NT8 = HB / 8, PK = HB / 16; + constexpr bool QREG = D <= 128; + extern __shared__ uint4 smem[]; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g8 = lane / 4, qd = lane % 4, mt = lane / 8; + const int64_t r = row; // the row whose stream and depth set the keys + const int kh = blockIdx.y, G = h / hk; + const int s = row_stream[r]; + const int p = streams[4 * s + 2], depth = depths[r]; + const int total = p + depth, k0 = gi * SPAN, k1 = min(total, k0 + SPAN); + if (k0 >= total) return; // the merge reads only the row's own groups + const __nv_bfloat16* kc = origin + offs[2 * s]; + const __nv_bfloat16* vc = origin + offs[2 * s + 1]; + uint4* ks = smem + warp * 2 * HB * C; + uint4* vs = ks + HB * C; + uint32_t qa[QREG ? K16 : 1][4]; + uint4* qs = smem + 4 * 2 * HB * C; // !QREG: the row's heads [16][D], shared by the warps + if constexpr (QREG) { + const int a0 = pf + g8, a1 = pf + g8 + 8; + const bool v0 = a0 < npairs, v1 = a1 < npairs; + const uint32_t* q0 = reinterpret_cast(q + ((rq + (v0 ? a0 : pf) / G) * h + kh * G + (v0 ? a0 : pf) % G) * D); + const uint32_t* q1 = reinterpret_cast(q + ((rq + (v1 ? a1 : pf) / G) * h + kh * G + (v1 ? a1 : pf) % G) * D); +#pragma unroll + for (int k = 0; k < K16; ++k) { + qa[k][0] = v0 ? q0[8 * k + qd] : 0u; + qa[k][1] = v1 ? q1[8 * k + qd] : 0u; + qa[k][2] = v0 ? q0[8 * k + qd + 4] : 0u; + qa[k][3] = v1 ? q1[8 * k + qd + 4] : 0u; + } + } else { + for (int i = threadIdx.x; i < 16 * C; i += blockDim.x) { + const int a = pf + i / C; + const bool in = a < npairs; + const int64_t row = rq + (in ? a : pf) / G; + cp16z(qs + sw(i / C, i % C), q + (row * h + kh * G + (in ? a : pf) % G) * D + (i % C) * 8, in); + } + commit(); + wait_group<0>(); + __syncthreads(); + } + 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 = k0 / HB + warp; t * HB < k1; t += 4) { + for (int i = lane; i < HB * C; i += 32) { + const int key = t * HB + i / C; + const bool in = key < k1; + const __nv_bfloat16* ksrc; + const __nv_bfloat16* vsrc; + if (key < p) { + ksrc = kc + ((int64_t)key * hk + kh) * D; + vsrc = vc + ((int64_t)key * hk + kh) * D; + } else { // the row's own path: key p + j is path row j + const int node = in ? paths[r * MAXD + (key - p)] : 0; + ksrc = kn + ((int64_t)node * hk + kh) * D; + vsrc = vn + ((int64_t)node * hk + kh) * D; + } + if (!in) ksrc = kn, vsrc = vn; + cp16z(ks + sw(i / C, i % C), ksrc + (i % C) * 8, in); + cp16z(vs + sw(i / C, i % C), vsrc + (i % C) * 8, in); + } + commit(); + wait_group<0>(); + __syncwarp(); + float sc[NT8][4]; +#pragma unroll + for (int k = 0; k < K16; ++k) { + uint32_t qf[4]; + if constexpr (QREG) { + qf[0] = qa[k][0], qf[1] = qa[k][1], qf[2] = qa[k][2], qf[3] = qa[k][3]; + } else { + ldmatrix4(qf, qs + sw(lane % 16, 2 * k + lane / 16)); + } +#pragma unroll + for (int n = 0; n < NT8; n += 2) { + uint32_t b[4]; + ldmatrix4(b, ks + sw(8 * (n + mt / 2) + lane % 8, 2 * k + mt % 2)); + if (k == 0) { + mma0(sc[n], qf, b[0], b[1]); + mma0(sc[n + 1], qf, b[2], b[3]); + } else { + mma(sc[n], qf, b[0], b[1]); + mma(sc[n + 1], qf, b[2], b[3]); + } + } + } + const int kb = t * HB + 2 * qd; + float tm0 = -INFINITY, tm1 = -INFINITY; +#pragma unroll + for (int n = 0; n < NT8; ++n) { + const int c = kb + 8 * n; + const bool i0 = c < k1, i1 = c + 1 < k1; + sc[n][0] = i0 ? __fmul_rn(sc[n][0], scale) : -INFINITY; + sc[n][1] = i1 ? __fmul_rn(sc[n][1], scale) : -INFINITY; + sc[n][2] = i0 ? __fmul_rn(sc[n][2], scale) : -INFINITY; + sc[n][3] = i1 ? __fmul_rn(sc[n][3], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1])); + tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); + } + const bool act = tm0 != -INFINITY; + const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; + const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + float c0 = 0.0f, c1 = 0.0f; +#pragma unroll + for (int n = 0; n < NT8; ++n) { + sc[n][0] = sc[n][0] != -INFINITY ? texp(__fsub_rn(sc[n][0], n0)) : 0.0f; + sc[n][1] = sc[n][1] != -INFINITY ? texp(__fsub_rn(sc[n][1], n0)) : 0.0f; + sc[n][2] = sc[n][2] != -INFINITY ? texp(__fsub_rn(sc[n][2], n1)) : 0.0f; + sc[n][3] = sc[n][3] != -INFINITY ? texp(__fsub_rn(sc[n][3], n1)) : 0.0f; + c0 = __fadd_rn(c0, __fadd_rn(sc[n][0], sc[n][1])); + c1 = __fadd_rn(c1, __fadd_rn(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + c0 = __fadd_rn(c0, __shfl_xor_sync(0xffffffffu, c0, x)); + c1 = __fadd_rn(c1, __shfl_xor_sync(0xffffffffu, c1, x)); + } + l0 = __fmaf_rn(l0, a0, c0); + l1 = __fmaf_rn(l1, a1, c1); + m0 = n0; + m1 = n1; +#pragma unroll + for (int i = 0; i < NB; ++i) { + o[i][0] = __fmul_rn(o[i][0], a0); + o[i][1] = __fmul_rn(o[i][1], a0); + o[i][2] = __fmul_rn(o[i][2], a1); + o[i][3] = __fmul_rn(o[i][3], a1); + } + uint32_t pa[PK][4]; +#pragma unroll + for (int j = 0; j < PK; ++j) { + pa[j][0] = pack(sc[2 * j][0], sc[2 * j][1]); + pa[j][1] = pack(sc[2 * j][2], sc[2 * j][3]); + pa[j][2] = pack(sc[2 * j + 1][0], sc[2 * j + 1][1]); + pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); + } +#pragma unroll + for (int j = 0; j < PK; ++j) +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, vs + sw(16 * j + 8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa[j], b[0], b[1]); + mma(o[i + 1], pa[j], b[2], b[3]); + } + __syncwarp(); + } + __syncthreads(); + float* os = reinterpret_cast(smem); // [4 warps][16 heads][D], then max and sum [4][16] + float* ms = os + 4 * 16 * D; + float* ls = ms + 4 * 16; +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + os[(warp * 16 + g8) * D + d] = o[i][0]; + os[(warp * 16 + g8) * D + d + 1] = o[i][1]; + os[(warp * 16 + g8 + 8) * D + d] = o[i][2]; + os[(warp * 16 + g8 + 8) * D + d + 1] = o[i][3]; + } + if (qd == 0) { + ms[warp * 16 + g8] = m0, ls[warp * 16 + g8] = l0; + ms[warp * 16 + g8 + 8] = m1, ls[warp * 16 + g8 + 8] = l1; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * D; idx += 128) { + const int hd = idx / D, e = idx % D; + if (pf + hd >= npairs) continue; + const int64_t part = ((rq + (pf + hd) / G) * hk + kh) * ng + gi; + const int head = (pf + hd) % G; + float mm = -INFINITY; + for (int w = 0; w < 4; ++w) mm = fmaxf(mm, ms[w * 16 + hd]); + float acc = 0.0f, l = 0.0f; + for (int w = 0; w < 4; ++w) { + const float mw = ms[w * 16 + hd]; + if (mw == -INFINITY) continue; + const float a = texp(__fsub_rn(mw, mm)); + acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); + l = __fmaf_rn(a, ls[w * 16 + hd], l); + } + po[(part * 16 + head) * D + e] = acc; + if (e == 0) pm[part * 16 + head] = mm, pl[part * 16 + head] = l; + } +} + +// A stream's full committed groups once for its rows, 16 (row, head) pairs a block: each partial a row block's bits. +template +__device__ __forceinline__ void tree_shared( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, + const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, + const int* __restrict__ depths, float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, + float scale, int ng, int n) { + const int G = h / hk; + for (int item = blockIdx.x; item < n * ng; item += gridDim.x) { // (row, group): a fixed grid, no empty blocks + const int64_t r = item / ng; + const int gi = item % ng; + const int st = row_stream[r]; + const int64_t r0 = streams[4 * st]; + const int pairs = streams[4 * st + 1] * G, p = streams[4 * st + 2]; + const int tile = static_cast(r - r0); // a stream's j-th row hosts its j-th pair tile (tiles <= rows) + if (tile * 16 >= pairs || (gi + 1) * SPAN > p) continue; + tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, r0, + tile * 16, pairs, r); + __syncthreads(); + } +} + +#define TF_TREE_SHARED(NAME, D, HB) \ + extern "C" __global__ void __launch_bounds__(128) NAME( \ + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, \ + const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ + const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ + float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { \ + tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ + } + +TF_TREE_SHARED(tf_tree_shared, 128, 32) +TF_TREE_SHARED(tf_tree_shared_256, 256, 16) + +// Block (row, key head, group slot): the row's groups past its stream's full committed ones, slot, slot + slots, ... +template +__device__ __forceinline__ void tree_part( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, + const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, + const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, + float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng) { + const int64_t r = blockIdx.x; + const int p = streams[4 * row_stream[r] + 2], total = p + depths[r]; + for (int gi = blockIdx.z + p / SPAN; gi < ng && gi * SPAN < total; gi += gridDim.z) { // full committed: shared + tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, + r, 0, h / hk, r); + __syncthreads(); // the next group's tiles reuse the partials' shared memory + } +} + +#define TF_TREE_PART(NAME, D, HB) \ + extern "C" __global__ void __launch_bounds__(128) NAME( \ + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, \ + const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, \ + const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ + const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ + float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng) { \ + tree_part(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng); \ + } + +TF_TREE_PART(tf_tree_part, 128, 32) +TF_TREE_PART(tf_tree_part_256, 256, 16) + +// out [W, H, D] bf16: block (row, key head) folds the row's groups in order. +extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( + const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, + const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ depths, + __nv_bfloat16* __restrict__ out, int h, int hk, int ng, int d) { + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, G = h / hk; + const int total = streams[4 * row_stream[r] + 2] + depths[r]; + const int groups = (total + SPAN - 1) / SPAN; + for (int idx = threadIdx.x; idx < G * d; idx += 256) { + const int hd = idx / d, e = idx % d; + float m = -INFINITY, l = 0.0f, acc = 0.0f; + for (int gi = 0; gi < groups; ++gi) { + const int64_t part = (r * hk + kh) * ng + gi; + const float cm = pm[part * 16 + hd], cl = pl[part * 16 + hd]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : texp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? texp(__fsub_rn(cm, next)) : 0.0f; + acc = __fmaf_rn(acc, a, __fmul_rn(po[(part * 16 + hd) * d + e], b)); + l = __fmaf_rn(l, a, __fmul_rn(cl, b)); + m = next; + } + out[(r * h + kh * G + hd) * d + e] = __float2bfloat16_rn(__fdiv_rn(acc, l)); + } +} diff --git a/zig/src/core/attention_ref.zig b/zig/src/core/attention_ref.zig new file mode 100644 index 0000000000..b9f07beb2e --- /dev/null +++ b/zig/src/core/attention_ref.zig @@ -0,0 +1,38 @@ +//! Attention of one query head in float64, the reference for every backend's attention kernels. +const std = @import("std"); + +/// softmax(scale q . k_j) v over `keys`; tol 2^-7 of sum p_j |v_j| / sum p_j (bf16 weights) + a bf16 half-ulp. +pub fn head(q: []const f64, k: []const f64, v: []const f64, keys: []const usize, d: usize, scale: f64, out: []f64, tol: []f64) void { + var m: f64 = -std.math.inf(f64); + for (keys) |j| { + var s: f64 = 0; + for (0..d) |e| s += q[e] * k[j * d + e]; + m = @max(m, s * scale); + } + var l: f64 = 0; + @memset(out, 0); + @memset(tol, 0); + for (keys) |j| { + var s: f64 = 0; + for (0..d) |e| s += q[e] * k[j * d + e]; + const p = @exp(s * scale - m); + l += p; + for (0..d) |e| { + out[e] += p * v[j * d + e]; + tol[e] += p * @abs(v[j * d + e]); + } + } + for (out, tol) |*o, *t| { + o.* /= l; + t.* = t.* / l * std.math.pow(f64, 2, -7) + @abs(o.*) * std.math.pow(f64, 2, -8) + 1e-6; + } +} + +test "one key returns its value; equal scores average" { + var out: [2]f64 = undefined; + var tol: [2]f64 = undefined; + head(&.{ 1, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, 4, 6, 8 }, &.{1}, 2, 1, &out, &tol); + try std.testing.expectEqualSlices(f64, &.{ 6, 8 }, &out); + head(&.{ 0, 0 }, &.{ 1, 0, 1, 0 }, &.{ 2, 4, 6, 8 }, &.{ 0, 1 }, 2, 1, &out, &tol); + try std.testing.expectEqualSlices(f64, &.{ 4, 6 }, &out); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 122ffe9e4c..26cdc72a62 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -28,6 +28,7 @@ pub const modelopt = @import("modelopt.zig"); pub const picture = @import("picture/picture.zig"); /// Qwen-VL image input every Qwen3.5-VL checkpoint shares: settings, geometry, M-RoPE positions, the tower's tables. pub const qwen_vl = @import("qwen_vl/qwen_vl.zig"); +pub const attention_ref = @import("attention_ref.zig"); test { _ = safetensors; @@ -41,6 +42,7 @@ test { _ = norms_ref; _ = act_ref; _ = contract; + _ = attention_ref; _ = exl3_format; _ = exlr8; _ = eaq_container; diff --git a/zig/src/cuda/block_attention.zig b/zig/src/cuda/block_attention.zig new file mode 100644 index 0000000000..fd63c57d42 --- /dev/null +++ b/zig/src/cuda/block_attention.zig @@ -0,0 +1,48 @@ +//! block_attention.cu for every CUDA family: a DFlash-style drafter's block rows over their context and the block. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// A 32-key K and V tile (head dim 128), shared by the block's warps. +const smem: u32 = 2 * 32 * 16 * 16; + +/// q [h, rows, 128], kb and vb [hk, rows, 128] (head-major); stream j's block rows j L .. j L + L - 1. +pub const Shape = struct { h: usize, hk: usize, rows: usize, block: usize, window: usize, scale: f32, causal: bool = false }; + +pub const BlockAttention = struct { + mod: Module, + fn_: Function, + + pub fn load(d: *const Driver) !BlockAttention { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.block_attention); + errdefer mod.unload(); + const f = try mod.function("tf_block_attention"); + try f.allowDynamicShared(smem); + return .{ .mod = mod, .fn_ = f }; + } + + pub fn unload(g: *BlockAttention) void { + g.mod.unload(); + } + + /// out [rows, h, 128]; table i64 [streams, 2]: stream j's context K and V ([hk, lens[j], 128] bf16); lens i32. + pub fn run(g: *const BlockAttention, s: Stream, q: u64, kb: u64, vb: u64, table: u64, lens: u64, out: u64, sh: Shape) !void { + if (sh.hk == 0 or sh.h % sh.hk != 0 or sh.h / sh.hk > 8 or sh.block == 0 or sh.block > 16 or sh.rows % sh.block != 0) return error.BlockShape; + if (sh.rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ q, kb, vb, table, lens, out }) |v| a.add(v); + a.add(sh.scale); + for ([_]usize{ sh.window, sh.rows, sh.h, sh.hk, sh.block }) |v| a.add(@as(i32, @intCast(v))); + a.add(@as(i32, @intFromBool(sh.causal))); + const grid: launch_.Dim3 = .{ .x = @intCast(sh.rows / sh.block), .y = @intCast(sh.hk), .z = 1 }; + try launch_.launch(g.fn_, .{ .grid = grid, .block = .{ .x = @intCast(32 * (sh.h / sh.hk)) }, .shared = smem }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..7e0dd2b10a 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -35,6 +35,8 @@ pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; +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 modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6c4e0991b6..f48dd67738 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -41,6 +41,8 @@ pub const tp_group = @import("tp_group.zig"); pub const tp_watch = @import("tp_watch.zig"); pub const ck = @import("ck.zig"); pub const modelopt = @import("modelopt.zig"); +pub const tree_attention = @import("tree_attention.zig"); +pub const block_attention = @import("block_attention.zig"); test { _ = @import("memory.zig"); @@ -63,4 +65,6 @@ test { _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); + _ = @import("tree_attention.zig"); + _ = @import("block_attention.zig"); } diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig new file mode 100644 index 0000000000..dd8e56bdcc --- /dev/null +++ b/zig/src/cuda/tree_attention.zig @@ -0,0 +1,100 @@ +//! tree_attention.cu for every CUDA family: drafted rows over committed keys and their paths, groups in order. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Keys a group, a path's most rows, the part kernels' shared memory (four warps' K/V slots; at 256 also the queries). +pub const span = 2048; +pub const max_depth = 128; +const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; + +/// row_stream [n], streams [s, 4] (committed at [2]), offs [s, 2] (K, V from origin), paths [n, max_depth], depths. +pub const Rows = struct { row_stream: u64, streams: u64, offs: u64, origin: u64, paths: u64, depths: u64 }; + +pub const Scratch = struct { o: u64, m: u64, l: u64 }; + +pub const TreeAttention = struct { + mod: Module, + paths_fn: Function, + part_fns: [2]Function, // head dim 128, 256 + shared_fns: [2]Function, + merge_fn: Function, + /// The shared kernel's blocks a key head: a fixed grid walking every (row, group) (bits do not depend on it). + shared_blocks: usize = 96, + + pub fn load(d: *const Driver) !TreeAttention { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.tree_attention); + errdefer mod.unload(); + var parts: [2]Function = .{ try mod.function("tf_tree_part"), try mod.function("tf_tree_part_256") }; + for (&parts, part_smem) |*f, b| try f.allowDynamicShared(b); + var shared: [2]Function = .{ try mod.function("tf_tree_shared"), try mod.function("tf_tree_shared_256") }; + for (&shared, part_smem) |*f, b| try f.allowDynamicShared(b); + return .{ .mod = mod, .paths_fn = try mod.function("tf_tree_paths"), .part_fns = parts, .shared_fns = shared, .merge_fn = try mod.function("tf_tree_merge") }; + } + + pub fn unload(g: *TreeAttention) void { + g.mod.unload(); + } + + /// Groups a row with `keys` committed and path keys needs. + pub fn groups(keys: usize) usize { + return (keys + span - 1) / span; + } + + /// The fp32 partials of `rows` rows, `hk` key heads and `ng` groups: o [.., 16, dim], then max and sum [.., 16]. + pub fn scratchBytes(rows: usize, hk: usize, ng: usize, dim: usize) usize { + return rows * hk * ng * 16 * (dim + 2) * 4; + } + + pub fn scratch(base: u64, rows: usize, hk: usize, ng: usize, dim: usize) Scratch { + const parts = rows * hk * ng * 16; + return .{ .o = base, .m = base + parts * dim * 4, .l = base + parts * dim * 4 + parts * 4 }; + } + + /// Each row's path (root first) and depth from `parents` int32 [n] (global window rows, -1 at a root). + pub fn paths(g: *const TreeAttention, s: Stream, parents: u64, out_paths: u64, depths: u64, n: usize) !void { + if (n == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ parents, out_paths, depths }) |v| a.add(v); + a.add(@as(i32, @intCast(n))); + try launch_.launch(g.paths_fn, .{ .grid = .{ .x = @intCast((n + 127) / 128), .y = 1, .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } + + /// q [n, h, dim] (128 or 256) over committed keys and paths (kn, vn: window rows) -> out; ng >= a row's groups. + pub fn run(g: *const TreeAttention, s: Stream, q: u64, kn: u64, vn: u64, rows: Rows, sc: Scratch, out: u64, n: usize, h: usize, hk: usize, dim: usize, ng: usize, scale: f32) !void { + if (hk == 0 or h % hk != 0 or h / hk > 16 or ng == 0 or (dim != 128 and dim != 256)) return error.TreeShape; + const which: usize = if (dim == 256) 1 else 0; + if (n == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ q, kn, vn, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, sc.o, sc.m, sc.l }) |v| a.add(v); + a.add(@as(i32, @intCast(h))); + a.add(@as(i32, @intCast(hk))); + a.add(scale); + a.add(@as(i32, @intCast(ng))); + var c: launch_.Args = .{}; + for ([_]u64{ q, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, sc.o, sc.m, sc.l }) |v| c.add(v); + c.add(@as(i32, @intCast(h))); + c.add(@as(i32, @intCast(hk))); + c.add(scale); + c.add(@as(i32, @intCast(ng))); + c.add(@as(i32, @intCast(n))); + try launch_.launch(g.shared_fns[which], .{ .grid = .{ .x = @intCast(@min(n * ng, g.shared_blocks)), .y = @intCast(hk), .z = 1 }, .block = .{ .x = 128 }, .shared = part_smem[which] }, s, &c); + // a row's own block: its boundary group and at most the next (a path is at most 128 keys) + try launch_.launch(g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2)) }, .block = .{ .x = 128 }, .shared = part_smem[which] }, s, &a); + var b: launch_.Args = .{}; + for ([_]u64{ sc.o, sc.m, sc.l, rows.row_stream, rows.streams, rows.depths, out }) |v| b.add(v); + b.add(@as(i32, @intCast(h))); + b.add(@as(i32, @intCast(hk))); + b.add(@as(i32, @intCast(ng))); + b.add(@as(i32, @intCast(dim))); + try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = 1 }, .block = .{ .x = 256 } }, s, &b); + } +}; diff --git a/zig/tests/cuda/block_tests.zig b/zig/tests/cuda/block_tests.zig new file mode 100644 index 0000000000..48d07a021c --- /dev/null +++ b/zig/tests/cuda/block_tests.zig @@ -0,0 +1,152 @@ +//! block_attention.cu against core/attention_ref.zig: streams' blocks over windowed contexts, a stream's bits alone. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const D = 128; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Bufs = struct { + gpu: check.Gpu, + list: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn up(b: *Bufs, bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(b.gpu.d, bytes); + try b.list.append(b.gpu.gpa, d); + return d.ptr; + } + + fn out(b: *Bufs, s: cuda.Stream, n: usize) !u64 { + const d = try cuda.DeviceBuffer.alloc(b.gpu.d, n); + try b.list.append(b.gpu.gpa, d); + try d.fill8(0xff, s.handle); + return d.ptr; + } + + fn back(b: *Bufs, a: std.mem.Allocator, ptr: u64, n: usize) ![]u16 { + const o = try a.alloc(u16, n); + const d: cuda.DeviceBuffer = .{ .d = b.gpu.d, .ptr = ptr, .len = n * 2 }; + try d.download(0, std.mem.sliceAsBytes(o)); + return o; + } +}; + +/// Streams `which` of the case: their blocks packed head-major, their contexts; the output rows. +fn call(b: *Bufs, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.block_attention.BlockAttention, which: []const usize, q: []const u16, kb: []const u16, vb: []const u16, ctx_k: []const u64, ctx_v: []const u64, lens: []const i32, h: usize, hk: usize, L: usize, all_rows: usize, window: usize, causal: bool) ![]u16 { + const rows = which.len * L; + const pq = try a.alloc(u16, h * rows * D); + const pk = try a.alloc(u16, hk * rows * D); + const pv = try a.alloc(u16, hk * rows * D); + const table = try a.alloc(i64, 2 * which.len); + const ls = try a.alloc(i32, which.len); + for (which, 0..) |j, i| { + for (0..h) |hd| @memcpy(pq[(hd * rows + i * L) * D ..][0 .. L * D], q[(hd * all_rows + j * L) * D ..][0 .. L * D]); + for (0..hk) |hd| { + @memcpy(pk[(hd * rows + i * L) * D ..][0 .. L * D], kb[(hd * all_rows + j * L) * D ..][0 .. L * D]); + @memcpy(pv[(hd * rows + i * L) * D ..][0 .. L * D], vb[(hd * all_rows + j * L) * D ..][0 .. L * D]); + } + table[2 * i] = @intCast(ctx_k[j]); + table[2 * i + 1] = @intCast(ctx_v[j]); + ls[i] = lens[j]; + } + const o = try b.out(s, rows * h * D * 2); + const sc: f32 = @floatCast(1 / @sqrt(@as(f64, D))); + try g.run(s, try b.up(std.mem.sliceAsBytes(pq)), try b.up(std.mem.sliceAsBytes(pk)), try b.up(std.mem.sliceAsBytes(pv)), try b.up(std.mem.sliceAsBytes(table)), try b.up(std.mem.sliceAsBytes(ls)), o, .{ .h = h, .hk = hk, .rows = rows, .block = L, .window = window, .scale = sc, .causal = causal }); + try s.synchronize(); + return b.back(a, o, rows * h * D); +} + +fn case(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.block_attention.BlockAttention, rng: std.Random, L: usize, causal: bool, worst: *f64) !void { + var b: Bufs = .{ .gpu = gpu }; + defer { + for (b.list.items) |*d| d.free(); + b.list.deinit(gpu.gpa); + } + const h = 32; + const hk = 8; + const window = 2048; + const lens = [_]i32{ 5, 700, 2100 }; + const S = lens.len; + const rows = S * L; + const q = try a.alloc(u16, h * rows * D); + const kb = try a.alloc(u16, hk * rows * D); + const vb = try a.alloc(u16, hk * rows * D); + for (q) |*x| x.* = bf16(rng.floatNorm(f32)); + for (kb) |*x| x.* = bf16(rng.floatNorm(f32)); + for (vb) |*x| x.* = bf16(rng.floatNorm(f32)); + var ck: [S][]u16 = undefined; + var cv: [S][]u16 = undefined; + var pk: [S]u64 = undefined; + var pv: [S]u64 = undefined; + for (lens, 0..) |n, j| { + const m: usize = @intCast(n); + ck[j] = try a.alloc(u16, hk * m * D); + cv[j] = try a.alloc(u16, hk * m * D); + for (ck[j]) |*x| x.* = bf16(rng.floatNorm(f32)); + for (cv[j]) |*x| x.* = bf16(rng.floatNorm(f32)); + pk[j] = try b.up(std.mem.sliceAsBytes(ck[j])); + pv[j] = try b.up(std.mem.sliceAsBytes(cv[j])); + } + const all = [_]usize{ 0, 1, 2 }; + const got = try call(&b, a, s, g, &all, q, kb, vb, &pk, &pv, &lens, h, hk, L, rows, window, causal); + const scale = 1 / @sqrt(@as(f64, D)); + const kd = try a.alloc(f64, (window + 1 + L) * D); + const vd = try a.alloc(f64, kd.len); + const keys = try a.alloc(usize, window + 1 + L); + const qv = try a.alloc(f64, D); + const want = try a.alloc(f64, D); + const tol = try a.alloc(f64, D); + var most: f64 = 0; + for (0..S) |j| for (0..L) |r| for (0..h) |hd| { + const kh = hd / (h / hk); + const n: usize = @intCast(lens[j]); + var nk: usize = 0; + for (0..n) |k| if (n + r - k <= window) { + for (0..D) |e| { + kd[nk * D + e] = bf(ck[j][(kh * n + k) * D + e]); + vd[nk * D + e] = bf(cv[j][(kh * n + k) * D + e]); + } + keys[nk] = nk; + nk += 1; + }; + for (0..L) |c| if (!causal or c <= r) { + for (0..D) |e| { + kd[nk * D + e] = bf(kb[(kh * rows + j * L + c) * D + e]); + vd[nk * D + e] = bf(vb[(kh * rows + j * L + c) * D + e]); + } + keys[nk] = nk; + nk += 1; + }; + for (qv, 0..) |*x, e| x.* = bf(q[(hd * rows + j * L + r) * D + e]); + core.attention_ref.head(qv, kd, vd, keys[0..nk], D, scale, want, tol); + for (want, tol, 0..) |w, t, e| most = @max(most, @abs(bf(got[((j * L + r) * h + hd) * D + e]) - w) / t); + }; + worst.* = @max(worst.*, most); + try check.expect(most <= 1, "block attention L {d}{s}: within float64 tolerance (worst {d:.3})", .{ L, if (causal) " causal" else "", most }); + const one = [_]usize{2}; + const alone = try call(&b, a, s, g, &one, q, kb, vb, &pk, &pv, &lens, h, hk, L, rows, window, causal); + try check.sameBytes("block attention: one stream alone", std.mem.sliceAsBytes(alone), std.mem.sliceAsBytes(got[2 * L * h * D ..][0 .. L * h * D])); +} + +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 cuda.block_attention.BlockAttention.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(16); + var worst: f64 = 0; + try case(gpu, arena.allocator(), s, &g, prng.random(), 16, false, &worst); // Qwen3.8-27B's DFlash2: 32/8 heads, a 2,048 window + try case(gpu, arena.allocator(), s, &g, prng.random(), 7, true, &worst); + check.pass("block attention: windowed contexts then the block, within float64 tolerance (worst {d:.3}), a stream's bits alone", .{worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..f1e2682483 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,8 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const tree_tests = @import("tree_tests.zig"); +const block_tests = @import("block_tests.zig"); const usage = \\usage: tf-cuda-test @@ -87,6 +89,8 @@ const usage = \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage \\ exlr8-dense [DIR] EXLR8 dense formats: decode bit-equal, FP64, row windows; DIR adds the real matrices + \\ tree tree attention (committed keys and paths) against float64 references (synthetic) + \\ block DFlash block attention over windowed contexts against float64 (synthetic) \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR @@ -188,6 +192,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { } if (std.mem.eql(u8, cmd, "kernel-contract")) return contract_tests.run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); + 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, "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); if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); diff --git a/zig/tests/cuda/tree_tests.zig b/zig/tests/cuda/tree_tests.zig new file mode 100644 index 0000000000..c11088413c --- /dev/null +++ b/zig/tests/cuda/tree_tests.zig @@ -0,0 +1,201 @@ +//! tree_attention.cu against core/attention_ref.zig, and one stream's rows alone giving the same bits. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const tree = cuda.tree_attention; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Set = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn up(x: *Set, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(x.gpu.d, bytes); + try x.bufs.append(x.gpu.gpa, b); + return b.ptr; + } + + fn zero(x: *Set, s: cuda.Stream, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(x.gpu.d, @max(n, 16)); + try x.bufs.append(x.gpu.gpa, b); + try b.fill8(0xff, s.handle); + return b.ptr; + } + + fn down(x: *Set, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try x.a.alloc(T, n); + const b: cuda.DeviceBuffer = .{ .d = x.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } +}; + +/// One call over `rows` (global window rows of the given streams, in order): their outputs. +fn call(x: *Set, s: cuda.Stream, g: *const tree.TreeAttention, d: usize, rows: []const usize, row_stream_all: []const i32, parents_all: []const i32, qs: []const u16, kns: []const u16, vns: []const u16, h: usize, hk: usize, streams_all: []const i32, offs: u64, origin: u64, ng: usize, scale: f32) ![]u16 { + const n = rows.len; + var local: std.AutoHashMap(usize, i32) = .init(x.a); + for (rows, 0..) |r, i| try local.put(r, @intCast(i)); + const rs = try x.a.alloc(i32, n); + const par = try x.a.alloc(i32, n); + const q = try x.a.alloc(u16, n * h * d); + const kn = try x.a.alloc(u16, n * hk * d); + const vn = try x.a.alloc(u16, n * hk * d); + for (rows, 0..) |r, i| { + rs[i] = row_stream_all[r]; + par[i] = if (parents_all[r] < 0) -1 else local.get(@intCast(parents_all[r])).?; + @memcpy(q[i * h * d ..][0 .. h * d], qs[r * h * d ..][0 .. h * d]); + @memcpy(kn[i * hk * d ..][0 .. hk * d], kns[r * hk * d ..][0 .. hk * d]); + @memcpy(vn[i * hk * d ..][0 .. hk * d], vns[r * hk * d ..][0 .. hk * d]); + } + // the streams' table for these rows: each stream's first row and rows among them, its committed keys + const st = try x.a.dupe(i32, streams_all); + for (0..st.len / 4) |k| { + st[4 * k] = 0; + st[4 * k + 1] = 0; + } + for (rows, 0..) |r, i| { + const k: usize = @intCast(row_stream_all[r]); + if (st[4 * k + 1] == 0) st[4 * k] = @intCast(i); + st[4 * k + 1] += 1; + } + const streams_dev = try x.up(std.mem.sliceAsBytes(st)); + const paths = try x.zero(s, n * tree.max_depth * 4); + const depths = try x.zero(s, n * 4); + try g.paths(s, try x.up(std.mem.sliceAsBytes(par)), paths, depths, n); + const sc_base = try x.zero(s, tree.TreeAttention.scratchBytes(n, hk, ng, d)); + const out = try x.zero(s, n * h * d * 2); + const rws: tree.Rows = .{ .row_stream = try x.up(std.mem.sliceAsBytes(rs)), .streams = streams_dev, .offs = offs, .origin = origin, .paths = paths, .depths = depths }; + try g.run(s, try x.up(std.mem.sliceAsBytes(q)), try x.up(std.mem.sliceAsBytes(kn)), try x.up(std.mem.sliceAsBytes(vn)), rws, tree.TreeAttention.scratch(sc_base, n, hk, ng, d), out, n, h, hk, d, ng, scale); + try s.synchronize(); + return x.down(u16, out, n * h * d); +} + +fn case(gpu: check.Gpu, d: usize, h: usize, hk: usize, seed: u64, worst_all: *f64) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try tree.TreeAttention.load(gpu.d); + defer g.unload(); + var x: Set = .{ .gpu = gpu, .a = a }; + defer { + for (x.bufs.items) |*b| b.free(); + x.bufs.deinit(gpu.gpa); + } + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const committed = [_]usize{ 5, 1500, 4200 }; + const widths = [_]usize{ 9, 33, 17 }; + // the caches: stream s's rows at element offset cache_at[s] of one buffer (keys, then values) + var cache_rows: usize = 0; + for (committed) |p| cache_rows += p; + const kcache = try a.alloc(u16, 2 * cache_rows * hk * d); + for (kcache) |*v| v.* = bf16(rng.floatNorm(f32)); + var offs: [2 * committed.len]i64 = undefined; + var at: usize = 0; + for (committed, 0..) |p, si| { + offs[2 * si] = @intCast(at * hk * d); + offs[2 * si + 1] = @intCast((cache_rows + at) * hk * d); + at += p; + } + var total: usize = 0; + for (widths) |w| total += w; + const row_stream = try a.alloc(i32, total); + const parents = try a.alloc(i32, total); + var streams: [4 * committed.len]i32 = undefined; + var start: usize = 0; + for (widths, committed, 0..) |w, p, si| { + for (0..w) |j| { + row_stream[start + j] = @intCast(si); + parents[start + j] = if (j == 0 or rng.uintLessThan(u32, 6) == 0) -1 else @intCast(start + rng.uintLessThan(usize, j)); + } + for ([_]usize{ start, w, p, 0 }, 0..) |v, k| streams[4 * si + k] = @intCast(v); + start += w; + } + const qs = try a.alloc(u16, total * h * d); + const kns = try a.alloc(u16, total * hk * d); + const vns = try a.alloc(u16, total * hk * d); + for (qs) |*v| v.* = bf16(rng.floatNorm(f32)); + for (kns) |*v| v.* = bf16(rng.floatNorm(f32)); + for (vns) |*v| v.* = bf16(rng.floatNorm(f32)); + const origin = try x.up(std.mem.sliceAsBytes(kcache)); + const offs_dev = try x.up(std.mem.sliceAsBytes(&offs)); + var ng: usize = 0; + for (committed, widths) |p, w| ng = @max(ng, tree.TreeAttention.groups(p + w)); + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(d)))); + const all = try a.alloc(usize, total); + for (all, 0..) |*r, i| r.* = i; + const got = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + // the reference: committed keys, then the path root first + var worst: f64 = 0; + const kd = try a.alloc(f64, (4200 + 128) * d); + const vd = try a.alloc(f64, kd.len); + const keys = try a.alloc(usize, 4200 + 128); + const q = try a.alloc(f64, d); + const want = try a.alloc(f64, d); + const tol = try a.alloc(f64, d); + var path: [tree.max_depth]usize = undefined; + for (0..total) |r| { + const si: usize = @intCast(row_stream[r]); + const p = committed[si]; + var depth: usize = 0; + var cur: i32 = @intCast(r); + while (cur >= 0) : (cur = parents[@intCast(cur)]) { + path[depth] = @intCast(cur); + depth += 1; + } + std.mem.reverse(usize, path[0..depth]); + for (0..hk) |kh| { + for (0..p) |j| for (0..d) |e| { + const ko: usize = @intCast(offs[2 * si]); + const vo: usize = @intCast(offs[2 * si + 1]); + kd[j * d + e] = bf(kcache[ko + (j * hk + kh) * d + e]); + vd[j * d + e] = bf(kcache[vo + (j * hk + kh) * d + e]); + }; + for (0..depth) |j| for (0..d) |e| { + kd[(p + j) * d + e] = bf(kns[(path[j] * hk + kh) * d + e]); + vd[(p + j) * d + e] = bf(vns[(path[j] * hk + kh) * d + e]); + }; + for (0..p + depth) |j| keys[j] = j; + for (0..h / hk) |hh| { + const hd = kh * (h / hk) + hh; + for (q, 0..) |*v, e| v.* = bf(qs[(r * h + hd) * d + e]); + core.attention_ref.head(q, kd, vd, keys[0 .. p + depth], d, scale, want, tol); + for (want, tol, 0..) |w, t, e| worst = @max(worst, @abs(bf(got[(r * h + hd) * d + e]) - w) / t); + } + } + } + try check.expect(worst <= 1, "tree attention dim {d}: within float64 tolerance (worst {d:.3})", .{ d, worst }); + worst_all.* = @max(worst_all.*, worst); + // stream 1's rows alone (re-indexed): the same bits + const only = try a.alloc(usize, widths[1]); + for (only, 0..) |*r, i| r.* = widths[0] + i; + const alone = try call(&x, s, &g, d, only, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one stream's rows alone", std.mem.sliceAsBytes(alone), std.mem.sliceAsBytes(got[widths[0] * h * d ..][0 .. widths[1] * h * d])); + // one shared block walking every (row, group), and a grid sized for a long context: the same bits + g.shared_blocks = 1; + const strided = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one shared block", std.mem.sliceAsBytes(strided), std.mem.sliceAsBytes(got)); + g.shared_blocks = 96; + const wide = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, 129, scale); + try check.sameBytes("tree attention: groups for a 262k context", std.mem.sliceAsBytes(wide), std.mem.sliceAsBytes(got)); +} + +pub fn run(gpu: check.Gpu) !void { + var worst: f64 = 0; + try case(gpu, 128, 48, 4, 41, &worst); // Kolibri + try case(gpu, 256, 24, 4, 42, &worst); // Qwen3.8-27B + check.pass("tree attention: dim 128 and 256, committed keys and paths within float64 tolerance (worst {d:.3}), a stream's bits alone, any group slots", .{worst}); +} From de31ba9bffe55e1587b77de5226521ffae4c69f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:56:17 +0200 Subject: [PATCH 02/41] cuda/tree_attention: a group's 4 slices write their own partials, the merge combines them in order; a slice takes its own one-warp block where four-warp blocks would leave SMs idle (TreeAttention.sms), the same bits --- zig/kernels/cuda/tree_attention.cu | 118 ++++++++++++++--------------- zig/src/cuda/tree_attention.zig | 31 ++++++-- zig/tests/cuda/tree_tests.zig | 11 ++- 3 files changed, 91 insertions(+), 69 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 035b1d642f..3056dbfd5a 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -58,17 +58,18 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ } // Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -template +template __device__ __forceinline__ void tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int gi, int ng, - int64_t rq, int pf, int npairs, int64_t row) { + int64_t rq, int pf, int npairs, int64_t row, int slice) { constexpr int C = D / 8, K16 = D / 16, NB = D / 8, NT8 = HB / 8, PK = HB / 16; constexpr bool QREG = D <= 128; extern __shared__ uint4 smem[]; const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g8 = lane / 4, qd = lane % 4, mt = lane / 8; + const int wi = W == 4 ? warp : slice; // the group's slice: key tiles wi, wi + 4, ... const int64_t r = row; // the row whose stream and depth set the keys const int kh = blockIdx.y, G = h / hk; const int s = row_stream[r]; @@ -80,7 +81,7 @@ __device__ __forceinline__ void tree_group( uint4* ks = smem + warp * 2 * HB * C; uint4* vs = ks + HB * C; uint32_t qa[QREG ? K16 : 1][4]; - uint4* qs = smem + 4 * 2 * HB * C; // !QREG: the row's heads [16][D], shared by the warps + uint4* qs = smem + W * 2 * HB * C; // !QREG: the row's heads [16][D], shared by the warps if constexpr (QREG) { const int a0 = pf + g8, a1 = pf + g8 + 8; const bool v0 = a0 < npairs, v1 = a1 < npairs; @@ -108,7 +109,7 @@ __device__ __forceinline__ void tree_group( #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 = k0 / HB + warp; t * HB < k1; t += 4) { + for (int t = k0 / HB + wi; t * HB < k1; t += 4) { for (int i = lane; i < HB * C; i += 32) { const int key = t * HB + i / C; const bool in = key < k1; @@ -216,79 +217,62 @@ __device__ __forceinline__ void tree_group( } __syncwarp(); } - __syncthreads(); - float* os = reinterpret_cast(smem); // [4 warps][16 heads][D], then max and sum [4][16] - float* ms = os + 4 * 16 * D; - float* ls = ms + 4 * 16; + // each slice's partial (o, max, sum) of its 16 pairs: the merge combines a group's 4 slices in order #pragma unroll - for (int i = 0; i < NB; ++i) { - const int d = 8 * i + 2 * qd; - os[(warp * 16 + g8) * D + d] = o[i][0]; - os[(warp * 16 + g8) * D + d + 1] = o[i][1]; - os[(warp * 16 + g8 + 8) * D + d] = o[i][2]; - os[(warp * 16 + g8 + 8) * D + d + 1] = o[i][3]; - } - if (qd == 0) { - ms[warp * 16 + g8] = m0, ls[warp * 16 + g8] = l0; - ms[warp * 16 + g8 + 8] = m1, ls[warp * 16 + g8 + 8] = l1; - } - __syncthreads(); - for (int idx = threadIdx.x; idx < 16 * D; idx += 128) { - const int hd = idx / D, e = idx % D; + for (int half = 0; half < 2; ++half) { + const int hd = g8 + 8 * half; if (pf + hd >= npairs) continue; - const int64_t part = ((rq + (pf + hd) / G) * hk + kh) * ng + gi; + const int64_t part = (((rq + (pf + hd) / G) * hk + kh) * ng + gi) * 4 + wi; const int head = (pf + hd) % G; - float mm = -INFINITY; - for (int w = 0; w < 4; ++w) mm = fmaxf(mm, ms[w * 16 + hd]); - float acc = 0.0f, l = 0.0f; - for (int w = 0; w < 4; ++w) { - const float mw = ms[w * 16 + hd]; - if (mw == -INFINITY) continue; - const float a = texp(__fsub_rn(mw, mm)); - acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); - l = __fmaf_rn(a, ls[w * 16 + hd], l); +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + po[(part * 16 + head) * D + d] = o[i][2 * half]; + po[(part * 16 + head) * D + d + 1] = o[i][2 * half + 1]; } - po[(part * 16 + head) * D + e] = acc; - if (e == 0) pm[part * 16 + head] = mm, pl[part * 16 + head] = l; + if (qd == 0) pm[part * 16 + head] = half ? m1 : m0, pl[part * 16 + head] = half ? l1 : l0; } } // A stream's full committed groups once for its rows, 16 (row, head) pairs a block: each partial a row block's bits. -template +template __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { const int G = h / hk; - for (int item = blockIdx.x; item < n * ng; item += gridDim.x) { // (row, group): a fixed grid, no empty blocks - const int64_t r = item / ng; - const int gi = item % ng; + const int per = W == 4 ? 1 : 4; // slices an item takes a block + for (int item = blockIdx.x; item < n * ng * per; item += gridDim.x) { // (row, group[, slice]): a fixed grid + const int64_t r = item / per / ng; + const int gi = item / per % ng; const int st = row_stream[r]; const int64_t r0 = streams[4 * st]; const int pairs = streams[4 * st + 1] * G, p = streams[4 * st + 2]; const int tile = static_cast(r - r0); // a stream's j-th row hosts its j-th pair tile (tiles <= rows) if (tile * 16 >= pairs || (gi + 1) * SPAN > p) continue; - tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, r0, - tile * 16, pairs, r); + tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, + r0, tile * 16, pairs, r, item % per); __syncthreads(); } } -#define TF_TREE_SHARED(NAME, D, HB) \ - extern "C" __global__ void __launch_bounds__(128) NAME( \ +#define TF_TREE_SHARED(NAME, D, HB, W) \ + extern "C" __global__ void __launch_bounds__(32 * W) NAME( \ const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { \ - tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ + tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ } -TF_TREE_SHARED(tf_tree_shared, 128, 32) -TF_TREE_SHARED(tf_tree_shared_256, 256, 16) +TF_TREE_SHARED(tf_tree_shared, 128, 32, 4) +TF_TREE_SHARED(tf_tree_shared_256, 256, 16, 4) +TF_TREE_SHARED(tf_tree_shared_w1, 128, 32, 1) +TF_TREE_SHARED(tf_tree_shared_256_w1, 256, 16, 1) // Block (row, key head, group slot): the row's groups past its stream's full committed ones, slot, slot + slots, ... -template +template __device__ __forceinline__ void tree_part( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, @@ -296,46 +280,58 @@ __device__ __forceinline__ void tree_part( float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng) { const int64_t r = blockIdx.x; const int p = streams[4 * row_stream[r] + 2], total = p + depths[r]; - for (int gi = blockIdx.z + p / SPAN; gi < ng && gi * SPAN < total; gi += gridDim.z) { // full committed: shared - tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, - r, 0, h / hk, r); + const int per = W == 4 ? 1 : 4, slot = blockIdx.z / per, slots = gridDim.z / per; + for (int gi = slot + p / SPAN; gi < ng && gi * SPAN < total; gi += slots) { // full committed: shared + tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + ng, r, 0, h / hk, r, blockIdx.z % per); __syncthreads(); // the next group's tiles reuse the partials' shared memory } } -#define TF_TREE_PART(NAME, D, HB) \ - extern "C" __global__ void __launch_bounds__(128) NAME( \ +#define TF_TREE_PART(NAME, D, HB, W) \ + extern "C" __global__ void __launch_bounds__(32 * W) NAME( \ const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, \ const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng) { \ - tree_part(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng); \ + tree_part(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, \ + ng); \ } -TF_TREE_PART(tf_tree_part, 128, 32) -TF_TREE_PART(tf_tree_part_256, 256, 16) +TF_TREE_PART(tf_tree_part, 128, 32, 4) +TF_TREE_PART(tf_tree_part_256, 256, 16, 4) +TF_TREE_PART(tf_tree_part_w1, 128, 32, 1) +TF_TREE_PART(tf_tree_part_256_w1, 256, 16, 1) -// out [W, H, D] bf16: block (row, key head) folds the row's groups in order. +// out [W, H, D] bf16: block (row, key head, head) folds the row's groups in order, each group's 4 slices first. extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ depths, __nv_bfloat16* __restrict__ out, int h, int hk, int ng, int d) { const int64_t r = blockIdx.x; - const int kh = blockIdx.y, G = h / hk; + const int kh = blockIdx.y, G = h / hk, hd = blockIdx.z; // block (row, key head, head): each element alone const int total = streams[4 * row_stream[r] + 2] + depths[r]; const int groups = (total + SPAN - 1) / SPAN; - for (int idx = threadIdx.x; idx < G * d; idx += 256) { - const int hd = idx / d, e = idx % d; + for (int e = threadIdx.x; e < d; e += 256) { float m = -INFINITY, l = 0.0f, acc = 0.0f; for (int gi = 0; gi < groups; ++gi) { - const int64_t part = (r * hk + kh) * ng + gi; - const float cm = pm[part * 16 + hd], cl = pl[part * 16 + hd]; + const int64_t part = ((r * hk + kh) * ng + gi) * 4; + float cm = -INFINITY; + for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * 16 + hd]); + float co = 0.0f, cl = 0.0f; + for (int w = 0; w < 4; ++w) { + const float mw = pm[(part + w) * 16 + hd]; + if (mw == -INFINITY) continue; + const float aw = texp(__fsub_rn(mw, cm)); + co = __fmaf_rn(aw, po[((part + w) * 16 + hd) * d + e], co); + cl = __fmaf_rn(aw, pl[(part + w) * 16 + hd], cl); + } const bool active = cl > 0.0f; const float next = active ? fmaxf(m, cm) : m; const float a = active ? (m == -INFINITY ? 0.0f : texp(__fsub_rn(m, next))) : 1.0f; const float b = active ? texp(__fsub_rn(cm, next)) : 0.0f; - acc = __fmaf_rn(acc, a, __fmul_rn(po[(part * 16 + hd) * d + e], b)); + acc = __fmaf_rn(acc, a, __fmul_rn(co, b)); l = __fmaf_rn(l, a, __fmul_rn(cl, b)); m = next; } diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index dd8e56bdcc..5f1f9ad418 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -14,6 +14,10 @@ const Stream = stream_.Stream; pub const span = 2048; pub const max_depth = 128; const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; +/// The one-warp blocks' shared memory: one warp's K/V slots (at 256 also the queries). +const part_smem_w1 = [2]u32{ 2 * 32 * 16 * 16, 2 * 16 * 32 * 16 + 16 * 32 * 16 }; +/// A group's slices: its key tiles dealt to 4 warps, one partial each, combined in order by the merge. +const slices = 4; /// row_stream [n], streams [s, 4] (committed at [2]), offs [s, 2] (K, V from origin), paths [n, max_depth], depths. pub const Rows = struct { row_stream: u64, streams: u64, offs: u64, origin: u64, paths: u64, depths: u64 }; @@ -25,9 +29,13 @@ pub const TreeAttention = struct { paths_fn: Function, part_fns: [2]Function, // head dim 128, 256 shared_fns: [2]Function, + part_w1: [2]Function, // the same slices, one warp a block + shared_w1: [2]Function, merge_fn: Function, /// The shared kernel's blocks a key head: a fixed grid walking every (row, group) (bits do not depend on it). shared_blocks: usize = 96, + /// The device's SMs: where the four-warp grid would leave some idle, a slice takes its own block (the same bits). + sms: usize = 0, pub fn load(d: *const Driver) !TreeAttention { if (!kernels.available) return error.BuiltWithoutKernels; @@ -37,7 +45,11 @@ pub const TreeAttention = struct { for (&parts, part_smem) |*f, b| try f.allowDynamicShared(b); var shared: [2]Function = .{ try mod.function("tf_tree_shared"), try mod.function("tf_tree_shared_256") }; for (&shared, part_smem) |*f, b| try f.allowDynamicShared(b); - return .{ .mod = mod, .paths_fn = try mod.function("tf_tree_paths"), .part_fns = parts, .shared_fns = shared, .merge_fn = try mod.function("tf_tree_merge") }; + var parts1: [2]Function = .{ try mod.function("tf_tree_part_w1"), try mod.function("tf_tree_part_256_w1") }; + for (&parts1, part_smem_w1) |*f, b| try f.allowDynamicShared(b); + var shared1: [2]Function = .{ try mod.function("tf_tree_shared_w1"), try mod.function("tf_tree_shared_256_w1") }; + for (&shared1, part_smem_w1) |*f, b| try f.allowDynamicShared(b); + return .{ .mod = mod, .paths_fn = try mod.function("tf_tree_paths"), .part_fns = parts, .shared_fns = shared, .part_w1 = parts1, .shared_w1 = shared1, .merge_fn = try mod.function("tf_tree_merge") }; } pub fn unload(g: *TreeAttention) void { @@ -49,13 +61,13 @@ pub const TreeAttention = struct { return (keys + span - 1) / span; } - /// The fp32 partials of `rows` rows, `hk` key heads and `ng` groups: o [.., 16, dim], then max and sum [.., 16]. + /// fp32 partials of `rows` rows, `hk` key heads, `ng` groups, 4 slices each: o [.., 16, dim], max and sum [.., 16]. pub fn scratchBytes(rows: usize, hk: usize, ng: usize, dim: usize) usize { - return rows * hk * ng * 16 * (dim + 2) * 4; + return rows * hk * ng * slices * 16 * (dim + 2) * 4; } pub fn scratch(base: u64, rows: usize, hk: usize, ng: usize, dim: usize) Scratch { - const parts = rows * hk * ng * 16; + const parts = rows * hk * ng * slices * 16; return .{ .o = base, .m = base + parts * dim * 4, .l = base + parts * dim * 4 + parts * 4 }; } @@ -86,15 +98,20 @@ pub const TreeAttention = struct { c.add(scale); c.add(@as(i32, @intCast(ng))); c.add(@as(i32, @intCast(n))); - try launch_.launch(g.shared_fns[which], .{ .grid = .{ .x = @intCast(@min(n * ng, g.shared_blocks)), .y = @intCast(hk), .z = 1 }, .block = .{ .x = 128 }, .shared = part_smem[which] }, s, &c); + // one warp a block where four-warp blocks would leave SMs idle: each slice its own block, the same partials + const one = @min(n * ng, g.shared_blocks) * hk < g.sms; + const per: usize = if (one) slices else 1; + const block: u32 = if (one) 32 else 128; + const smem = if (one) part_smem_w1[which] else part_smem[which]; + try launch_.launch(if (one) g.shared_w1[which] else g.shared_fns[which], .{ .grid = .{ .x = @intCast(@min(n * ng, g.shared_blocks) * per), .y = @intCast(hk), .z = 1 }, .block = .{ .x = block }, .shared = smem }, s, &c); // a row's own block: its boundary group and at most the next (a path is at most 128 keys) - try launch_.launch(g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2)) }, .block = .{ .x = 128 }, .shared = part_smem[which] }, s, &a); + try launch_.launch(if (one) g.part_w1[which] else g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2) * per) }, .block = .{ .x = block }, .shared = smem }, s, &a); var b: launch_.Args = .{}; for ([_]u64{ sc.o, sc.m, sc.l, rows.row_stream, rows.streams, rows.depths, out }) |v| b.add(v); b.add(@as(i32, @intCast(h))); b.add(@as(i32, @intCast(hk))); b.add(@as(i32, @intCast(ng))); b.add(@as(i32, @intCast(dim))); - try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = 1 }, .block = .{ .x = 256 } }, s, &b); + try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(h / hk) }, .block = .{ .x = 256 } }, s, &b); } }; diff --git a/zig/tests/cuda/tree_tests.zig b/zig/tests/cuda/tree_tests.zig index c11088413c..203c68fdce 100644 --- a/zig/tests/cuda/tree_tests.zig +++ b/zig/tests/cuda/tree_tests.zig @@ -191,11 +191,20 @@ fn case(gpu: check.Gpu, d: usize, h: usize, hk: usize, seed: u64, worst_all: *f6 g.shared_blocks = 96; const wide = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, 129, scale); try check.sameBytes("tree attention: groups for a 262k context", std.mem.sliceAsBytes(wide), std.mem.sliceAsBytes(got)); + // one warp a block (as on a device with more SMs than four-warp blocks): every slice's partial the same + g.sms = 1 << 20; + defer g.sms = 0; + const one = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one warp a block", std.mem.sliceAsBytes(one), std.mem.sliceAsBytes(got)); + g.shared_blocks = 1; + const one_strided = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one warp a block, one shared block", std.mem.sliceAsBytes(one_strided), std.mem.sliceAsBytes(got)); + g.shared_blocks = 96; } pub fn run(gpu: check.Gpu) !void { var worst: f64 = 0; try case(gpu, 128, 48, 4, 41, &worst); // Kolibri try case(gpu, 256, 24, 4, 42, &worst); // Qwen3.8-27B - check.pass("tree attention: dim 128 and 256, committed keys and paths within float64 tolerance (worst {d:.3}), a stream's bits alone, any group slots", .{worst}); + check.pass("tree attention: dim 128 and 256, committed keys and paths within float64 tolerance (worst {d:.3}), a stream's bits alone, any group slots, four-warp or one-warp blocks", .{worst}); } From 6c3d49d5cd976404a1babf86e1c68965633dcd52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:56:22 +0200 Subject: [PATCH 03/41] cuda/tree_attention: four-warp blocks combine a group's slices themselves again, only one-warp blocks leave them to the merge (a small fixed slice area), so the scratch stays at one partial a group --- zig/kernels/cuda/tree_attention.cu | 84 ++++++++++++++++++++++-------- zig/src/cuda/tree_attention.zig | 28 ++++++---- zig/tests/cuda/tree_tests.zig | 13 +++-- 3 files changed, 88 insertions(+), 37 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 3056dbfd5a..137fb8fef7 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -217,20 +217,59 @@ __device__ __forceinline__ void tree_group( } __syncwarp(); } - // each slice's partial (o, max, sum) of its 16 pairs: the merge combines a group's 4 slices in order + if constexpr (W == 1) { + // one warp a block: each slice's partial (o, max, sum) of its 16 pairs; the merge combines the 4 in order #pragma unroll - for (int half = 0; half < 2; ++half) { - const int hd = g8 + 8 * half; - if (pf + hd >= npairs) continue; - const int64_t part = (((rq + (pf + hd) / G) * hk + kh) * ng + gi) * 4 + wi; - const int head = (pf + hd) % G; + for (int half = 0; half < 2; ++half) { + const int hd = g8 + 8 * half; + if (pf + hd >= npairs) continue; + const int64_t part = (((rq + (pf + hd) / G) * hk + kh) * ng + gi) * 4 + wi; + const int head = (pf + hd) % G; +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + po[(part * 16 + head) * D + d] = o[i][2 * half]; + po[(part * 16 + head) * D + d + 1] = o[i][2 * half + 1]; + } + if (qd == 0) pm[part * 16 + head] = half ? m1 : m0, pl[part * 16 + head] = half ? l1 : l0; + } + return; + } else { + __syncthreads(); + float* os = reinterpret_cast(smem); // [4 warps][16 heads][D], then max and sum [4][16] + float* ms = os + 4 * 16 * D; + float* ls = ms + 4 * 16; #pragma unroll for (int i = 0; i < NB; ++i) { const int d = 8 * i + 2 * qd; - po[(part * 16 + head) * D + d] = o[i][2 * half]; - po[(part * 16 + head) * D + d + 1] = o[i][2 * half + 1]; + os[(warp * 16 + g8) * D + d] = o[i][0]; + os[(warp * 16 + g8) * D + d + 1] = o[i][1]; + os[(warp * 16 + g8 + 8) * D + d] = o[i][2]; + os[(warp * 16 + g8 + 8) * D + d + 1] = o[i][3]; + } + if (qd == 0) { + ms[warp * 16 + g8] = m0, ls[warp * 16 + g8] = l0; + ms[warp * 16 + g8 + 8] = m1, ls[warp * 16 + g8 + 8] = l1; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * D; idx += 128) { + const int hd = idx / D, e = idx % D; + if (pf + hd >= npairs) continue; + const int64_t part = ((rq + (pf + hd) / G) * hk + kh) * ng + gi; + const int head = (pf + hd) % G; + float mm = -INFINITY; + for (int w = 0; w < 4; ++w) mm = fmaxf(mm, ms[w * 16 + hd]); + float acc = 0.0f, l = 0.0f; + for (int w = 0; w < 4; ++w) { + const float mw = ms[w * 16 + hd]; + if (mw == -INFINITY) continue; + const float a = texp(__fsub_rn(mw, mm)); + acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); + l = __fmaf_rn(a, ls[w * 16 + hd], l); + } + po[(part * 16 + head) * D + e] = acc; + if (e == 0) pm[part * 16 + head] = mm, pl[part * 16 + head] = l; } - if (qd == 0) pm[part * 16 + head] = half ? m1 : m0, pl[part * 16 + head] = half ? l1 : l0; } } @@ -304,11 +343,11 @@ TF_TREE_PART(tf_tree_part_256, 256, 16, 4) TF_TREE_PART(tf_tree_part_w1, 128, 32, 1) TF_TREE_PART(tf_tree_part_256_w1, 256, 16, 1) -// out [W, H, D] bf16: block (row, key head, head) folds the row's groups in order, each group's 4 slices first. +// out [W, H, D] bf16: block (row, key head, head) folds the row's groups in order (sliced: a group's 4 slices first). extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ depths, - __nv_bfloat16* __restrict__ out, int h, int hk, int ng, int d) { + __nv_bfloat16* __restrict__ out, int h, int hk, int ng, int d, int sliced) { const int64_t r = blockIdx.x; const int kh = blockIdx.y, G = h / hk, hd = blockIdx.z; // block (row, key head, head): each element alone const int total = streams[4 * row_stream[r] + 2] + depths[r]; @@ -316,16 +355,19 @@ extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( for (int e = threadIdx.x; e < d; e += 256) { float m = -INFINITY, l = 0.0f, acc = 0.0f; for (int gi = 0; gi < groups; ++gi) { - const int64_t part = ((r * hk + kh) * ng + gi) * 4; - float cm = -INFINITY; - for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * 16 + hd]); - float co = 0.0f, cl = 0.0f; - for (int w = 0; w < 4; ++w) { - const float mw = pm[(part + w) * 16 + hd]; - if (mw == -INFINITY) continue; - const float aw = texp(__fsub_rn(mw, cm)); - co = __fmaf_rn(aw, po[((part + w) * 16 + hd) * d + e], co); - cl = __fmaf_rn(aw, pl[(part + w) * 16 + hd], cl); + const int64_t part = ((r * hk + kh) * ng + gi) * (sliced ? 4 : 1); + float cm = -INFINITY, co = 0.0f, cl = 0.0f; + if (!sliced) { // four-warp blocks combined the group's slices themselves + cm = pm[part * 16 + hd], co = po[(part * 16 + hd) * d + e], cl = pl[part * 16 + hd]; + } else { // one-warp blocks left the 4 slices: their combine, in the same order + for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * 16 + hd]); + for (int w = 0; w < 4; ++w) { + const float mw = pm[(part + w) * 16 + hd]; + if (mw == -INFINITY) continue; + const float aw = texp(__fsub_rn(mw, cm)); + co = __fmaf_rn(aw, po[((part + w) * 16 + hd) * d + e], co); + cl = __fmaf_rn(aw, pl[(part + w) * 16 + hd], cl); + } } const bool active = cl > 0.0f; const float next = active ? fmaxf(m, cm) : m; diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index 5f1f9ad418..1647c07a77 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -16,13 +16,15 @@ pub const max_depth = 128; const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; /// The one-warp blocks' shared memory: one warp's K/V slots (at 256 also the queries). const part_smem_w1 = [2]u32{ 2 * 32 * 16 * 16, 2 * 16 * 32 * 16 + 16 * 32 * 16 }; -/// A group's slices: its key tiles dealt to 4 warps, one partial each, combined in order by the merge. +/// A group's slices: its key tiles dealt to 4 warps, combined in order (in the block, or by the merge). const slices = 4; +/// (row, key head, group) items one-warp blocks may take: their slices' partials sit apart from the combined ones. +pub const slice_items = 256; /// row_stream [n], streams [s, 4] (committed at [2]), offs [s, 2] (K, V from origin), paths [n, max_depth], depths. pub const Rows = struct { row_stream: u64, streams: u64, offs: u64, origin: u64, paths: u64, depths: u64 }; -pub const Scratch = struct { o: u64, m: u64, l: u64 }; +pub const Scratch = struct { o: u64, m: u64, l: u64, so: u64, sm: u64, sl: u64 }; pub const TreeAttention = struct { mod: Module, @@ -61,14 +63,16 @@ pub const TreeAttention = struct { return (keys + span - 1) / span; } - /// fp32 partials of `rows` rows, `hk` key heads, `ng` groups, 4 slices each: o [.., 16, dim], max and sum [.., 16]. + /// fp32 partials of `rows` rows, `hk` key heads, `ng` groups (o [.., 16, dim], max, sum [.., 16]), then slices. pub fn scratchBytes(rows: usize, hk: usize, ng: usize, dim: usize) usize { - return rows * hk * ng * slices * 16 * (dim + 2) * 4; + return (rows * hk * ng + slice_items * slices) * 16 * (dim + 2) * 4; } pub fn scratch(base: u64, rows: usize, hk: usize, ng: usize, dim: usize) Scratch { - const parts = rows * hk * ng * slices * 16; - return .{ .o = base, .m = base + parts * dim * 4, .l = base + parts * dim * 4 + parts * 4 }; + const parts = rows * hk * ng * 16; + const sb = base + parts * (dim + 2) * 4; + const sp = slice_items * slices * 16; + return .{ .o = base, .m = base + parts * dim * 4, .l = base + parts * dim * 4 + parts * 4, .so = sb, .sm = sb + sp * dim * 4, .sl = sb + sp * dim * 4 + sp * 4 }; } /// Each row's path (root first) and depth from `parents` int32 [n] (global window rows, -1 at a root). @@ -85,21 +89,22 @@ pub const TreeAttention = struct { if (hk == 0 or h % hk != 0 or h / hk > 16 or ng == 0 or (dim != 128 and dim != 256)) return error.TreeShape; const which: usize = if (dim == 256) 1 else 0; if (n == 0) return; + // one warp a block where four-warp blocks would leave SMs idle: each slice its own block, its own partial + const one = @min(n * ng, g.shared_blocks) * hk < g.sms and n * ng * hk <= slice_items; + const o = if (one) [3]u64{ sc.so, sc.sm, sc.sl } else [3]u64{ sc.o, sc.m, sc.l }; var a: launch_.Args = .{}; - for ([_]u64{ q, kn, vn, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, sc.o, sc.m, sc.l }) |v| a.add(v); + for ([_]u64{ q, kn, vn, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, o[0], o[1], o[2] }) |v| a.add(v); a.add(@as(i32, @intCast(h))); a.add(@as(i32, @intCast(hk))); a.add(scale); a.add(@as(i32, @intCast(ng))); var c: launch_.Args = .{}; - for ([_]u64{ q, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, sc.o, sc.m, sc.l }) |v| c.add(v); + for ([_]u64{ q, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, o[0], o[1], o[2] }) |v| c.add(v); c.add(@as(i32, @intCast(h))); c.add(@as(i32, @intCast(hk))); c.add(scale); c.add(@as(i32, @intCast(ng))); c.add(@as(i32, @intCast(n))); - // one warp a block where four-warp blocks would leave SMs idle: each slice its own block, the same partials - const one = @min(n * ng, g.shared_blocks) * hk < g.sms; const per: usize = if (one) slices else 1; const block: u32 = if (one) 32 else 128; const smem = if (one) part_smem_w1[which] else part_smem[which]; @@ -107,11 +112,12 @@ pub const TreeAttention = struct { // a row's own block: its boundary group and at most the next (a path is at most 128 keys) try launch_.launch(if (one) g.part_w1[which] else g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2) * per) }, .block = .{ .x = block }, .shared = smem }, s, &a); var b: launch_.Args = .{}; - for ([_]u64{ sc.o, sc.m, sc.l, rows.row_stream, rows.streams, rows.depths, out }) |v| b.add(v); + for ([_]u64{ o[0], o[1], o[2], rows.row_stream, rows.streams, rows.depths, out }) |v| b.add(v); b.add(@as(i32, @intCast(h))); b.add(@as(i32, @intCast(hk))); b.add(@as(i32, @intCast(ng))); b.add(@as(i32, @intCast(dim))); + b.add(@as(i32, @intFromBool(one))); try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(h / hk) }, .block = .{ .x = 256 } }, s, &b); } }; diff --git a/zig/tests/cuda/tree_tests.zig b/zig/tests/cuda/tree_tests.zig index 203c68fdce..3b8de8d5cb 100644 --- a/zig/tests/cuda/tree_tests.zig +++ b/zig/tests/cuda/tree_tests.zig @@ -191,14 +191,17 @@ fn case(gpu: check.Gpu, d: usize, h: usize, hk: usize, seed: u64, worst_all: *f6 g.shared_blocks = 96; const wide = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, 129, scale); try check.sameBytes("tree attention: groups for a 262k context", std.mem.sliceAsBytes(wide), std.mem.sliceAsBytes(got)); - // one warp a block (as on a device with more SMs than four-warp blocks): every slice's partial the same + // one warp a block (a device with SMs to spare): stream 2's rows alone, its full groups too, the batch's bits g.sms = 1 << 20; defer g.sms = 0; - const one = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); - try check.sameBytes("tree attention: one warp a block", std.mem.sliceAsBytes(one), std.mem.sliceAsBytes(got)); + const third = try a.alloc(usize, widths[2]); + for (third, 0..) |*r, i| r.* = widths[0] + widths[1] + i; + try check.expect(widths[2] * ng * hk <= tree.slice_items, "tree attention: the one-warp case fits the slices", .{}); + const one = try call(&x, s, &g, d, third, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one warp a block", std.mem.sliceAsBytes(one), std.mem.sliceAsBytes(got[(widths[0] + widths[1]) * h * d ..][0 .. widths[2] * h * d])); g.shared_blocks = 1; - const one_strided = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); - try check.sameBytes("tree attention: one warp a block, one shared block", std.mem.sliceAsBytes(one_strided), std.mem.sliceAsBytes(got)); + const one_strided = try call(&x, s, &g, d, third, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one warp a block, one shared block", std.mem.sliceAsBytes(one_strided), std.mem.sliceAsBytes(one)); g.shared_blocks = 96; } From 197cacc79baa099a8c347fa172de83a12665b202 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:16:03 +0200 Subject: [PATCH 04/41] cuda/tree_attention: one-warp blocks fetch their tiles 2 ahead (3 cp.async stages), the same arithmetic in the same order --- zig/kernels/cuda/tree_attention.cu | 36 +++++++++++++++++++++++++----- zig/src/cuda/tree_attention.zig | 5 +++-- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 137fb8fef7..cdca510755 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -16,6 +16,9 @@ using qmm_frag::mma; using qmm_frag::mma0; constexpr int SPAN = 2048; // keys a group: one partial +#ifndef TF_TREE_STAGES +#define TF_TREE_STAGES 3 +#endif constexpr int MAXD = 128; // a path's most rows constexpr float LOG2E = 1.4426950408889634f; @@ -78,10 +81,12 @@ __device__ __forceinline__ void tree_group( if (k0 >= total) return; // the merge reads only the row's own groups const __nv_bfloat16* kc = origin + offs[2 * s]; const __nv_bfloat16* vc = origin + offs[2 * s + 1]; - uint4* ks = smem + warp * 2 * HB * C; + constexpr int ST = W == 1 ? TF_TREE_STAGES : 1; // tile stages: a one-warp block fetches ahead while it computes + uint4* const kv0 = smem + warp * ST * 2 * HB * C; + uint4* ks = kv0; uint4* vs = ks + HB * C; uint32_t qa[QREG ? K16 : 1][4]; - uint4* qs = smem + W * 2 * HB * C; // !QREG: the row's heads [16][D], shared by the warps + uint4* qs = smem + W * ST * 2 * HB * C; // !QREG: the row's heads [16][D], shared by the warps if constexpr (QREG) { const int a0 = pf + g8, a1 = pf + g8 + 8; const bool v0 = a0 < npairs, v1 = a1 < npairs; @@ -109,7 +114,9 @@ __device__ __forceinline__ void tree_group( #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 = k0 / HB + wi; t * HB < k1; t += 4) { + // tile t's keys and values into stage st (committed keys, then the row's path), one cp.async group + auto load = [&](int t, int st) { + uint4* const kd = kv0 + st * 2 * HB * C; for (int i = lane; i < HB * C; i += 32) { const int key = t * HB + i / C; const bool in = key < k1; @@ -124,11 +131,28 @@ __device__ __forceinline__ void tree_group( vsrc = vn + ((int64_t)node * hk + kh) * D; } if (!in) ksrc = kn, vsrc = vn; - cp16z(ks + sw(i / C, i % C), ksrc + (i % C) * 8, in); - cp16z(vs + sw(i / C, i % C), vsrc + (i % C) * 8, in); + cp16z(kd + sw(i / C, i % C), ksrc + (i % C) * 8, in); + cp16z(kd + HB * C + sw(i / C, i % C), vsrc + (i % C) * 8, in); } commit(); - wait_group<0>(); + }; + // ST - 1 tiles ahead: one cp.async group a tile, an empty one past the end, so the wait count never changes + const int t0 = k0 / HB + wi; +#pragma unroll + for (int a = 0; a < ST - 1; ++a) { + if ((t0 + 4 * a) * HB < k1) load(t0 + 4 * a, a); + else commit(); + } + int st = 0; + for (int t = t0; t * HB < k1; t += 4) { + const int ahead = t + 4 * (ST - 1); + if (ahead * HB < k1) load(ahead, (st + ST - 1) % ST); + else commit(); + wait_group(); // tile t landed; the next ST - 1 are in flight + ks = kv0 + st * 2 * HB * C; + vs = ks + HB * C; + st = (st + 1) % ST; + __syncwarp(); float sc[NT8][4]; #pragma unroll diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index 1647c07a77..c58e4ba9f7 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -14,8 +14,9 @@ const Stream = stream_.Stream; pub const span = 2048; pub const max_depth = 128; const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; -/// The one-warp blocks' shared memory: one warp's K/V slots (at 256 also the queries). -const part_smem_w1 = [2]u32{ 2 * 32 * 16 * 16, 2 * 16 * 32 * 16 + 16 * 32 * 16 }; +/// The one-warp blocks' tile stages (tree_attention.cu's TF_TREE_STAGES) and shared memory (at 256 also the queries). +const stages_w1 = 3; +const part_smem_w1 = [2]u32{ stages_w1 * 2 * 32 * 16 * 16, stages_w1 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; /// A group's slices: its key tiles dealt to 4 warps, combined in order (in the block, or by the merge). const slices = 4; /// (row, key head, group) items one-warp blocks may take: their slices' partials sit apart from the combined ones. From 403549d1200ef7f59a2456593adace5eb8568244 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:40:39 +0200 Subject: [PATCH 05/41] cuda/tree_attention: 512-key groups (4x the warps in flight at one row), partials sized to the query heads, one-warp blocks with 3 tile stages where they fit one wave of the SMs, else 2 --- zig/kernels/cuda/tree_attention.cu | 65 +++++++++++++++--------------- zig/src/cuda/tree_attention.zig | 51 ++++++++++++----------- zig/tests/cuda/tree_tests.zig | 10 +++-- 3 files changed, 68 insertions(+), 58 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index cdca510755..310ab7d659 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -15,10 +15,7 @@ using qmm_frag::ldmatrix4; using qmm_frag::mma; using qmm_frag::mma0; -constexpr int SPAN = 2048; // keys a group: one partial -#ifndef TF_TREE_STAGES -#define TF_TREE_STAGES 3 -#endif +constexpr int SPAN = 512; // keys a group: one partial (4 slices of its tiles) constexpr int MAXD = 128; // a path's most rows constexpr float LOG2E = 1.4426950408889634f; @@ -61,7 +58,7 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ } // Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -template +template __device__ __forceinline__ void tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, @@ -81,7 +78,7 @@ __device__ __forceinline__ void tree_group( if (k0 >= total) return; // the merge reads only the row's own groups const __nv_bfloat16* kc = origin + offs[2 * s]; const __nv_bfloat16* vc = origin + offs[2 * s + 1]; - constexpr int ST = W == 1 ? TF_TREE_STAGES : 1; // tile stages: a one-warp block fetches ahead while it computes + // ST tile stages: a one-warp block fetches ST - 1 tiles ahead while it computes one (stages never change bits) uint4* const kv0 = smem + warp * ST * 2 * HB * C; uint4* ks = kv0; uint4* vs = ks + HB * C; @@ -252,10 +249,10 @@ __device__ __forceinline__ void tree_group( #pragma unroll for (int i = 0; i < NB; ++i) { const int d = 8 * i + 2 * qd; - po[(part * 16 + head) * D + d] = o[i][2 * half]; - po[(part * 16 + head) * D + d + 1] = o[i][2 * half + 1]; + po[(part * G + head) * D + d] = o[i][2 * half]; + po[(part * G + head) * D + d + 1] = o[i][2 * half + 1]; } - if (qd == 0) pm[part * 16 + head] = half ? m1 : m0, pl[part * 16 + head] = half ? l1 : l0; + if (qd == 0) pm[part * G + head] = half ? m1 : m0, pl[part * G + head] = half ? l1 : l0; } return; } else { @@ -291,14 +288,14 @@ __device__ __forceinline__ void tree_group( acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); l = __fmaf_rn(a, ls[w * 16 + hd], l); } - po[(part * 16 + head) * D + e] = acc; - if (e == 0) pm[part * 16 + head] = mm, pl[part * 16 + head] = l; + po[(part * G + head) * D + e] = acc; + if (e == 0) pm[part * G + head] = mm, pl[part * G + head] = l; } } } // A stream's full committed groups once for its rows, 16 (row, head) pairs a block: each partial a row block's bits. -template +template __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, @@ -314,28 +311,30 @@ __device__ __forceinline__ void tree_shared( const int pairs = streams[4 * st + 1] * G, p = streams[4 * st + 2]; const int tile = static_cast(r - r0); // a stream's j-th row hosts its j-th pair tile (tiles <= rows) if (tile * 16 >= pairs || (gi + 1) * SPAN > p) continue; - tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, + tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, r0, tile * 16, pairs, r, item % per); __syncthreads(); } } -#define TF_TREE_SHARED(NAME, D, HB, W) \ +#define TF_TREE_SHARED(NAME, D, HB, W, ST) \ extern "C" __global__ void __launch_bounds__(32 * W) NAME( \ const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { \ - tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ + tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ } -TF_TREE_SHARED(tf_tree_shared, 128, 32, 4) -TF_TREE_SHARED(tf_tree_shared_256, 256, 16, 4) -TF_TREE_SHARED(tf_tree_shared_w1, 128, 32, 1) -TF_TREE_SHARED(tf_tree_shared_256_w1, 256, 16, 1) +TF_TREE_SHARED(tf_tree_shared, 128, 32, 4, 1) +TF_TREE_SHARED(tf_tree_shared_256, 256, 16, 4, 1) +TF_TREE_SHARED(tf_tree_shared_w1, 128, 32, 1, 2) +TF_TREE_SHARED(tf_tree_shared_256_w1, 256, 16, 1, 2) +TF_TREE_SHARED(tf_tree_shared_w1s3, 128, 32, 1, 3) +TF_TREE_SHARED(tf_tree_shared_256_w1s3, 256, 16, 1, 3) // Block (row, key head, group slot): the row's groups past its stream's full committed ones, slot, slot + slots, ... -template +template __device__ __forceinline__ void tree_part( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, @@ -345,27 +344,29 @@ __device__ __forceinline__ void tree_part( const int p = streams[4 * row_stream[r] + 2], total = p + depths[r]; const int per = W == 4 ? 1 : 4, slot = blockIdx.z / per, slots = gridDim.z / per; for (int gi = slot + p / SPAN; gi < ng && gi * SPAN < total; gi += slots) { // full committed: shared - tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, r, 0, h / hk, r, blockIdx.z % per); __syncthreads(); // the next group's tiles reuse the partials' shared memory } } -#define TF_TREE_PART(NAME, D, HB, W) \ +#define TF_TREE_PART(NAME, D, HB, W, ST) \ extern "C" __global__ void __launch_bounds__(32 * W) NAME( \ const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, \ const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng) { \ - tree_part(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, \ + tree_part(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, \ ng); \ } -TF_TREE_PART(tf_tree_part, 128, 32, 4) -TF_TREE_PART(tf_tree_part_256, 256, 16, 4) -TF_TREE_PART(tf_tree_part_w1, 128, 32, 1) -TF_TREE_PART(tf_tree_part_256_w1, 256, 16, 1) +TF_TREE_PART(tf_tree_part, 128, 32, 4, 1) +TF_TREE_PART(tf_tree_part_256, 256, 16, 4, 1) +TF_TREE_PART(tf_tree_part_w1, 128, 32, 1, 2) +TF_TREE_PART(tf_tree_part_256_w1, 256, 16, 1, 2) +TF_TREE_PART(tf_tree_part_w1s3, 128, 32, 1, 3) +TF_TREE_PART(tf_tree_part_256_w1s3, 256, 16, 1, 3) // out [W, H, D] bf16: block (row, key head, head) folds the row's groups in order (sliced: a group's 4 slices first). extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( @@ -382,15 +383,15 @@ extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( const int64_t part = ((r * hk + kh) * ng + gi) * (sliced ? 4 : 1); float cm = -INFINITY, co = 0.0f, cl = 0.0f; if (!sliced) { // four-warp blocks combined the group's slices themselves - cm = pm[part * 16 + hd], co = po[(part * 16 + hd) * d + e], cl = pl[part * 16 + hd]; + cm = pm[part * G + hd], co = po[(part * G + hd) * d + e], cl = pl[part * G + hd]; } else { // one-warp blocks left the 4 slices: their combine, in the same order - for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * 16 + hd]); + for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * G + hd]); for (int w = 0; w < 4; ++w) { - const float mw = pm[(part + w) * 16 + hd]; + const float mw = pm[(part + w) * G + hd]; if (mw == -INFINITY) continue; const float aw = texp(__fsub_rn(mw, cm)); - co = __fmaf_rn(aw, po[((part + w) * 16 + hd) * d + e], co); - cl = __fmaf_rn(aw, pl[(part + w) * 16 + hd], cl); + co = __fmaf_rn(aw, po[((part + w) * G + hd) * d + e], co); + cl = __fmaf_rn(aw, pl[(part + w) * G + hd], cl); } } const bool active = cl > 0.0f; diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index c58e4ba9f7..69b2ad598f 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -11,16 +11,17 @@ const Function = module.Function; const Stream = stream_.Stream; /// Keys a group, a path's most rows, the part kernels' shared memory (four warps' K/V slots; at 256 also the queries). -pub const span = 2048; +pub const span = 512; pub const max_depth = 128; const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; -/// The one-warp blocks' tile stages (tree_attention.cu's TF_TREE_STAGES) and shared memory (at 256 also the queries). -const stages_w1 = 3; -const part_smem_w1 = [2]u32{ stages_w1 * 2 * 32 * 16 * 16, stages_w1 * 2 * 16 * 32 * 16 + 16 * 32 * 16 }; +/// One-warp blocks' shared memory by tile stages (2: two blocks an SM; 3: one, deeper), at 256 also the queries. +fn smemW1(stages: u32, which: usize) u32 { + return if (which == 0) stages * 2 * 32 * 16 * 16 else stages * 2 * 16 * 32 * 16 + 16 * 32 * 16; +} /// A group's slices: its key tiles dealt to 4 warps, combined in order (in the block, or by the merge). const slices = 4; /// (row, key head, group) items one-warp blocks may take: their slices' partials sit apart from the combined ones. -pub const slice_items = 256; +pub const slice_items = 1024; /// row_stream [n], streams [s, 4] (committed at [2]), offs [s, 2] (K, V from origin), paths [n, max_depth], depths. pub const Rows = struct { row_stream: u64, streams: u64, offs: u64, origin: u64, paths: u64, depths: u64 }; @@ -32,12 +33,12 @@ pub const TreeAttention = struct { paths_fn: Function, part_fns: [2]Function, // head dim 128, 256 shared_fns: [2]Function, - part_w1: [2]Function, // the same slices, one warp a block - shared_w1: [2]Function, + part_w1: [2][2]Function, // the same slices, one warp a block: [2 stages, 3 stages][head dim] + shared_w1: [2][2]Function, merge_fn: Function, /// The shared kernel's blocks a key head: a fixed grid walking every (row, group) (bits do not depend on it). shared_blocks: usize = 96, - /// The device's SMs: where the four-warp grid would leave some idle, a slice takes its own block (the same bits). + /// The device's SMs (0: four-warp blocks only): small grids take pipelined one-warp blocks (the same bits). sms: usize = 0, pub fn load(d: *const Driver) !TreeAttention { @@ -48,10 +49,12 @@ pub const TreeAttention = struct { for (&parts, part_smem) |*f, b| try f.allowDynamicShared(b); var shared: [2]Function = .{ try mod.function("tf_tree_shared"), try mod.function("tf_tree_shared_256") }; for (&shared, part_smem) |*f, b| try f.allowDynamicShared(b); - var parts1: [2]Function = .{ try mod.function("tf_tree_part_w1"), try mod.function("tf_tree_part_256_w1") }; - for (&parts1, part_smem_w1) |*f, b| try f.allowDynamicShared(b); - var shared1: [2]Function = .{ try mod.function("tf_tree_shared_w1"), try mod.function("tf_tree_shared_256_w1") }; - for (&shared1, part_smem_w1) |*f, b| try f.allowDynamicShared(b); + var parts1: [2][2]Function = .{ .{ try mod.function("tf_tree_part_w1"), try mod.function("tf_tree_part_256_w1") }, .{ try mod.function("tf_tree_part_w1s3"), try mod.function("tf_tree_part_256_w1s3") } }; + var shared1: [2][2]Function = .{ .{ try mod.function("tf_tree_shared_w1"), try mod.function("tf_tree_shared_256_w1") }, .{ try mod.function("tf_tree_shared_w1s3"), try mod.function("tf_tree_shared_256_w1s3") } }; + for (0..2) |st| for (0..2) |which| { + try parts1[st][which].allowDynamicShared(smemW1(@intCast(st + 2), which)); + try shared1[st][which].allowDynamicShared(smemW1(@intCast(st + 2), which)); + }; return .{ .mod = mod, .paths_fn = try mod.function("tf_tree_paths"), .part_fns = parts, .shared_fns = shared, .part_w1 = parts1, .shared_w1 = shared1, .merge_fn = try mod.function("tf_tree_merge") }; } @@ -64,15 +67,15 @@ pub const TreeAttention = struct { return (keys + span - 1) / span; } - /// fp32 partials of `rows` rows, `hk` key heads, `ng` groups (o [.., 16, dim], max, sum [.., 16]), then slices. - pub fn scratchBytes(rows: usize, hk: usize, ng: usize, dim: usize) usize { - return (rows * hk * ng + slice_items * slices) * 16 * (dim + 2) * 4; + /// fp32 partials of `rows` rows' `h` heads (`hk` key heads) in `ng` groups: o [.., dim], max, sum; then slices. + pub fn scratchBytes(rows: usize, h: usize, hk: usize, ng: usize, dim: usize) usize { + return (rows * h * ng + slice_items * slices * (h / hk)) * (dim + 2) * 4; } - pub fn scratch(base: u64, rows: usize, hk: usize, ng: usize, dim: usize) Scratch { - const parts = rows * hk * ng * 16; + pub fn scratch(base: u64, rows: usize, h: usize, hk: usize, ng: usize, dim: usize) Scratch { + const parts = rows * h * ng; const sb = base + parts * (dim + 2) * 4; - const sp = slice_items * slices * 16; + const sp = slice_items * slices * (h / hk); return .{ .o = base, .m = base + parts * dim * 4, .l = base + parts * dim * 4 + parts * 4, .so = sb, .sm = sb + sp * dim * 4, .sl = sb + sp * dim * 4 + sp * 4 }; } @@ -90,8 +93,8 @@ pub const TreeAttention = struct { if (hk == 0 or h % hk != 0 or h / hk > 16 or ng == 0 or (dim != 128 and dim != 256)) return error.TreeShape; const which: usize = if (dim == 256) 1 else 0; if (n == 0) return; - // one warp a block where four-warp blocks would leave SMs idle: each slice its own block, its own partial - const one = @min(n * ng, g.shared_blocks) * hk < g.sms and n * ng * hk <= slice_items; + // small grids (a few rows): each slice its own pipelined one-warp block and partial; else four-warp blocks + const one = g.sms > 0 and n * ng * hk <= slice_items; const o = if (one) [3]u64{ sc.so, sc.sm, sc.sl } else [3]u64{ sc.o, sc.m, sc.l }; var a: launch_.Args = .{}; for ([_]u64{ q, kn, vn, rows.origin, rows.offs, rows.row_stream, rows.streams, rows.paths, rows.depths, o[0], o[1], o[2] }) |v| a.add(v); @@ -108,10 +111,12 @@ pub const TreeAttention = struct { c.add(@as(i32, @intCast(n))); const per: usize = if (one) slices else 1; const block: u32 = if (one) 32 else 128; - const smem = if (one) part_smem_w1[which] else part_smem[which]; - try launch_.launch(if (one) g.shared_w1[which] else g.shared_fns[which], .{ .grid = .{ .x = @intCast(@min(n * ng, g.shared_blocks) * per), .y = @intCast(hk), .z = 1 }, .block = .{ .x = block }, .shared = smem }, s, &c); + // three stages where the one-warp grid fits one wave of the SMs (a block each), else two (two blocks an SM) + const deep: usize = if (@min(n * ng, g.shared_blocks) * slices * hk <= g.sms) 1 else 0; + const smem = if (one) smemW1(@intCast(deep + 2), which) else part_smem[which]; + try launch_.launch(if (one) g.shared_w1[deep][which] else g.shared_fns[which], .{ .grid = .{ .x = @intCast(@min(n * ng, g.shared_blocks) * per), .y = @intCast(hk), .z = 1 }, .block = .{ .x = block }, .shared = smem }, s, &c); // a row's own block: its boundary group and at most the next (a path is at most 128 keys) - try launch_.launch(if (one) g.part_w1[which] else g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2) * per) }, .block = .{ .x = block }, .shared = smem }, s, &a); + try launch_.launch(if (one) g.part_w1[deep][which] else g.part_fns[which], .{ .grid = .{ .x = @intCast(n), .y = @intCast(hk), .z = @intCast(@min(ng, 2) * per) }, .block = .{ .x = block }, .shared = smem }, s, &a); var b: launch_.Args = .{}; for ([_]u64{ o[0], o[1], o[2], rows.row_stream, rows.streams, rows.depths, out }) |v| b.add(v); b.add(@as(i32, @intCast(h))); diff --git a/zig/tests/cuda/tree_tests.zig b/zig/tests/cuda/tree_tests.zig index 3b8de8d5cb..fedfdb134f 100644 --- a/zig/tests/cuda/tree_tests.zig +++ b/zig/tests/cuda/tree_tests.zig @@ -73,10 +73,10 @@ fn call(x: *Set, s: cuda.Stream, g: *const tree.TreeAttention, d: usize, rows: [ const paths = try x.zero(s, n * tree.max_depth * 4); const depths = try x.zero(s, n * 4); try g.paths(s, try x.up(std.mem.sliceAsBytes(par)), paths, depths, n); - const sc_base = try x.zero(s, tree.TreeAttention.scratchBytes(n, hk, ng, d)); + const sc_base = try x.zero(s, tree.TreeAttention.scratchBytes(n, h, hk, ng, d)); const out = try x.zero(s, n * h * d * 2); const rws: tree.Rows = .{ .row_stream = try x.up(std.mem.sliceAsBytes(rs)), .streams = streams_dev, .offs = offs, .origin = origin, .paths = paths, .depths = depths }; - try g.run(s, try x.up(std.mem.sliceAsBytes(q)), try x.up(std.mem.sliceAsBytes(kn)), try x.up(std.mem.sliceAsBytes(vn)), rws, tree.TreeAttention.scratch(sc_base, n, hk, ng, d), out, n, h, hk, d, ng, scale); + try g.run(s, try x.up(std.mem.sliceAsBytes(q)), try x.up(std.mem.sliceAsBytes(kn)), try x.up(std.mem.sliceAsBytes(vn)), rws, tree.TreeAttention.scratch(sc_base, n, h, hk, ng, d), out, n, h, hk, d, ng, scale); try s.synchronize(); return x.down(u16, out, n * h * d); } @@ -189,7 +189,7 @@ fn case(gpu: check.Gpu, d: usize, h: usize, hk: usize, seed: u64, worst_all: *f6 const strided = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); try check.sameBytes("tree attention: one shared block", std.mem.sliceAsBytes(strided), std.mem.sliceAsBytes(got)); g.shared_blocks = 96; - const wide = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, 129, scale); + const wide = try call(&x, s, &g, d, all, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, 513, scale); try check.sameBytes("tree attention: groups for a 262k context", std.mem.sliceAsBytes(wide), std.mem.sliceAsBytes(got)); // one warp a block (a device with SMs to spare): stream 2's rows alone, its full groups too, the batch's bits g.sms = 1 << 20; @@ -203,6 +203,10 @@ fn case(gpu: check.Gpu, d: usize, h: usize, hk: usize, seed: u64, worst_all: *f6 const one_strided = try call(&x, s, &g, d, third, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); try check.sameBytes("tree attention: one warp a block, one shared block", std.mem.sliceAsBytes(one_strided), std.mem.sliceAsBytes(one)); g.shared_blocks = 96; + g.sms = 1; // more one-warp blocks than SMs: two tile stages instead of three + const two = try call(&x, s, &g, d, third, row_stream, parents, qs, kns, vns, h, hk, &streams, offs_dev, origin, ng, scale); + try check.sameBytes("tree attention: one warp a block, two stages", std.mem.sliceAsBytes(two), std.mem.sliceAsBytes(one)); + g.shared_blocks = 96; } pub fn run(gpu: check.Gpu) !void { From 591065e0978e20159bc19ba1f231b1c42c9dced9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:14:41 +0200 Subject: [PATCH 06/41] cuda/tree_attention: the merge issues a chunk's partial loads (16 slots) before folding them, the same arithmetic in the same order --- zig/kernels/cuda/tree_attention.cu | 55 +++++++++++++++++++----------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 310ab7d659..eea5d2645d 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -369,6 +369,8 @@ TF_TREE_PART(tf_tree_part_w1s3, 128, 32, 1, 3) TF_TREE_PART(tf_tree_part_256_w1s3, 256, 16, 1, 3) // out [W, H, D] bf16: block (row, key head, head) folds the row's groups in order (sliced: a group's 4 slices first). +constexpr int MERGE_CHUNK = 16; // partials per chunk: 16 groups, or 4 groups of 4 slices + extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ depths, @@ -377,30 +379,43 @@ extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( const int kh = blockIdx.y, G = h / hk, hd = blockIdx.z; // block (row, key head, head): each element alone const int total = streams[4 * row_stream[r] + 2] + depths[r]; const int groups = (total + SPAN - 1) / SPAN; + const int per = sliced ? 4 : 1, slots = groups * per; + const int64_t base = (r * hk + kh) * ng * per; for (int e = threadIdx.x; e < d; e += 256) { float m = -INFINITY, l = 0.0f, acc = 0.0f; - for (int gi = 0; gi < groups; ++gi) { - const int64_t part = ((r * hk + kh) * ng + gi) * (sliced ? 4 : 1); - float cm = -INFINITY, co = 0.0f, cl = 0.0f; - if (!sliced) { // four-warp blocks combined the group's slices themselves - cm = pm[part * G + hd], co = po[(part * G + hd) * d + e], cl = pl[part * G + hd]; - } else { // one-warp blocks left the 4 slices: their combine, in the same order - for (int w = 0; w < 4; ++w) cm = fmaxf(cm, pm[(part + w) * G + hd]); - for (int w = 0; w < 4; ++w) { - const float mw = pm[(part + w) * G + hd]; - if (mw == -INFINITY) continue; - const float aw = texp(__fsub_rn(mw, cm)); - co = __fmaf_rn(aw, po[((part + w) * G + hd) * d + e], co); - cl = __fmaf_rn(aw, pl[(part + w) * G + hd], cl); + for (int s0 = 0; s0 < slots; s0 += MERGE_CHUNK) { // a chunk's loads all issued before any is used + float cm_[MERGE_CHUNK], co_[MERGE_CHUNK], cl_[MERGE_CHUNK]; +#pragma unroll + for (int i = 0; i < MERGE_CHUNK; ++i) { + const int64_t q = (base + s0 + i) * G + hd; + const bool in = s0 + i < slots; + cm_[i] = in ? pm[q] : -INFINITY, cl_[i] = in ? pl[q] : 0.0f, co_[i] = in ? po[q * d + e] : 0.0f; + } +#pragma unroll + for (int i = 0; i < MERGE_CHUNK; ++i) { + if (s0 + i >= slots) break; + if (sliced && (i & 3)) continue; + float cm = -INFINITY, co = 0.0f, cl = 0.0f; + if (!sliced) { // four-warp blocks combined the group's slices themselves + cm = cm_[i], co = co_[i], cl = cl_[i]; + } else { // one-warp blocks left the 4 slices: their combine, in the same order + for (int w = 0; w < 4; ++w) cm = fmaxf(cm, cm_[i + w]); + for (int w = 0; w < 4; ++w) { + const float mw = cm_[i + w]; + if (mw == -INFINITY) continue; + const float aw = texp(__fsub_rn(mw, cm)); + co = __fmaf_rn(aw, co_[i + w], co); + cl = __fmaf_rn(aw, cl_[i + w], cl); + } } + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : texp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? texp(__fsub_rn(cm, next)) : 0.0f; + acc = __fmaf_rn(acc, a, __fmul_rn(co, b)); + l = __fmaf_rn(l, a, __fmul_rn(cl, b)); + m = next; } - const bool active = cl > 0.0f; - const float next = active ? fmaxf(m, cm) : m; - const float a = active ? (m == -INFINITY ? 0.0f : texp(__fsub_rn(m, next))) : 1.0f; - const float b = active ? texp(__fsub_rn(cm, next)) : 0.0f; - acc = __fmaf_rn(acc, a, __fmul_rn(co, b)); - l = __fmaf_rn(l, a, __fmul_rn(cl, b)); - m = next; } out[(r * h + kh * G + hd) * d + e] = __float2bfloat16_rn(__fdiv_rn(acc, l)); } From 1dd63e39a9180054df02b59a5d86b2d0b9b22c04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:41:51 +0200 Subject: [PATCH 07/41] cuda: act_rows, gdn_prep and qk_rows for every family (SwiGLU, output gates, DFlash2's conv, Gated DeltaNet inputs, QK norm and RoPE), checked against float64 --- zig/build/cuda.zig | 3 + zig/kernels/cuda/act_rows.cu | 142 ++++++++++++++++++ zig/kernels/cuda/gdn_prep.cu | 99 +++++++++++++ zig/kernels/cuda/qk_rows.cu | 194 +++++++++++++++++++++++++ zig/src/core/act_ref.zig | 11 ++ zig/src/core/gdn_prep_ref.zig | 45 ++++++ zig/src/core/qk_ref.zig | 30 ++++ zig/src/core/root.zig | 4 + zig/src/cuda/act_rows.zig | 95 ++++++++++++ zig/src/cuda/gdn_prep.zig | 72 +++++++++ zig/src/cuda/kernels.zig | 3 + zig/src/cuda/qk_rows.zig | 107 ++++++++++++++ zig/src/cuda/root.zig | 6 + zig/tests/cuda/act_rows_tests.zig | 132 +++++++++++++++++ zig/tests/cuda/gdn_prep_tests.zig | 156 ++++++++++++++++++++ zig/tests/cuda/main.zig | 9 ++ zig/tests/cuda/norm_rows_tests.zig | 24 +-- zig/tests/cuda/qk_rows_tests.zig | 226 +++++++++++++++++++++++++++++ 18 files changed, 1346 insertions(+), 12 deletions(-) create mode 100644 zig/kernels/cuda/act_rows.cu create mode 100644 zig/kernels/cuda/gdn_prep.cu create mode 100644 zig/kernels/cuda/qk_rows.cu create mode 100644 zig/src/core/gdn_prep_ref.zig create mode 100644 zig/src/core/qk_ref.zig create mode 100644 zig/src/cuda/act_rows.zig create mode 100644 zig/src/cuda/gdn_prep.zig create mode 100644 zig/src/cuda/qk_rows.zig create mode 100644 zig/tests/cuda/act_rows_tests.zig create mode 100644 zig/tests/cuda/gdn_prep_tests.zig create mode 100644 zig/tests/cuda/qk_rows_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..35a77ebc58 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -45,6 +45,9 @@ const kernels = [_]Kernel{ .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle + .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) + .{ .name = "gdn_prep", .flags = &.{"-O3"} }, // ours: Gated DeltaNet conv, norms, g and beta (cuda/gdn_prep.zig) + .{ .name = "qk_rows", .flags = &.{"-O3"} }, // ours: QK norm, RoPE and the cache write (cuda/qk_rows.zig) .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu new file mode 100644 index 0000000000..515cc6bf67 --- /dev/null +++ b/zig/kernels/cuda/act_rows.cu @@ -0,0 +1,142 @@ +// Activation rows for every CUDA family (float64 references: core/act_ref.zig): SwiGLU, output gates, group sums. + +// A thread 8 values with 16-byte loads; a row's bits depend on that row alone. + +#include +#include + +namespace tf_act { + +constexpr int NT = 256; + +__device__ __forceinline__ void unpack8(const uint4& u, float* f) { + const __nv_bfloat162* h = reinterpret_cast(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 t = __bfloat1622float2(h[i]); + f[2 * i] = t.x; + f[2 * i + 1] = t.y; + } +} + +__device__ __forceinline__ uint4 pack8(const float* f) { + uint4 u; + __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(f[2 * i], f[2 * i + 1]); + return u; +} + +__device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } + +// A thread's 8 stored values in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). +__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) { + float f[8]; + unpack8(packed, f); + float s = 0.f; + if (live) +#pragma unroll + for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4)); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2)); + return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1)); +} + +} // namespace tf_act + +using namespace tf_act; + +// out[r, i] = bf16(silu(gate[r * g_ld + i]) * up[r * u_ld + i]) for i < n (n % 64 == 0); xs (or null) out's sums. +extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows(const __nv_bfloat16* __restrict__ gate, long long g_ld, + const __nv_bfloat16* __restrict__ up, long long u_ld, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, + int n) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const bool live = 8 * c < n; + float g[8], u[8], o[8]; + if (live) { + unpack8(*reinterpret_cast(gate + r * g_ld + 8 * c), g); + unpack8(*reinterpret_cast(up + r * u_ld + 8 * c), u); + } +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = live ? __fmul_rn(__fmul_rn(g[i], sigmoid(g[i])), u[i]) : 0.f; + const uint4 op = pack8(o); + if (live) *reinterpret_cast(out + r * n + 8 * c) = op; + if (xs != nullptr) { + const float s = group_sum8(op, live); + if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s; + } +} + +// out[r, h * d + e] = bf16(o[r, h * d + e] * sigmoid(g[r * g_ld + h * g_hs + e])) (d % 64 == 0); xs (or null) its sums. +extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat16* __restrict__ o, + const __nv_bfloat16* __restrict__ g, long long g_ld, + long long g_hs, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, int heads, int d) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int n = heads * d; + const bool live = 8 * c < n; + const int h = 8 * c / d, e = 8 * c % d; + float ov[8], gv[8], y[8]; + if (live) { + unpack8(*reinterpret_cast(o + r * n + 8 * c), ov); + unpack8(*reinterpret_cast(g + r * g_ld + h * g_hs + e), gv); + } +#pragma unroll + for (int i = 0; i < 8; ++i) y[i] = live ? __fmul_rn(ov[i], sigmoid(gv[i])) : 0.f; + const uint4 yp = pack8(y); + if (live) *reinterpret_cast(out + r * n + 8 * c) = yp; + if (xs != nullptr) { + const float s = group_sum8(yp, live); + if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s; + } +} + +// xs[r, j] = the sum of x[r * ld + 64 j ..][0 .. 64] (k % 64 == 0), in group_sum8's order. +extern "C" __global__ void __launch_bounds__(NT) tf_group_sums_rows(const __nv_bfloat16* __restrict__ x, long long ld, + float* __restrict__ xs, int k) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + const bool live = 8 * c < k; + const uint4 v = live ? *reinterpret_cast(x + r * ld + 8 * c) : make_uint4(0, 0, 0, 0); + const float s = group_sum8(v, live); + if (live && (c & 7) == 0) xs[r * (k / 64) + c / 8] = s; +} + +// DFlash2's two-tap conv: out = bf16(x k0 + x[r - 1] k1) [+ res], no r - 1 at a seg block's start, k = bf16(b + dyn). +extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ dyn, + const __nv_bfloat16* __restrict__ base, + const __nv_bfloat16* __restrict__ res, + __nv_bfloat16* __restrict__ out, int d, int gs, int branch, + int seg) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + if (8 * c >= d) return; + const int g = d / gs, grp = 8 * c / gs; + float xv[8], pv[8], b0[8], b1[8], y[8]; + unpack8(*reinterpret_cast(x + r * d + 8 * c), xv); + if (r % seg > 0) unpack8(*reinterpret_cast(x + (r - 1) * d + 8 * c), pv); + else +#pragma unroll + for (int i = 0; i < 8; ++i) pv[i] = 0.f; + unpack8(*reinterpret_cast(base + (int64_t)(branch * 2) * d + 8 * c), b0); + unpack8(*reinterpret_cast(base + (int64_t)(branch * 2 + 1) * d + 8 * c), b1); + const float d0 = __bfloat162float(dyn[((r * 2 + branch) * 2) * g + grp]); + const float d1 = __bfloat162float(dyn[((r * 2 + branch) * 2 + 1) * g + grp]); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float k0 = __bfloat162float(__float2bfloat16_rn(__fadd_rn(b0[i], d0))); + const float k1 = __bfloat162float(__float2bfloat16_rn(__fadd_rn(b1[i], d1))); + y[i] = __fmaf_rn(pv[i], k1, __fmul_rn(xv[i], k0)); + } + if (res != nullptr) { + float rv[8]; + unpack8(*reinterpret_cast(res + r * d + 8 * c), rv); +#pragma unroll + for (int i = 0; i < 8; ++i) y[i] = __fadd_rn(rv[i], __bfloat162float(__float2bfloat16_rn(y[i]))); + } + *reinterpret_cast(out + r * d + 8 * c) = pack8(y); +} diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu new file mode 100644 index 0000000000..04f330f4e7 --- /dev/null +++ b/zig/kernels/cuda/gdn_prep.cu @@ -0,0 +1,99 @@ +// Gated DeltaNet inputs for every CUDA family (float64 reference: core/gdn_prep_ref.zig): conv, norms, g, beta. + +// A warp a head (q heads, k heads, then v heads), lane l holding channels J l .. J l + J - 1; a row's bits are its own. + +#include +#include +#include + +namespace tf_gdn_prep { + +constexpr int MAXT = 8; // conv taps at most + +struct Io { + const __nv_bfloat16* qkv; // [rows, qkv_ld]: q (kh dk), k (kh dk), v (vh dk) channels + long long qkv_ld; + const __nv_bfloat16* state; // [streams * keep, c]: each stream's kept conv rows + const __nv_bfloat16* conv; // [c, keep + 1] taps a channel + const int* win; // [rows, keep + 1]: < keep a kept row, else qkv row win - keep + const int* sid; // [rows] stream (null: stream 0) + const __nv_bfloat16* a; // [rows, a_ld] the v heads' dt inputs + long long a_ld; + const __nv_bfloat16* b; // [rows, b_ld] the v heads' beta inputs + long long b_ld; + const float* a_log; // [vh] + const float* dt_bias; // [vh] + __nv_bfloat16* q; // [rows, kh, dk] + __nv_bfloat16* k; // [rows, kh, dk] + __nv_bfloat16* v; // [rows, vh, dk] + float* g; // [rows, vh] + float* beta; // [rows, vh] +}; + +__device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } + +template +__device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int keep, float eps) { + constexpr int DK = 32 * J; + const int lane = threadIdx.x & 31; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y * (blockDim.x / 32) + threadIdx.x / 32; + if (hd >= 2 * kh + vh) return; + const int ch = hd * DK + J * lane; + const int64_t srow = (io.sid != nullptr ? static_cast(io.sid[r]) : 0) * keep; + int src[MAXT]; +#pragma unroll + for (int t = 0; t < MAXT; ++t) src[t] = t <= keep ? io.win[r * (keep + 1) + t] : 0; + float acc[J]; +#pragma unroll + for (int i = 0; i < J; ++i) acc[i] = 0.f; +#pragma unroll + for (int t = 0; t < MAXT; ++t) { + if (t > keep) break; + const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; +#pragma unroll + for (int i = 0; i < J; ++i) + acc[i] = __fmaf_rn(__bfloat162float(x[ch + i]), __bfloat162float(io.conv[(ch + i) * (keep + 1) + t]), acc[i]); + } + float cv[J]; +#pragma unroll + for (int i = 0; i < J; ++i) cv[i] = __bfloat162float(__float2bfloat16_rn(__fmul_rn(acc[i], sigmoid(acc[i])))); + if (hd < 2 * kh) { // q: l2norm / sqrt(dk), k: l2norm (rms with eps, scaled) + float part = 0.f; +#pragma unroll + for (int i = 0; i < J; ++i) part = __fmaf_rn(cv[i], cv[i], part); + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + const float inv = rsqrtf(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps)); + const bool is_q = hd < kh; + const float scale = is_q ? 1.f / DK : rsqrtf(static_cast(DK)); + __nv_bfloat16* out = (is_q ? io.q : io.k) + (r * kh + (is_q ? hd : hd - kh)) * DK + J * lane; +#pragma unroll + for (int i = 0; i < J; ++i) out[i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(cv[i], inv), scale)); + return; + } + const int hv = hd - 2 * kh; +#pragma unroll + for (int i = 0; i < J; ++i) io.v[(r * vh + hv) * DK + J * lane + i] = __float2bfloat16_rn(cv[i]); + if (lane == 0) { + const float x = __fadd_rn(__bfloat162float(io.a[r * io.a_ld + hv]), io.dt_bias[hv]); + const float sp = x > 20.f ? x : log1pf(expf(x)); + io.g[r * vh + hv] = expf(-__fmul_rn(expf(io.a_log[hv]), sp)); + io.beta[r * vh + hv] = sigmoid(__bfloat162float(io.b[r * io.b_ld + hv])); + } +} + +} // namespace tf_gdn_prep + +using namespace tf_gdn_prep; + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_64(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<2>(io, c, kh, vh, keep, eps); +} + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_128(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<4>(io, c, kh, vh, keep, eps); +} + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_256(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<8>(io, c, kh, vh, keep, eps); +} diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu new file mode 100644 index 0000000000..146e14017c --- /dev/null +++ b/zig/kernels/cuda/qk_rows.cu @@ -0,0 +1,194 @@ +// Attention inputs for every CUDA family (float64 references: core/qk_ref.zig): QK norm, RoPE and the cache write. + +// A block a (row, head), a thread an element; a head's norm sums in one fixed order, never depending on the row count. + +#include +#include + +namespace tf_qk { + +constexpr int MOST = 256; // head_dim at most + +// Where a row's heads come from: q, k and v (null: none) at row stride *_ld and head stride *_hs (elements). +struct QkSrc { + const __nv_bfloat16* q; + long long q_ld, q_hs; + const __nv_bfloat16* k; + long long k_ld, k_hs; + const __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +// The q and k norm weights: fp32 or bf16 (null q: no norm). +struct QkW { + const void* q; + const void* k; + int bf16; +}; + +__device__ __forceinline__ float wv(const QkW& w, bool is_k, int i) { + const void* p = is_k ? w.k : w.q; + return w.bf16 ? __bfloat162float(static_cast(p)[i]) : static_cast(p)[i]; +} + +// cos and sin of p's angle on rotated dim d: from tables [positions, rot], or computed from inv [rot / 2]. +__device__ __forceinline__ void rope_at(const float* cosv, const float* sinv, const float* inv, int p, int rot, int d, + float& c, float& sn) { + if (inv != nullptr) { + const float a = __fmul_rn(static_cast(p), inv[d < rot / 2 ? d : d - rot / 2]); + c = cosf(a); + sn = sinf(a); + } else { + c = cosv[static_cast(p) * rot + d]; + sn = sinv[static_cast(p) * rot + d]; + } +} + +// Where a row's q, k and v (null: none) go, at row stride *_ld and head stride *_hs (elements). +struct QkDst { + __nv_bfloat16* q; + long long q_ld, q_hs; + __nv_bfloat16* k; + long long k_ld, k_hs; + __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +__device__ __forceinline__ const __nv_bfloat16* head_src(const QkSrc& s, int64_t r, int hd, int h, int hk) { + if (hd < h) return s.q + r * s.q_ld + hd * s.q_hs; + if (hd < h + hk) return s.k + r * s.k_ld + (hd - h) * s.k_hs; + return s.v + r * s.v_ld + (hd - h - hk) * s.v_hs; +} + +} // namespace tf_qk + +using namespace tf_qk; + +// QkSrc's heads -> q, k, v (bf16); q, k normed by nw, rotate-half RoPE on `rot` dims; k, v also into kc/vc. +extern "C" __global__ void __launch_bounds__(MOST) tf_qk_rows(const QkSrc in, const QkW nw, const float* cosv, + const float* sinv, const float* inv, int rot, const int* pos, + const int* slot, const QkDst dst, __nv_bfloat16* kc, __nv_bfloat16* vc, + long long len, int h, int hk, float eps) { + __shared__ float warps[MOST / 32]; + const int d = threadIdx.x; + const int dim = blockDim.x; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y; + const __nv_bfloat16* src = head_src(in, r, hd, h, hk); + const float x = __bfloat162float(src[d]); + const int p = pos[r]; + const int64_t pl = p < len ? p : p % len; // a 64-bit modulo only past the cache's end + const int64_t at = slot != nullptr ? static_cast(slot[r]) * len + pl : pl; + if (hd >= h + hk) { // v: as is + const int k = hd - h - hk; + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + d] = src[d]; + if (vc != nullptr) vc[(at * hk + k) * dim + d] = src[d]; + return; + } + float y = x; + float yp = 0.f; + const int half = rot / 2; + const int partner = d < half ? d + half : d - half; + if (nw.q != nullptr) { + float part = __fmul_rn(x, x); + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + if ((d & 31) == 0) warps[d >> 5] = part; + __syncthreads(); + float s = 0.f; + for (int w = 0; w < dim / 32; ++w) s = __fadd_rn(s, warps[w]); + const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps)); + y = __fmul_rn(__fmul_rn(x, rstd), wv(nw, hd >= h, d)); + if (d < rot) yp = __fmul_rn(__fmul_rn(__bfloat162float(src[partner]), rstd), wv(nw, hd >= h, partner)); + } else if (d < rot) { + yp = __bfloat162float(src[partner]); + } + if ((cosv != nullptr || inv != nullptr) && d < rot) { + float c, sn; + rope_at(cosv, sinv, inv, p, rot, d, c, sn); + y = __fadd_rn(__fmul_rn(y, c), __fmul_rn(d < half ? -yp : yp, sn)); + } + const __nv_bfloat16 out = __float2bfloat16_rn(y); + if (hd < h) { + dst.q[r * dst.q_ld + hd * dst.q_hs + d] = out; + } else { + const int k = hd - h; + dst.k[r * dst.k_ld + k * dst.k_hs + d] = out; + if (kc != nullptr) kc[(at * hk + k) * dim + d] = out; + } +} + +// tf_qk_rows at dim 128, no or full rotation: a warp a head, lane l dims l + 32 j (same fold order, partner j ^ 2). +extern "C" __global__ void __launch_bounds__(128) tf_qk_rows_w(const QkSrc in, const QkW nw, + const float* cosv, const float* sinv, const float* inv, int rot, + const int* pos, const int* slot, const QkDst dst, + __nv_bfloat16* kc, __nv_bfloat16* vc, long long len, + int h, int hk, float eps) { + constexpr int D = 128, J = D / 32; + const int lane = threadIdx.x & 31; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y * 4 + (threadIdx.x >> 5); + if (hd >= h + (in.v != nullptr ? 2 * hk : hk)) return; + const __nv_bfloat16* src = head_src(in, r, hd, h, hk); + const int p = pos[r]; + __nv_bfloat16 raw[J]; +#pragma unroll + for (int j = 0; j < J; ++j) raw[j] = src[lane + 32 * j]; + // every load up front (the rotation's row waits on pos, then on memory) + float c[J], sn[J], w[J]; + const bool qk = hd < h + hk; +#pragma unroll + for (int j = 0; j < J; ++j) { + c[j] = sn[j] = 0.f; + if (qk && rot != 0 && (cosv != nullptr || inv != nullptr)) rope_at(cosv, sinv, inv, p, D, lane + 32 * j, c[j], sn[j]); + w[j] = qk && nw.q != nullptr ? wv(nw, hd >= h, lane + 32 * j) : 0.f; + } + const int64_t pl = p < len ? p : p % len; + const int64_t at = slot != nullptr ? static_cast(slot[r]) * len + pl : pl; + if (hd >= h + hk) { // v: as is + const int k = hd - h - hk; +#pragma unroll + for (int j = 0; j < J; ++j) { + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + lane + 32 * j] = raw[j]; + if (vc != nullptr) vc[(at * hk + k) * D + lane + 32 * j] = raw[j]; + } + return; + } + float x[J], y[J], yp[J]; +#pragma unroll + for (int j = 0; j < J; ++j) x[j] = __bfloat162float(raw[j]), y[j] = x[j], yp[j] = 0.f; + if (nw.q != nullptr) { + float part[J]; +#pragma unroll + for (int j = 0; j < J; ++j) part[j] = __fmul_rn(x[j], x[j]); + for (int m = 16; m > 0; m >>= 1) +#pragma unroll + for (int j = 0; j < J; ++j) part[j] = __fadd_rn(part[j], __shfl_xor_sync(0xFFFFFFFFu, part[j], m)); + float s = 0.f; +#pragma unroll + for (int j = 0; j < J; ++j) s = __fadd_rn(s, part[j]); + const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps)); +#pragma unroll + for (int j = 0; j < J; ++j) { + y[j] = __fmul_rn(__fmul_rn(x[j], rstd), w[j]); + if (rot != 0) yp[j] = __fmul_rn(__fmul_rn(x[j ^ 2], rstd), w[j ^ 2]); + } + } else if (rot != 0) { +#pragma unroll + for (int j = 0; j < J; ++j) yp[j] = x[j ^ 2]; + } + if (rot != 0 && (cosv != nullptr || inv != nullptr)) { +#pragma unroll + for (int j = 0; j < J; ++j) y[j] = __fadd_rn(__fmul_rn(y[j], c[j]), __fmul_rn(j < 2 ? -yp[j] : yp[j], sn[j])); + } +#pragma unroll + for (int j = 0; j < J; ++j) { + const __nv_bfloat16 out = __float2bfloat16_rn(y[j]); + if (hd < h) { + dst.q[r * dst.q_ld + hd * dst.q_hs + lane + 32 * j] = out; + } else { + const int k = hd - h; + dst.k[r * dst.k_ld + k * dst.k_hs + lane + 32 * j] = out; + if (kc != nullptr) kc[(at * hk + k) * D + lane + 32 * j] = out; + } + } +} diff --git a/zig/src/core/act_ref.zig b/zig/src/core/act_ref.zig index ccab2c254e..f56f282175 100644 --- a/zig/src/core/act_ref.zig +++ b/zig/src/core/act_ref.zig @@ -5,6 +5,16 @@ pub fn sigmoid(x: f64) f64 { return 1 / (1 + @exp(-x)); } +/// silu(gate) * up. +pub fn swiglu(gate: f64, up: f64) f64 { + return gate * sigmoid(gate) * up; +} + +/// DFlash2's dynamic conv at one value: x k0 + prev k1 with the taps as rounded (k_t = bf16(base_t + dyn_t)). +pub fn dconv(x: f64, prev: f64, k0: f64, k1: f64) f64 { + return x * k0 + prev * k1; +} + /// A 4-bit affine table's value: q * scale + bias. pub fn affine4(q: u4, scale: f64, bias: f64) f64 { return @as(f64, @floatFromInt(q)) * scale + bias; @@ -18,5 +28,6 @@ pub fn tolerance(want: f64) f64 { test "activations against hand values" { try std.testing.expectEqual(@as(f64, 0.5), sigmoid(0)); try std.testing.expectEqual(@as(f64, 7 * 0.5 - 1), affine4(7, 0.5, -1)); + try std.testing.expectApproxEqRel(@as(f64, 2 * sigmoid(2) * 3), swiglu(2, 3), 1e-15); try std.testing.expect(tolerance(1) > tolerance(0.5)); } diff --git a/zig/src/core/gdn_prep_ref.zig b/zig/src/core/gdn_prep_ref.zig new file mode 100644 index 0000000000..59a812d6df --- /dev/null +++ b/zig/src/core/gdn_prep_ref.zig @@ -0,0 +1,45 @@ +//! Gated DeltaNet inputs in float64, the reference every backend's GDN prep kernels are checked against. +const std = @import("std"); + +/// A head's conv over `taps` rows (x [taps][dk], w [dk][taps]), then SiLU -> c; err: a channel's fp32 bound. +pub fn conv(x: []const []const f64, w: []const f64, taps: usize, c: []f64, err: []f64) void { + for (c, err, 0..) |*o, *e, i| { + var acc: f64 = 0; + var mag: f64 = 0; + for (0..taps) |t| { + acc += x[t][i] * w[i * taps + t]; + mag += @abs(x[t][i] * w[i * taps + t]); + } + o.* = acc / (1 + @exp(-acc)); + e.* = 1.1 * @as(f64, @floatFromInt(taps)) * std.math.pow(f64, 2, -24) * mag; // silu' is at most 1.1 + } +} + +/// q (is_q) or k from a head's c: c / sqrt(mean(c^2) + eps) * (1 / dk for q, 1 / sqrt(dk) for k). +pub fn norm(c: []const f64, is_q: bool, eps: f64, out: []f64) f64 { + var ss: f64 = 0; + for (c) |v| ss += v * v; + const dk: f64 = @floatFromInt(c.len); + const f = 1 / @sqrt(ss / dk + eps) * (if (is_q) 1 / dk else 1 / @sqrt(dk)); + for (out, c) |*o, v| o.* = v * f; + return f; +} + +/// g = exp(-exp(a_log) * softplus(a + dt_bias)). +pub fn gate(a: f64, dt_bias: f64, a_log: f64) f64 { + const x = a + dt_bias; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + return @exp(-@exp(a_log) * sp); +} + +test "conv, norms and the gate against hand values" { + var c: [2]f64 = undefined; + var e: [2]f64 = undefined; + conv(&.{ &.{ 1, 2 }, &.{ 3, 4 } }, &.{ 1, 0, 0, 1 }, 2, &c, &e); // channel 0 tap 0, channel 1 tap 1 + try std.testing.expectApproxEqRel(@as(f64, 1.0 / (1.0 + @exp(-1.0))), c[0], 1e-15); + try std.testing.expectApproxEqRel(@as(f64, 4.0 / (1.0 + @exp(-4.0))), c[1], 1e-15); + var o: [2]f64 = undefined; + _ = norm(&.{ 3, 4 }, false, 0, &o); + try std.testing.expectApproxEqRel(@as(f64, 0.6), o[0], 1e-15); + try std.testing.expectApproxEqRel(@as(f64, @exp(-@exp(0.5) * 25.0)), gate(25, 0, 0.5), 1e-15); +} diff --git a/zig/src/core/qk_ref.zig b/zig/src/core/qk_ref.zig new file mode 100644 index 0000000000..b1763d9f34 --- /dev/null +++ b/zig/src/core/qk_ref.zig @@ -0,0 +1,30 @@ +//! QK norm and RoPE of one head in float64, the reference for every backend's attention inputs (cuda/qk_rows.zig). +const std = @import("std"); + +/// One head: RMS-normed by w (none: as is), then rotate-half RoPE on its first `rot` dims (no cos/sin: none). +pub fn head(x: []const f64, w: ?[]const f64, eps: f64, cos: ?[]const f64, sin: ?[]const f64, rot: usize, y: []f64) void { + const n = x.len; + var ss: f64 = 0; + for (x) |v| ss += v * v; + const inv = if (w != null) 1 / @sqrt(ss / @as(f64, @floatFromInt(n)) + eps) else 1; + var z: [256]f64 = undefined; + for (x, 0..) |v, i| z[i] = v * inv * (if (w) |ww| ww[i] else 1); + const half = rot / 2; + for (y, 0..) |*o, i| { + if (cos == null or i >= rot) { + o.* = z[i]; + continue; + } + const p = if (i < half) i + half else i - half; + const rotated = if (i < half) -z[p] else z[p]; + o.* = z[i] * cos.?[i] + rotated * sin.?[i]; + } +} + +test "rotate-half pairs (d, d + rot / 2) and leave the rest" { + var y: [4]f64 = undefined; + head(&.{ 1, 2, 3, 4 }, null, 0, &.{ 0, 0 }, &.{ 1, 1 }, 2, &y); // angle 90 degrees on the first pair + try std.testing.expectEqualSlices(f64, &.{ -2, 1, 3, 4 }, &y); + head(&.{ 2, 2, 2, 2 }, &.{ 1, 1, 1, 3 }, 0, null, null, 0, &y); + try std.testing.expectEqualSlices(f64, &.{ 1, 1, 1, 3 }, &y); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 122ffe9e4c..01934d154d 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -28,6 +28,8 @@ pub const modelopt = @import("modelopt.zig"); pub const picture = @import("picture/picture.zig"); /// Qwen-VL image input every Qwen3.5-VL checkpoint shares: settings, geometry, M-RoPE positions, the tower's tables. pub const qwen_vl = @import("qwen_vl/qwen_vl.zig"); +pub const gdn_prep_ref = @import("gdn_prep_ref.zig"); +pub const qk_ref = @import("qk_ref.zig"); test { _ = safetensors; @@ -41,6 +43,8 @@ test { _ = norms_ref; _ = act_ref; _ = contract; + _ = gdn_prep_ref; + _ = qk_ref; _ = exl3_format; _ = exlr8; _ = eaq_container; diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig new file mode 100644 index 0000000000..832844ce01 --- /dev/null +++ b/zig/src/cuda/act_rows.zig @@ -0,0 +1,95 @@ +//! act_rows.cu's activation rows for every CUDA family: SwiGLU, sigmoid output gates and 64-value group sums. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +const threads = 256; + +fn aligned(ptrs: []const u64, strides: []const usize) bool { + for (ptrs) |p| if (p % 16 != 0) return false; + for (strides) |v| if (v % 8 != 0) return false; + return true; +} + +fn grid(rows: usize, n: usize) launch_.Dim3 { + return .{ .x = @intCast(rows), .y = @intCast((n / 8 + threads - 1) / threads), .z = 1 }; +} + +pub const ActRows = struct { + mod: Module, + swiglu_fn: Function, + gate_fn: Function, + sums_fn: Function, + dconv_fn: Function, + + pub fn load(d: *const Driver) !ActRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.act_rows); + errdefer mod.unload(); + return .{ .mod = mod, .swiglu_fn = try mod.function("tf_swiglu_rows"), .gate_fn = try mod.function("tf_gate_rows"), .sums_fn = try mod.function("tf_group_sums_rows"), .dconv_fn = try mod.function("tf_dconv_rows") }; + } + + pub fn unload(g: *ActRows) void { + g.mod.unload(); + } + + /// out [rows, n] = bf16(silu(gate) * up), rows at strides g_ld, u_ld; xs (or null): out's 64-value sums. + pub fn swiglu(g: *const ActRows, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { + if (n % 64 != 0 or !aligned(&.{ gate, up, out }, &.{ g_ld, u_ld })) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(gate); + a.add(@as(i64, @intCast(g_ld))); + a.add(up); + a.add(@as(i64, @intCast(u_ld))); + a.add(out); + a.add(xs orelse 0); + a.add(@as(i32, @intCast(n))); + try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n), .block = .{ .x = threads } }, s, &a); + } + + /// out = bf16(o * sigmoid(gate)), head h's gate at r * g_ld + h * g_hs ([q | gate]: q + d, 2 d); xs: sums. + pub fn outputGate(g: *const ActRows, s: Stream, o: u64, gate_: u64, g_ld: usize, g_hs: usize, out: u64, xs: ?u64, rows: usize, heads: usize, d: usize) !void { + if (d % 64 != 0 or !aligned(&.{ o, gate_, out }, &.{ g_ld, g_hs })) return error.RowShape; + if (rows == 0 or heads == 0) return; + var a: launch_.Args = .{}; + a.add(o); + a.add(gate_); + a.add(@as(i64, @intCast(g_ld))); + a.add(@as(i64, @intCast(g_hs))); + a.add(out); + a.add(xs orelse 0); + a.add(@as(i32, @intCast(heads))); + a.add(@as(i32, @intCast(d))); + try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d), .block = .{ .x = threads } }, s, &a); + } + + /// xs [rows, k / 64]: x's rows (stride ld) summed over each 64 values. k % 64 == 0. + pub fn groupSums(g: *const ActRows, s: Stream, x: u64, ld: usize, xs: u64, rows: usize, k: usize) !void { + if (k % 64 != 0 or !aligned(&.{x}, &.{ld})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(ld))); + a.add(xs); + a.add(@as(i32, @intCast(k))); + try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k), .block = .{ .x = threads } }, s, &a); + } + + /// DFlash2's two-tap conv over [rows, d] bf16: dyn [rows, 2, 2, d / gs], base [2, 2, d]; res (or null) added. + pub fn dconv(g: *const ActRows, s: Stream, x: u64, dyn: u64, base: u64, res: ?u64, out: u64, rows: usize, d: usize, gs: usize, branch: usize, seg: usize) !void { + if (d % 8 != 0 or gs % 8 != 0 or d % gs != 0 or branch > 1 or seg == 0 or !aligned(&.{ x, base, res orelse 0, out }, &.{})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ x, dyn, base, res orelse 0, out }) |v| a.add(v); + for ([_]usize{ d, gs, branch, seg }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d), .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/cuda/gdn_prep.zig b/zig/src/cuda/gdn_prep.zig new file mode 100644 index 0000000000..d11f869096 --- /dev/null +++ b/zig/src/cuda/gdn_prep.zig @@ -0,0 +1,72 @@ +//! gdn_prep.cu's Gated DeltaNet inputs for every CUDA family: each row's conv window, SiLU, q/k norms, g and beta. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// The kernel's buffers (addresses; strides in elements), as gdn_prep.cu's Io. +pub const Io = extern struct { + qkv: u64, // [rows, qkv_ld] bf16: q (kh dk), k (kh dk), v (vh dk) channels + qkv_ld: i64, + state: u64, // [streams * keep, c] bf16: each stream's kept conv rows + conv: u64, // [c, keep + 1] bf16 taps a channel + win: u64, // [rows, keep + 1] i32: < keep a kept row of the row's stream, else qkv row (win - keep) + sid: u64 = 0, // [rows] i32 stream (0: every row stream 0) + a: u64, // [rows, a_ld] bf16 + a_ld: i64, + b: u64, // [rows, b_ld] bf16 + b_ld: i64, + a_log: u64, // [vh] fp32 + dt_bias: u64, // [vh] fp32 + q: u64, // [rows, kh, dk] bf16 + k: u64, + v: u64, // [rows, vh, dk] bf16 + g: u64, // [rows, vh] fp32 + beta: u64, +}; + +/// Channels a row (c = 2 kh dk + vh dk), heads and kept conv rows (taps - 1). +pub const Shape = struct { kh: usize, vh: usize, dk: usize, keep: usize, eps: f32 = 1e-6 }; + +pub const GdnPrep = struct { + mod: Module, + fns: [3]Function, // dk 64, 128, 256 + + pub fn load(d: *const Driver) !GdnPrep { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.gdn_prep); + errdefer mod.unload(); + return .{ .mod = mod, .fns = .{ try mod.function("tf_gdn_prep_64"), try mod.function("tf_gdn_prep_128"), try mod.function("tf_gdn_prep_256") } }; + } + + pub fn unload(g: *GdnPrep) void { + g.mod.unload(); + } + + /// Per row's conv window c = silu(conv): q = l2norm(c) / sqrt(dk), k = l2norm(c), v = c, g and beta. + pub fn run(g: *const GdnPrep, s: Stream, io: Io, rows: usize, sh: Shape) !void { + const which: usize = switch (sh.dk) { + 64 => 0, + 128 => 1, + 256 => 2, + else => return error.GdnShape, + }; + if (sh.keep + 1 > 8 or sh.keep == 0) return error.GdnShape; + if (rows == 0) return; + const heads = 2 * sh.kh + sh.vh; + var a: launch_.Args = .{}; + a.add(io); + a.add(@as(i32, @intCast(heads * sh.dk))); // c: a kept row's channels + a.add(@as(i32, @intCast(sh.kh))); + a.add(@as(i32, @intCast(sh.vh))); + a.add(@as(i32, @intCast(sh.keep))); + a.add(sh.eps); + try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + 3) / 4), .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..3326909329 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -56,6 +56,9 @@ pub const ck_lane: []const u8 = if (available) &Blob("fatbin_ck_lane").bytes els pub const ck_experts: []const u8 = if (available) &Blob("fatbin_ck_experts").bytes else &.{}; pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; +pub const act_rows: []const u8 = if (available) &Blob("fatbin_act_rows").bytes else &.{}; +pub const gdn_prep: []const u8 = if (available) &Blob("fatbin_gdn_prep").bytes else &.{}; +pub const qk_rows: []const u8 = if (available) &Blob("fatbin_qk_rows").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/cuda/qk_rows.zig b/zig/src/cuda/qk_rows.zig new file mode 100644 index 0000000000..d1ccb1a4ed --- /dev/null +++ b/zig/src/cuda/qk_rows.zig @@ -0,0 +1,107 @@ +//! qk_rows.cu's attention inputs for every CUDA family: QK norm, RoPE (rotate-half, partial) and the cache write. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// A packed row's heads: h query, hk key and hk value heads of `dim` (a multiple of 32, at most 256). +pub const Heads = struct { h: usize, hk: usize, dim: usize }; + +/// Where a row's heads come from: q, k and v (0: none) at row stride *_ld and head stride *_hs (elements). +pub const Src = extern struct { q: u64, q_ld: i64, q_hs: i64, k: u64, k_ld: i64, k_hs: i64, v: u64 = 0, v_ld: i64 = 0, v_hs: i64 = 0 }; + +/// q/k norms [2, dim] (or q's with norm_k), fp32 or bf16; RoPE on `rot` dims: cos/sin tables or inv [rot / 2]. +pub const Shape = struct { norms: ?u64 = null, norm_k: ?u64 = null, norms_bf16: bool = false, cos: ?u64 = null, sin: ?u64 = null, inv: ?u64 = null, rot: usize = 0, eps: f32 = 1e-6 }; + +/// qk_rows.cu's QkW. +const W = extern struct { q: u64, k: u64, bf16: i32 }; + +/// Where the rows go: q [n, h, dim], k and v [n, hk, dim]; k and v also into kc/vc at (slot * len + pos % len). +pub const Out = struct { + q: u64, + k: u64, + v: u64, + kc: ?u64 = null, + vc: ?u64 = null, + len: usize = 1, + /// Row and head strides in elements (0: packed [n, heads, dim]); head-major [heads, n, dim]: ld dim, hs n dim. + q_ld: usize = 0, + q_hs: usize = 0, + kv_ld: usize = 0, + kv_hs: usize = 0, +}; + +/// qk_rows.cu's QkDst. +const Dst = extern struct { q: u64, q_ld: i64, q_hs: i64, k: u64, k_ld: i64, k_hs: i64, v: u64, v_ld: i64, v_hs: i64 }; + +pub const QkRows = struct { + mod: Module, + fn_: Function, + /// tf_qk_rows_w: dim 128 with no or full rotation, a warp a head. + warp_fn: Function, + + pub fn load(d: *const Driver) !QkRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.qk_rows); + errdefer mod.unload(); + return .{ .mod = mod, .fn_ = try mod.function("tf_qk_rows"), .warp_fn = try mod.function("tf_qk_rows_w") }; + } + + pub fn unload(g: *QkRows) void { + g.mod.unload(); + } + + /// n packed rows qkv [n, (h + 2 hk) dim] bf16 at positions pos (int32 [n]) in cache slots slot (int32 [n] or none). + pub fn rows(g: *const QkRows, s: Stream, qkv: u64, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + const ld: i64 = @intCast((hs.h + 2 * hs.hk) * hs.dim); + const dim: i64 = @intCast(hs.dim); + const src: Src = .{ .q = qkv, .q_ld = ld, .q_hs = dim, .k = qkv + hs.h * hs.dim * 2, .k_ld = ld, .k_hs = dim, .v = qkv + (hs.h + hs.hk) * hs.dim * 2, .v_ld = ld, .v_hs = dim }; + return g.rowsFrom(s, src, hs, shape, pos, slot, out, n); + } + + /// As `rows`, the heads read from `src` (e.g. [q | gate] rows: q_hs = 2 dim); no v source: no v out or vc. + pub fn rowsFrom(g: *const QkRows, s: Stream, src: Src, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + if (hs.dim % 32 != 0 or hs.dim > 256 or shape.rot % 2 != 0 or shape.rot > hs.dim) return error.HeadShape; + if ((shape.cos == null) != (shape.sin == null) or (shape.cos != null and shape.inv != null)) return error.HeadShape; + const roped = shape.cos != null or shape.inv != null; + if (roped and shape.rot == 0) return error.HeadShape; + if (src.v == 0 and out.vc != null) return error.HeadShape; + if (n == 0) return; + var a: launch_.Args = .{}; + a.add(src); + const wsize: usize = if (shape.norms_bf16) 2 else 4; + const nq = shape.norms orelse 0; + a.add(W{ .q = nq, .k = shape.norm_k orelse (if (nq != 0) nq + hs.dim * wsize else 0), .bf16 = @intFromBool(shape.norms_bf16) }); + a.add(shape.cos orelse 0); + a.add(shape.sin orelse 0); + a.add(shape.inv orelse 0); + a.add(@as(i32, @intCast(if (roped) shape.rot else 0))); + a.add(pos); + a.add(slot orelse 0); + const dim: i64 = @intCast(hs.dim); + const qld: i64 = if (out.q_ld != 0) @intCast(out.q_ld) else @as(i64, @intCast(hs.h)) * dim; + const qhs: i64 = if (out.q_hs != 0) @intCast(out.q_hs) else dim; + const kld: i64 = if (out.kv_ld != 0) @intCast(out.kv_ld) else @as(i64, @intCast(hs.hk)) * dim; + const khs: i64 = if (out.kv_hs != 0) @intCast(out.kv_hs) else dim; + a.add(Dst{ .q = out.q, .q_ld = qld, .q_hs = qhs, .k = out.k, .k_ld = kld, .k_hs = khs, .v = out.v, .v_ld = kld, .v_hs = khs }); + a.add(out.kc orelse 0); + a.add(out.vc orelse 0); + a.add(@as(i64, @intCast(out.len))); + a.add(@as(i32, @intCast(hs.h))); + a.add(@as(i32, @intCast(hs.hk))); + a.add(shape.eps); + const heads = hs.h + (if (src.v != 0) 2 * hs.hk else hs.hk); + if (hs.dim == 128 and (!roped or shape.rot == 128)) { + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast((heads + 3) / 4), .z = 1 }; + return launch_.launch(g.warp_fn, .{ .grid = grid, .block = .{ .x = 128 } }, s, &a); + } + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast(heads), .z = 1 }; + try launch_.launch(g.fn_, .{ .grid = grid, .block = .{ .x = @intCast(hs.dim) } }, s, &a); + } +}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6c4e0991b6..710c2e76a7 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -34,6 +34,9 @@ pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); pub const vit = @import("vit.zig"); +pub const act_rows = @import("act_rows.zig"); +pub const gdn_prep = @import("gdn_prep.zig"); +pub const qk_rows = @import("qk_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const rank_scatter = @import("rank_scatter.zig"); @@ -49,6 +52,9 @@ test { _ = @import("norms.zig"); _ = @import("embed_rows.zig"); _ = @import("vit.zig"); + _ = @import("act_rows.zig"); + _ = @import("gdn_prep.zig"); + _ = @import("qk_rows.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = @import("rank_scatter.zig"); diff --git a/zig/tests/cuda/act_rows_tests.zig b/zig/tests/cuda/act_rows_tests.zig new file mode 100644 index 0000000000..43ddaf4b7e --- /dev/null +++ b/zig/tests/cuda/act_rows_tests.zig @@ -0,0 +1,132 @@ +//! act_rows.cu's SwiGLU, output gate, group sums and DFlash2's dynamic conv against float64, bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rows = @import("norm_rows_tests.zig"); + +const aref = core.act_ref; +const counts = rows.counts; +const most = rows.most; +const Rig = rows.Rig; +const bf16 = rows.bf16; +const bf = rows.bf; + +fn swiglu(r: *Rig, g: *const cuda.act_rows.ActRows, n: usize) !void { + const gu = r.bfs(most * 2 * n, 3); // [gate | up] rows + const d = try r.dev(std.mem.sliceAsBytes(gu)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * n * 2); + xss[c] = try r.out(m * n / 64 * 4); + try g.swiglu(r.s, d, 2 * n, d + 2 * n, 2 * n, os[c], xss[c], m, n); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "swiglu {d}", .{n}); + const o = try r.back(u16, os[counts.len - 1], most * n); + const xs = try r.back(f32, xss[counts.len - 1], most * n / 64); + for (0..most) |row| for (0..n) |i| { + const w = aref.swiglu(bf(gu[row * 2 * n + i]), bf(gu[row * 2 * n + n + i])); + try r.within(label, bf(o[row * n + i]), w, aref.tolerance(w)); + }; + try r.sums(label, o, xs, most, n); + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * n)), std.mem.sliceAsBytes(o[0 .. m * n])); +} + +fn outputGate(r: *Rig, g: *const cuda.act_rows.ActRows, heads: usize, d: usize) !void { + const n = heads * d; + const o = r.bfs(most * n, 1); + const qg = r.bfs(most * 2 * n, 3); // [q_h | gate_h] a head + const do_ = try r.dev(std.mem.sliceAsBytes(o)); + const dq = try r.dev(std.mem.sliceAsBytes(qg)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * n * 2); + xss[c] = try r.out(m * n / 64 * 4); + try g.outputGate(r.s, do_, dq + 2 * d, 2 * n, 2 * d, os[c], xss[c], m, heads, d); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "output gate {d} heads of {d}", .{ heads, d }); + const got = try r.back(u16, os[counts.len - 1], most * n); + const xs = try r.back(f32, xss[counts.len - 1], most * n / 64); + for (0..most) |row| for (0..heads) |h| for (0..d) |e| { + const w = bf(o[row * n + h * d + e]) * aref.sigmoid(bf(qg[row * 2 * n + h * 2 * d + d + e])); + try r.within(label, bf(got[row * n + h * d + e]), w, aref.tolerance(w)); + }; + try r.sums(label, got, xs, most, n); + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * n)), std.mem.sliceAsBytes(got[0 .. m * n])); +} + +fn groupSums(r: *Rig, g: *const cuda.act_rows.ActRows, k: usize) !void { + const ld = k + 64; + const x = r.bfs(most * ld, 2); + const d = try r.dev(std.mem.sliceAsBytes(x)); + const xsd = try r.out(most * k / 64 * 4); + try g.groupSums(r.s, d, ld, xsd, most, k); + const xs = try r.back(f32, xsd, most * k / 64); + const dense = try r.a.alloc(u16, most * k); + for (0..most) |row| @memcpy(dense[row * k ..][0..k], x[row * ld ..][0..k]); + try r.sums("group sums", dense, xs, most, k); +} + +fn dconv(r: *Rig, g: *const cuda.act_rows.ActRows, d: usize, gs: usize, seg: usize, branch: usize, residual: bool) !void { + const x = r.bfs(most * d, 1); + const dyn = r.bfs(most * 4 * (d / gs), 0.3); + const base = r.bfs(4 * d, 0.7); + const res = r.bfs(most * d, 1); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const ddyn = try r.dev(std.mem.sliceAsBytes(dyn)); + const dbase = try r.dev(std.mem.sliceAsBytes(base)); + const dres = try r.dev(std.mem.sliceAsBytes(res)); + var os: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * d * 2); + try g.dconv(r.s, dx, ddyn, dbase, if (residual) dres else null, os[c], m, d, gs, branch, seg); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "dconv d {d} gs {d} seg {d} branch {d}{s}", .{ d, gs, seg, branch, if (residual) " +res" else "" }); + const got = try r.back(u16, os[counts.len - 1], most * d); + const half = std.math.pow(f64, 2, -8); + for (0..most) |row| for (0..d) |i| { + const grp = i / gs; + const k0 = bf(bf16(@as(f32, @floatCast(bf(base[(branch * 2) * d + i]))) + @as(f32, @floatCast(bf(dyn[((row * 2 + branch) * 2) * (d / gs) + grp]))))); + const k1 = bf(bf16(@as(f32, @floatCast(bf(base[(branch * 2 + 1) * d + i]))) + @as(f32, @floatCast(bf(dyn[((row * 2 + branch) * 2 + 1) * (d / gs) + grp]))))); + const xv = bf(x[row * d + i]); + const pv = if (row % seg > 0) bf(x[(row - 1) * d + i]) else 0; + const y = aref.dconv(xv, pv, k0, k1); + const yt = (@abs(xv * k0) + @abs(pv * k1)) * std.math.pow(f64, 2, -22); + if (residual) { + const rv = bf(res[row * d + i]); + try r.within(label, bf(got[row * d + i]), rv + y, @abs(rv + y) * half + @abs(y) * half + yt + 1e-30); + } else try r.within(label, bf(got[row * d + i]), y, @abs(y) * half + yt + 1e-30); + }; + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * d)), std.mem.sliceAsBytes(got[0 .. m * d])); +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(27); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var g = try cuda.act_rows.ActRows.load(gpu.d); + defer g.unload(); + try swiglu(&r, &g, 17408); // Qwen3.8-27B's MLP + try swiglu(&r, &g, 192); + try outputGate(&r, &g, 24, 256); // the 27B's attention output gate + try outputGate(&r, &g, 3, 64); + try groupSums(&r, &g, 5120); + try dconv(&r, &g, 5120, 16, 16, 0, false); // DFlash2's block conv (groups of 16) + try dconv(&r, &g, 5120, 16, 16, 1, true); + try dconv(&r, &g, 256, 64, 7, 0, true); + check.pass("act rows: swiglu, output gate, dconv within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); +} diff --git a/zig/tests/cuda/gdn_prep_tests.zig b/zig/tests/cuda/gdn_prep_tests.zig new file mode 100644 index 0000000000..1ad46ec614 --- /dev/null +++ b/zig/tests/cuda/gdn_prep_tests.zig @@ -0,0 +1,156 @@ +//! gdn_prep.cu against core/gdn_prep_ref.zig: random conv windows over two streams, bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const ref = core.gdn_prep_ref; +const counts = [_]usize{ 1, 3, 17, 40 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +fn bfs(a: std.mem.Allocator, rng: std.Random, n: usize, scale: f32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*x| x.* = bf16(rng.floatNorm(f32) * scale); + return v; +} + +const Bufs = struct { + gpu: check.Gpu, + list: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn up(b: *Bufs, bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(b.gpu.d, bytes); + try b.list.append(b.gpu.gpa, d); + return d.ptr; + } + + fn out(b: *Bufs, s: cuda.Stream, n: usize) !u64 { + const d = try cuda.DeviceBuffer.alloc(b.gpu.d, @max(n, 16)); + try b.list.append(b.gpu.gpa, d); + try d.fill8(0xff, s.handle); + return d.ptr; + } + + fn back(b: *Bufs, a: std.mem.Allocator, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try a.alloc(T, n); + const d: cuda.DeviceBuffer = .{ .d = b.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try d.download(0, std.mem.sliceAsBytes(o)); + return o; + } +}; + +fn case(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.gdn_prep.GdnPrep, rng: std.Random, sh: cuda.gdn_prep.Shape, worst: *f64) !void { + var b: Bufs = .{ .gpu = gpu }; + defer { + for (b.list.items) |*d| d.free(); + b.list.deinit(gpu.gpa); + } + const c = (2 * sh.kh + sh.vh) * sh.dk; + const taps = sh.keep + 1; + const ld = c + 64; // stacked rows: something after the channels + const ab_ld = sh.vh + 3; + const qkv = try bfs(a, rng, most * ld, 1); + const state = try bfs(a, rng, 2 * sh.keep * c, 1); + const conv = try bfs(a, rng, c * taps, 0.4); + const av = try bfs(a, rng, most * ab_ld, 2); + const bv = try bfs(a, rng, most * ab_ld, 2); + const a_log = try a.alloc(f32, sh.vh); + const dtb = try a.alloc(f32, sh.vh); + for (a_log, dtb) |*x, *y| { + x.* = rng.float(f32) * 2 - 1; + y.* = rng.floatNorm(f32); + } + const sid = try a.alloc(i32, most); + const win = try a.alloc(i32, most * taps); + for (sid, 0..) |*x, i| x.* = @intCast(i % 2); + for (win, 0..) |*w, i| { // kept rows or any qkv row (a tree's window is any rows) + w.* = if (rng.uintLessThan(u32, 3) == 0) @intCast(rng.uintLessThan(usize, sh.keep)) else @intCast(sh.keep + rng.uintLessThan(usize, most)); + if (i % taps == taps - 1) w.* = @intCast(sh.keep + i / taps); // the row itself last + } + const base: cuda.gdn_prep.Io = .{ .qkv = try b.up(std.mem.sliceAsBytes(qkv)), .qkv_ld = @intCast(ld), .state = try b.up(std.mem.sliceAsBytes(state)), .conv = try b.up(std.mem.sliceAsBytes(conv)), .win = try b.up(std.mem.sliceAsBytes(win)), .sid = try b.up(std.mem.sliceAsBytes(sid)), .a = try b.up(std.mem.sliceAsBytes(av)), .a_ld = @intCast(ab_ld), .b = try b.up(std.mem.sliceAsBytes(bv)), .b_ld = @intCast(ab_ld), .a_log = try b.up(std.mem.sliceAsBytes(a_log)), .dt_bias = try b.up(std.mem.sliceAsBytes(dtb)), .q = 0, .k = 0, .v = 0, .g = 0, .beta = 0 }; + var ios: [counts.len]cuda.gdn_prep.Io = undefined; + for (counts, &ios) |m, *io| { + io.* = base; + io.q = try b.out(s, m * sh.kh * sh.dk * 2); + io.k = try b.out(s, m * sh.kh * sh.dk * 2); + io.v = try b.out(s, m * sh.vh * sh.dk * 2); + io.g = try b.out(s, m * sh.vh * 4); + io.beta = try b.out(s, m * sh.vh * 4); + try g.run(s, io.*, m, sh); + } + try s.synchronize(); + const io = ios[counts.len - 1]; + const q = try b.back(a, u16, io.q, most * sh.kh * sh.dk); + const k = try b.back(a, u16, io.k, most * sh.kh * sh.dk); + const v = try b.back(a, u16, io.v, most * sh.vh * sh.dk); + const gg = try b.back(a, f32, io.g, most * sh.vh); + const be = try b.back(a, f32, io.beta, most * sh.vh); + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "gdn prep kh {d} vh {d} dk {d} taps {d}", .{ sh.kh, sh.vh, sh.dk, taps }); + const rows = try a.alloc([]f64, taps); + for (rows) |*x| x.* = try a.alloc(f64, sh.dk); + const xs = try a.alloc([]const f64, taps); + const w = try a.alloc(f64, sh.dk * taps); + const cv = try a.alloc(f64, sh.dk); + const err = try a.alloc(f64, sh.dk); + const want = try a.alloc(f64, sh.dk); + var most_q: f64 = 0; + const rel = 3 * std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -18); // c rounded, the norm through it, the output + for (0..most) |r| for (0..2 * sh.kh + sh.vh) |hd| { + for (rows, xs, 0..) |x, *p, t| { + const src: usize = @intCast(win[r * taps + t]); + for (x, 0..) |*o, i| o.* = if (src < sh.keep) bf(state[(@as(usize, @intCast(sid[r])) * sh.keep + src) * c + hd * sh.dk + i]) else bf(qkv[(src - sh.keep) * ld + hd * sh.dk + i]); + p.* = x; + } + for (0..sh.dk) |i| for (0..taps) |t| { + w[i * taps + t] = bf(conv[(hd * sh.dk + i) * taps + t]); + }; + ref.conv(xs, w, taps, cv, err); + if (hd < 2 * sh.kh) { + const is_q = hd < sh.kh; + const f = ref.norm(cv, is_q, sh.eps, want); + const got = if (is_q) q[(r * sh.kh + hd) * sh.dk ..] else k[(r * sh.kh + hd - sh.kh) * sh.dk ..]; + for (want, err, 0..) |x, e, i| most_q = @max(most_q, @abs(bf(got[i]) - x) / (@abs(x) * rel + f * e + 1e-30)); + } else { + const hv = hd - 2 * sh.kh; + for (cv, err, 0..) |x, e, i| most_q = @max(most_q, @abs(bf(v[(r * sh.vh + hv) * sh.dk + i]) - x) / (@abs(x) * (std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -20)) + e + 1e-30)); + const x = bf(av[r * ab_ld + hv]) + dtb[hv]; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + const wg = ref.gate(bf(av[r * ab_ld + hv]), dtb[hv], a_log[hv]); + most_q = @max(most_q, @abs(gg[r * sh.vh + hv] - wg) / (@abs(wg) * std.math.pow(f64, 2, -20) * (4 + 4 * @exp(@as(f64, a_log[hv])) * sp) + 1e-30)); + const wb = 1 / (1 + @exp(-bf(bv[r * ab_ld + hv]))); + most_q = @max(most_q, @abs(be[r * sh.vh + hv] - wb) / (wb * std.math.pow(f64, 2, -19) + 1e-30)); + } + }; + worst.* = @max(worst.*, most_q); + try check.expect(most_q <= 1, "{s}: within float64 tolerance (worst {d:.3})", .{ label, most_q }); + for (counts[0 .. counts.len - 1], ios[0 .. counts.len - 1]) |m, o| { + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, u16, o.q, m * sh.kh * sh.dk)), std.mem.sliceAsBytes(q[0 .. m * sh.kh * sh.dk])); + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, u16, o.v, m * sh.vh * sh.dk)), std.mem.sliceAsBytes(v[0 .. m * sh.vh * sh.dk])); + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, f32, o.g, m * sh.vh)), std.mem.sliceAsBytes(gg[0 .. m * sh.vh])); + } +} + +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 cuda.gdn_prep.GdnPrep.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(48); + var worst: f64 = 0; + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 16, .vh = 48, .dk = 128, .keep = 3 }, &worst); // Qwen3.8-27B + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 2, .vh = 3, .dk = 64, .keep = 3 }, &worst); + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 1, .vh = 2, .dk = 256, .keep = 1 }, &worst); + check.pass("gdn prep: conv windows, q/k norms, v, g and beta within float64 tolerance (worst {d:.3}), bits at every row count", .{worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..15e70e0162 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,9 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +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 usage = \\usage: tf-cuda-test @@ -59,6 +62,9 @@ const usage = \\ triton a Triton cubin's bits against the Python oracle's fixture \\ norms RMS norm and the residual step against float64 references (synthetic) \\ norm-rows bf16-residual and gated norms, token rows against float64 (synthetic) + \\ act-rows SwiGLU, output gates, group sums, DFlash2's dynamic conv against float64 (synthetic) + \\ gdn-prep Gated DeltaNet conv, norms, g and beta against float64 (synthetic) + \\ qk-rows QK norm, RoPE and the cache write against float64 references (synthetic) \\ sample sample.cu's keyed draws against the Metal rule's host references (synthetic rows) \\ glue Nemotron's glue kernels against their host references at its shapes (synthetic rows) \\ lfm2-kernels LFM2's glue kernels against their host references at d1's shapes (synthetic rows) @@ -164,6 +170,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "act-rows")) return act_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "gdn-prep")) return gdn_prep_tests.run(gpu); + if (std.mem.eql(u8, cmd, "qk-rows")) return qk_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "lfm2-kernels")) return lfm2_kernel_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig index 8f1d295ec1..6610270ace 100644 --- a/zig/tests/cuda/norm_rows_tests.zig +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -6,20 +6,20 @@ const check = @import("check.zig"); const nref = core.norms_ref; const aref = core.act_ref; -const counts = [_]usize{ 1, 3, 17, 64 }; -const most = counts[counts.len - 1]; +pub const counts = [_]usize{ 1, 3, 17, 64 }; +pub const most = counts[counts.len - 1]; const eps: f32 = 1e-6; -fn bf16(x: f32) u16 { +pub fn bf16(x: f32) u16 { const u: u32 = @bitCast(x); return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); } -fn bf(h: u16) f64 { +pub fn bf(h: u16) f64 { return @as(f32, @bitCast(@as(u32, h) << 16)); } -const Rig = struct { +pub const Rig = struct { gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, @@ -27,20 +27,20 @@ const Rig = struct { bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, worst: f64 = 0, - fn dev(r: *Rig, bytes: []const u8) !u64 { + pub fn dev(r: *Rig, bytes: []const u8) !u64 { const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); try r.bufs.append(r.gpu.gpa, b); return b.ptr; } - fn out(r: *Rig, bytes: usize) !u64 { + pub fn out(r: *Rig, bytes: usize) !u64 { const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); try r.bufs.append(r.gpu.gpa, b); try b.fill8(0xff, r.s.handle); // a value the kernel misses reads NaN return b.ptr; } - fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + pub fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { const o = try r.a.alloc(T, n); try r.s.synchronize(); const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; @@ -49,20 +49,20 @@ const Rig = struct { } /// bf16 values: mostly within [-scale, scale), one in 64 twenty times larger. - fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + pub fn bfs(r: *Rig, n: usize, scale: f32) []u16 { const v = r.a.alloc(u16, n) catch unreachable; for (v, 0..) |*x, i| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); return v; } /// Weights near 1. - fn weights(r: *Rig, n: usize) []f32 { + pub fn weights(r: *Rig, n: usize) []f32 { const v = r.a.alloc(f32, n) catch unreachable; for (v) |*x| x.* = 0.5 + r.rng.float(f32); return v; } - fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + pub fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { const q = @abs(got - want) / tol; if (!(q <= 1)) { try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); @@ -71,7 +71,7 @@ const Rig = struct { } /// Group sums `xs` of the stored bf16 rows `y` ([rows, k]). - fn sums(r: *Rig, what: []const u8, y: []const u16, xs: []const f32, rows: usize, k: usize) !void { + pub fn sums(r: *Rig, what: []const u8, y: []const u16, xs: []const f32, rows: usize, k: usize) !void { var row: [64]f64 = undefined; for (0..rows) |i| for (0..k / 64) |j| { for (&row, 0..) |*v, e| v.* = bf(y[i * k + 64 * j + e]); diff --git a/zig/tests/cuda/qk_rows_tests.zig b/zig/tests/cuda/qk_rows_tests.zig new file mode 100644 index 0000000000..cdb5105ef3 --- /dev/null +++ b/zig/tests/cuda/qk_rows_tests.zig @@ -0,0 +1,226 @@ +//! qk_rows.cu against core/qk_ref.zig: q, k within float64 tolerance, v and the caches exact, bits at any row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const counts = [_]usize{ 1, 3, 17, 300 }; +const most = counts[counts.len - 1]; +const positions = 4096; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Case = struct { h: usize, hk: usize, dim: usize, rot: usize, norm: bool }; + +fn run1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.qk_rows.QkRows, rng: std.Random, c: Case, worst: *f64) !void { + const heads = c.h + 2 * c.hk; + const width = heads * c.dim; + const xs = try a.alloc(u16, most * width); + for (xs, 0..) |*h, i| h.* = bf16(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const nw = try a.alloc(f32, 2 * c.dim); + for (nw) |*w| w.* = 0.5 + rng.float(f32); + const tc = try a.alloc(f32, positions * @max(c.rot, 1)); + const ts = try a.alloc(f32, positions * @max(c.rot, 1)); + for (0..positions) |p| for (0..c.rot) |i| { + const j = i % (c.rot / 2); + const ang = @as(f64, @floatFromInt(p)) / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(c.rot))); + tc[p * c.rot + i] = @floatCast(@cos(ang)); + ts[p * c.rot + i] = @floatCast(@sin(ang)); + }; + const len = 600; + const pos = try a.alloc(i32, most); + const slot = try a.alloc(i32, most); + for (pos, slot, 0..) |*p, *sl, i| { + p.* = @intCast((i * 7919 + 13) % positions); + sl.* = @intCast(i % 2); + } + // distinct (slot, pos % len) a row, so no row's cache entry overwrites another's + for (0..most) |i| pos[i] = @intCast((i / 2) % len + len * (i % 5)); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*b| b.free(); + bufs.deinit(gpu.gpa); + } + const up = struct { + fn f(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(gp.d, bytes); + try l.append(gp.gpa, b); + return b.ptr; + } + fn z(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), st: cuda.Stream, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(gp.d, n); + try l.append(gp.gpa, b); + try b.fill8(0xff, st.handle); + return b.ptr; + } + }; + const dx = try up.f(gpu, &bufs, std.mem.sliceAsBytes(xs)); + const dnw = try up.f(gpu, &bufs, std.mem.sliceAsBytes(nw)); + const dc = try up.f(gpu, &bufs, std.mem.sliceAsBytes(tc)); + const ds = try up.f(gpu, &bufs, std.mem.sliceAsBytes(ts)); + const dp = try up.f(gpu, &bufs, std.mem.sliceAsBytes(pos)); + const dsl = try up.f(gpu, &bufs, std.mem.sliceAsBytes(slot)); + const shape: cuda.qk_rows.Shape = .{ .norms = if (c.norm) dnw else null, .cos = if (c.rot > 0) dc else null, .sin = if (c.rot > 0) ds else null, .rot = c.rot }; + const cache_bytes = 2 * len * c.hk * c.dim * 2; + const outs = try a.alloc(cuda.qk_rows.Out, counts.len); + for (counts, outs) |m, *o| { + o.* = .{ .q = try up.z(gpu, &bufs, s, m * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, m * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, m * c.hk * c.dim * 2), .kc = try up.z(gpu, &bufs, s, cache_bytes), .vc = try up.z(gpu, &bufs, s, cache_bytes), .len = len }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o.*, m); + } + try s.synchronize(); + const o = outs[counts.len - 1]; + const back = struct { + fn f(gp: check.Gpu, al: std.mem.Allocator, ptr: u64, n: usize) ![]u16 { + const raw = try al.alloc(u16, n); + const b: cuda.DeviceBuffer = .{ .d = gp.d, .ptr = ptr, .len = n * 2 }; + try b.download(0, std.mem.sliceAsBytes(raw)); + return raw; + } + }; + const gq = try back.f(gpu, a, o.q, most * c.h * c.dim); + const gk = try back.f(gpu, a, o.k, most * c.hk * c.dim); + const gv = try back.f(gpu, a, o.v, most * c.hk * c.dim); + const gkc = try back.f(gpu, a, o.kc.?, cache_bytes / 2); + const gvc = try back.f(gpu, a, o.vc.?, cache_bytes / 2); + const xv = try a.alloc(f64, c.dim); + const wq = try a.alloc(f64, c.dim); + const wk = try a.alloc(f64, c.dim); + for (wq, wk, 0..) |*p, *q, i| { + p.* = nw[i]; + q.* = nw[c.dim + i]; + } + const cv = try a.alloc(f64, @max(c.rot, 1)); + const sv = try a.alloc(f64, @max(c.rot, 1)); + const y = try a.alloc(f64, c.dim); + var q_most: f64 = 0; + for (0..most) |r| { + const p: usize = @intCast(pos[r]); + for (0..c.rot) |i| { + cv[i] = tc[p * c.rot + i]; + sv[i] = ts[p * c.rot + i]; + } + const at = @as(usize, @intCast(slot[r])) * len + p % len; + for (0..heads) |hd| { + for (xv, 0..) |*v, i| v.* = bf(xs[r * width + hd * c.dim + i]); + if (hd >= c.h + c.hk) { // v: exact copies, in the output and the cache + const k = hd - c.h - c.hk; + for (0..c.dim) |i| { + try check.expect(gv[(r * c.hk + k) * c.dim + i] == xs[r * width + hd * c.dim + i], "qk rows: v copied", .{}); + try check.expect(gvc[(at * c.hk + k) * c.dim + i] == xs[r * width + hd * c.dim + i], "qk rows: v cached", .{}); + } + continue; + } + const is_q = hd < c.h; + core.qk_ref.head(xv, if (c.norm) (if (is_q) wq else wk) else null, 1e-6, if (c.rot > 0) cv else null, if (c.rot > 0) sv else null, c.rot, y); + for (y, 0..) |want, i| { + const got = if (is_q) bf(gq[(r * c.h + hd) * c.dim + i]) else bf(gk[(r * c.hk + hd - c.h) * c.dim + i]); + const tol = (@abs(want) + 2 * @abs(xv[i]) + 1e-3) * std.math.pow(f64, 2, -8); + q_most = @max(q_most, @abs(got - want) / tol); + if (!is_q) try check.expect(gkc[(at * c.hk + hd - c.h) * c.dim + i] == gk[(r * c.hk + hd - c.h) * c.dim + i], "qk rows: k cached as output", .{}); + } + } + } + worst.* = @max(worst.*, q_most); + // the same heads from [q | gate] rows and a padded k buffer, no v: the same q, k and k cache bits + const qg = try a.alloc(u16, most * c.h * 2 * c.dim); + const kb = try a.alloc(u16, most * (c.hk * c.dim + 64)); + @memset(qg, 0x7fc0); // the gate halves: NaN, never read + @memset(kb, 0x7fc0); + for (0..most) |r| { + for (0..c.h) |hd| @memcpy(qg[(r * c.h + hd) * 2 * c.dim ..][0..c.dim], xs[r * width + hd * c.dim ..][0..c.dim]); + @memcpy(kb[r * (c.hk * c.dim + 64) ..][0 .. c.hk * c.dim], xs[r * width + c.h * c.dim ..][0 .. c.hk * c.dim]); + } + const src: cuda.qk_rows.Src = .{ .q = try up.f(gpu, &bufs, std.mem.sliceAsBytes(qg)), .q_ld = @intCast(c.h * 2 * c.dim), .q_hs = @intCast(2 * c.dim), .k = try up.f(gpu, &bufs, std.mem.sliceAsBytes(kb)), .k_ld = @intCast(c.hk * c.dim + 64), .k_hs = @intCast(c.dim) }; + const o2: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = 0, .kc = try up.z(gpu, &bufs, s, cache_bytes), .len = len }; + try g.rowsFrom(s, src, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o2, most); + try s.synchronize(); + try check.sameBytes("qk rows from strided heads: q", std.mem.sliceAsBytes(try back.f(gpu, a, o2.q, most * c.h * c.dim)), std.mem.sliceAsBytes(gq)); + try check.sameBytes("qk rows from strided heads: k", std.mem.sliceAsBytes(try back.f(gpu, a, o2.k, most * c.hk * c.dim)), std.mem.sliceAsBytes(gk)); + try check.sameBytes("qk rows from strided heads: k cache", std.mem.sliceAsBytes(try back.f(gpu, a, o2.kc.?, cache_bytes / 2)), std.mem.sliceAsBytes(gkc)); + // head-major outputs [heads, rows, dim] (strides): the packed call's bits, rearranged + { + const o4: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .len = len, .q_ld = c.dim, .q_hs = most * c.dim, .kv_ld = c.dim, .kv_hs = most * c.dim }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o4, most); + try s.synchronize(); + const hq = try back.f(gpu, a, o4.q, most * c.h * c.dim); + const hk_ = try back.f(gpu, a, o4.k, most * c.hk * c.dim); + const hv = try back.f(gpu, a, o4.v, most * c.hk * c.dim); + for (0..most) |r| { + for (0..c.h) |hd| try check.sameBytes("qk rows head-major: q", std.mem.sliceAsBytes(hq[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gq[(r * c.h + hd) * c.dim ..][0..c.dim])); + for (0..c.hk) |hd| { + try check.sameBytes("qk rows head-major: k", std.mem.sliceAsBytes(hk_[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gk[(r * c.hk + hd) * c.dim ..][0..c.dim])); + try check.sameBytes("qk rows head-major: v", std.mem.sliceAsBytes(hv[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gv[(r * c.hk + hd) * c.dim ..][0..c.dim])); + } + } + } + // bf16 norm weights at their own addresses and RoPE from inv (angle = pos * inv in fp32): float64 tolerance + if (c.norm) { + const wqb = try a.alloc(u16, c.dim); + const wkb = try a.alloc(u16, c.dim); + for (wqb, wkb, 0..) |*p, *q, i| { + p.* = bf16(nw[i]); + q.* = bf16(nw[c.dim + i]); + } + const half = @max(c.rot / 2, 1); + const inv = try a.alloc(f32, half); + for (inv, 0..) |*v, j| v.* = @floatCast(1 / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(@max(c.rot, 1))))); + const sh3: cuda.qk_rows.Shape = .{ .norms = try up.f(gpu, &bufs, std.mem.sliceAsBytes(wqb)), .norm_k = try up.f(gpu, &bufs, std.mem.sliceAsBytes(wkb)), .norms_bf16 = true, .inv = if (c.rot > 0) try up.f(gpu, &bufs, std.mem.sliceAsBytes(inv)) else null, .rot = c.rot }; + const o3: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .len = len }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, sh3, dp, dsl, o3, most); + try s.synchronize(); + const q3 = try back.f(gpu, a, o3.q, most * c.h * c.dim); + const k3 = try back.f(gpu, a, o3.k, most * c.hk * c.dim); + for (wq, wk, 0..) |*p, *q, i| { + p.* = bf(wqb[i]); + q.* = bf(wkb[i]); + } + for (0..most) |r| { + const p: usize = @intCast(pos[r]); + for (0..c.rot) |i| { + const ang: f64 = @as(f32, @floatFromInt(p)) * inv[i % half]; + cv[i] = @cos(ang); + sv[i] = @sin(ang); + } + for (0..c.h + c.hk) |hd| { + for (xv, 0..) |*v, i| v.* = bf(xs[r * width + hd * c.dim + i]); + const is_q = hd < c.h; + core.qk_ref.head(xv, if (is_q) wq else wk, 1e-6, if (c.rot > 0) cv else null, if (c.rot > 0) sv else null, c.rot, y); + for (y, 0..) |want, i| { + const got = if (is_q) bf(q3[(r * c.h + hd) * c.dim + i]) else bf(k3[(r * c.hk + hd - c.h) * c.dim + i]); + q_most = @max(q_most, @abs(got - want) / ((@abs(want) + 2 * @abs(xv[i]) + 1e-3) * std.math.pow(f64, 2, -8))); + } + } + } + worst.* = @max(worst.*, q_most); + } + try check.expect(q_most <= 1, "qk rows {d}/{d}x{d}, rot {d}{s}: within float64 tolerance (worst {d:.3})", .{ c.h, c.hk, c.dim, c.rot, if (c.norm) ", normed" else "", q_most }); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, oo| { + try check.sameBytes("qk rows: q at a smaller row count", std.mem.sliceAsBytes(try back.f(gpu, a, oo.q, m * c.h * c.dim)), std.mem.sliceAsBytes(gq[0 .. m * c.h * c.dim])); + try check.sameBytes("qk rows: k at a smaller row count", std.mem.sliceAsBytes(try back.f(gpu, a, oo.k, m * c.hk * c.dim)), std.mem.sliceAsBytes(gk[0 .. m * c.hk * c.dim])); + } +} + +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 cuda.qk_rows.QkRows.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(31); + var worst: f64 = 0; + const a = arena.allocator(); + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .dim = 128, .rot = 128, .norm = true }, &worst); // Kolibri, sliding + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .dim = 128, .rot = 0, .norm = true }, &worst); // Kolibri, full + try run1(gpu, a, s, &g, prng.random(), .{ .h = 8, .hk = 2, .dim = 64, .rot = 32, .norm = false }, &worst); // partial rotary + try run1(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 2, .dim = 128, .rot = 128, .norm = false }, &worst); // a warp a head, no norm + try run1(gpu, a, s, &g, prng.random(), .{ .h = 24, .hk = 4, .dim = 256, .rot = 64, .norm = true }, &worst); // Qwen3.8-27B + check.pass("qk rows: q and k within float64 tolerance (worst {d:.3}), v and the caches exact, bits at any row count", .{worst}); +} From bffc7047c8769d4f39ba5ae3ed55a99cf2f0349e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:27:32 +0200 Subject: [PATCH 08/41] cuda: qk_rows and gdn_prep take 1/sqrt by IEEE division and square root, correctly rounded on every architecture --- zig/kernels/cuda/gdn_prep.cu | 4 ++-- zig/kernels/cuda/qk_rows.cu | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index 04f330f4e7..c1af9cb37e 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -63,9 +63,9 @@ __device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int ke #pragma unroll for (int i = 0; i < J; ++i) part = __fmaf_rn(cv[i], cv[i], part); for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); - const float inv = rsqrtf(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps)); + const float inv = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps))); const bool is_q = hd < kh; - const float scale = is_q ? 1.f / DK : rsqrtf(static_cast(DK)); + const float scale = is_q ? 1.f / DK : __fdiv_rn(1.f, __fsqrt_rn(static_cast(DK))); __nv_bfloat16* out = (is_q ? io.q : io.k) + (r * kh + (is_q ? hd : hd - kh)) * DK + J * lane; #pragma unroll for (int i = 0; i < J; ++i) out[i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(cv[i], inv), scale)); diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu index 146e14017c..395b1ca677 100644 --- a/zig/kernels/cuda/qk_rows.cu +++ b/zig/kernels/cuda/qk_rows.cu @@ -96,7 +96,7 @@ extern "C" __global__ void __launch_bounds__(MOST) tf_qk_rows(const QkSrc in, co __syncthreads(); float s = 0.f; for (int w = 0; w < dim / 32; ++w) s = __fadd_rn(s, warps[w]); - const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps)); + const float rstd = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps))); y = __fmul_rn(__fmul_rn(x, rstd), wv(nw, hd >= h, d)); if (d < rot) yp = __fmul_rn(__fmul_rn(__bfloat162float(src[partner]), rstd), wv(nw, hd >= h, partner)); } else if (d < rot) { @@ -166,7 +166,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_qk_rows_w(const QkSrc in, c float s = 0.f; #pragma unroll for (int j = 0; j < J; ++j) s = __fadd_rn(s, part[j]); - const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps)); + const float rstd = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps))); #pragma unroll for (int j = 0; j < J; ++j) { y[j] = __fmul_rn(__fmul_rn(x[j], rstd), w[j]); From a97b9c0ff97f2cef42cd97165bc8c0cc1936d900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:55:00 +0200 Subject: [PATCH 09/41] cuda: decode-sized row calls take small blocks (launch.rowThreads, Triton's tuned sizes): swiglu 128 and gate 64 threads, gdn_prep a warp a block pair; the same bits at any row count --- zig/kernels/cuda/act_rows.cu | 8 ++++---- zig/src/cuda/act_rows.zig | 14 ++++++++------ zig/src/cuda/gdn_prep.zig | 3 ++- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index 515cc6bf67..9fc24a05b6 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/zig/kernels/cuda/act_rows.cu @@ -52,7 +52,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows(const __nv_bfloa __nv_bfloat16* __restrict__ out, float* __restrict__ xs, int n) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values const bool live = 8 * c < n; float g[8], u[8], o[8]; if (live) { @@ -75,7 +75,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat1 long long g_hs, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, int heads, int d) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values const int n = heads * d; const bool live = 8 * c < n; const int h = 8 * c / d, e = 8 * c % d; @@ -98,7 +98,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat1 extern "C" __global__ void __launch_bounds__(NT) tf_group_sums_rows(const __nv_bfloat16* __restrict__ x, long long ld, float* __restrict__ xs, int k) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; const bool live = 8 * c < k; const uint4 v = live ? *reinterpret_cast(x + r * ld + 8 * c) : make_uint4(0, 0, 0, 0); const float s = group_sum8(v, live); @@ -113,7 +113,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat __nv_bfloat16* __restrict__ out, int d, int gs, int branch, int seg) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; if (8 * c >= d) return; const int g = d / gs, grp = 8 * c / gs; float xv[8], pv[8], b0[8], b1[8], y[8]; diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig index 832844ce01..1d343f85f3 100644 --- a/zig/src/cuda/act_rows.zig +++ b/zig/src/cuda/act_rows.zig @@ -18,8 +18,8 @@ fn aligned(ptrs: []const u64, strides: []const usize) bool { return true; } -fn grid(rows: usize, n: usize) launch_.Dim3 { - return .{ .x = @intCast(rows), .y = @intCast((n / 8 + threads - 1) / threads), .z = 1 }; +fn grid(rows: usize, n: usize, t: u32) launch_.Dim3 { + return .{ .x = @intCast(rows), .y = @intCast((n / 8 + t - 1) / t), .z = 1 }; } pub const ActRows = struct { @@ -52,7 +52,8 @@ pub const ActRows = struct { a.add(out); a.add(xs orelse 0); a.add(@as(i32, @intCast(n))); - try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n), .block = .{ .x = threads } }, s, &a); + const t = launch_.rowThreads(rows, threads, 128); + try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n, t), .block = .{ .x = t } }, s, &a); } /// out = bf16(o * sigmoid(gate)), head h's gate at r * g_ld + h * g_hs ([q | gate]: q + d, 2 d); xs: sums. @@ -68,7 +69,8 @@ pub const ActRows = struct { a.add(xs orelse 0); a.add(@as(i32, @intCast(heads))); a.add(@as(i32, @intCast(d))); - try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d), .block = .{ .x = threads } }, s, &a); + const t = launch_.rowThreads(rows, threads, 64); + try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d, t), .block = .{ .x = t } }, s, &a); } /// xs [rows, k / 64]: x's rows (stride ld) summed over each 64 values. k % 64 == 0. @@ -80,7 +82,7 @@ pub const ActRows = struct { a.add(@as(i64, @intCast(ld))); a.add(xs); a.add(@as(i32, @intCast(k))); - try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k), .block = .{ .x = threads } }, s, &a); + try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k, threads), .block = .{ .x = threads } }, s, &a); } /// DFlash2's two-tap conv over [rows, d] bf16: dyn [rows, 2, 2, d / gs], base [2, 2, d]; res (or null) added. @@ -90,6 +92,6 @@ pub const ActRows = struct { var a: launch_.Args = .{}; for ([_]u64{ x, dyn, base, res orelse 0, out }) |v| a.add(v); for ([_]usize{ d, gs, branch, seg }) |v| a.add(@as(i32, @intCast(v))); - try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d), .block = .{ .x = threads } }, s, &a); + try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d, threads), .block = .{ .x = threads } }, s, &a); } }; diff --git a/zig/src/cuda/gdn_prep.zig b/zig/src/cuda/gdn_prep.zig index d11f869096..641a01cd8a 100644 --- a/zig/src/cuda/gdn_prep.zig +++ b/zig/src/cuda/gdn_prep.zig @@ -67,6 +67,7 @@ pub const GdnPrep = struct { a.add(@as(i32, @intCast(sh.vh))); a.add(@as(i32, @intCast(sh.keep))); a.add(sh.eps); - try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + 3) / 4), .z = 1 }, .block = .{ .x = 128 } }, s, &a); + const t = launch_.rowThreads(rows, 128, 64); // a warp a head either way + try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + t / 32 - 1) / (t / 32)), .z = 1 }, .block = .{ .x = t } }, s, &a); } }; From 7bf17c648eb8dba94f559184f738bf56201af480 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:07:48 +0200 Subject: [PATCH 10/41] cuda: act_rows and gdn_prep say why their sigmoid divides exactly --- zig/kernels/cuda/act_rows.cu | 1 + zig/kernels/cuda/gdn_prep.cu | 1 + 2 files changed, 2 insertions(+) diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index 9fc24a05b6..a10b0ed2c9 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/zig/kernels/cuda/act_rows.cu @@ -27,6 +27,7 @@ __device__ __forceinline__ uint4 pack8(const float* f) { return u; } +// Exactly rounded division on purpose: the same bits on every GPU, where div.full and ex2.approx are not. __device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } // A thread's 8 stored values in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index c1af9cb37e..05b3f140f6 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -30,6 +30,7 @@ struct Io { float* beta; // [rows, vh] }; +// Exactly rounded division on purpose: the same bits on every GPU, where div.full and ex2.approx are not. __device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } template From 011c851c9a7671f5bcf14a92b8bc61fa0369b928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:47:35 +0200 Subject: [PATCH 11/41] cuda/gdn_prep: every conv tap's loads issued before any is used, the same bits --- zig/kernels/cuda/gdn_prep.cu | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index 05b3f140f6..7915bea2f7 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -45,16 +45,22 @@ __device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int ke int src[MAXT]; #pragma unroll for (int t = 0; t < MAXT; ++t) src[t] = t <= keep ? io.win[r * (keep + 1) + t] : 0; + __nv_bfloat16 xv[MAXT][J], kv[MAXT][J]; +#pragma unroll + for (int t = 0; t < MAXT; ++t) { // every tap's loads issued before any is used + if (t > keep) break; + const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; +#pragma unroll + for (int i = 0; i < J; ++i) xv[t][i] = x[ch + i], kv[t][i] = io.conv[(ch + i) * (keep + 1) + t]; + } float acc[J]; #pragma unroll for (int i = 0; i < J; ++i) acc[i] = 0.f; #pragma unroll for (int t = 0; t < MAXT; ++t) { if (t > keep) break; - const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; #pragma unroll - for (int i = 0; i < J; ++i) - acc[i] = __fmaf_rn(__bfloat162float(x[ch + i]), __bfloat162float(io.conv[(ch + i) * (keep + 1) + t]), acc[i]); + for (int i = 0; i < J; ++i) acc[i] = __fmaf_rn(__bfloat162float(xv[t][i]), __bfloat162float(kv[t][i]), acc[i]); } float cv[J]; #pragma unroll From df0ef6bcae57d79cb454045f92f4d868b22750cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:15:46 +0200 Subject: [PATCH 12/41] cuda: sliding-window attention, sigmoid top-k MoE routing and combine, and dense rows for every family, checked against float64 --- zig/build/cuda.zig | 3 + zig/kernels/cuda/dense_rows.cu | 70 ++++ zig/kernels/cuda/moe_glue.cu | 184 +++++++++ zig/kernels/cuda/window_attention.cu | 570 +++++++++++++++++++++++++++ zig/src/core/dense_ref.zig | 24 ++ zig/src/core/moe_ref.zig | 55 +++ zig/src/core/root.zig | 4 + zig/src/cuda/dense_rows.zig | 51 +++ zig/src/cuda/kernels.zig | 3 + zig/src/cuda/moe_glue.zig | 76 ++++ zig/src/cuda/root.zig | 6 + zig/src/cuda/window_attention.zig | 115 ++++++ zig/tests/cuda/dense_rows_tests.zig | 79 ++++ zig/tests/cuda/main.zig | 9 + zig/tests/cuda/moe_glue_tests.zig | 161 ++++++++ zig/tests/cuda/window_tests.zig | 217 ++++++++++ 16 files changed, 1627 insertions(+) create mode 100644 zig/kernels/cuda/dense_rows.cu create mode 100644 zig/kernels/cuda/moe_glue.cu create mode 100644 zig/kernels/cuda/window_attention.cu create mode 100644 zig/src/core/dense_ref.zig create mode 100644 zig/src/core/moe_ref.zig create mode 100644 zig/src/cuda/dense_rows.zig create mode 100644 zig/src/cuda/moe_glue.zig create mode 100644 zig/src/cuda/window_attention.zig create mode 100644 zig/tests/cuda/dense_rows_tests.zig create mode 100644 zig/tests/cuda/moe_glue_tests.zig create mode 100644 zig/tests/cuda/window_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4ca82f73c0..f3d50c01de 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -49,6 +49,9 @@ const kernels = [_]Kernel{ .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .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 = "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 .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op diff --git a/zig/kernels/cuda/dense_rows.cu b/zig/kernels/cuda/dense_rows.cu new file mode 100644 index 0000000000..acfc61efe3 --- /dev/null +++ b/zig/kernels/cuda/dense_rows.cu @@ -0,0 +1,70 @@ +// bf16 rows times bf16 weight rows into fp32 for every CUDA family: routers and output heads (core/dense_ref.zig). + +// A warp an output column, RT rows a launch tile; each row's sum runs in one order whatever RT and the row count. + +#include +#include + +namespace tf_dense { + +constexpr int WARPS = 8; +constexpr int VEC = 8; // bf16 a 16-byte load + +__device__ __forceinline__ void unpack(const uint4& u, float* f) { + const __nv_bfloat162* h = reinterpret_cast(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 t = __bfloat1622float2(h[i]); + f[2 * i] = t.x; + f[2 * i + 1] = t.y; + } +} + +// out[r, n] = x[r] . w[n] for a tile's rows: lane l sums k = 8 l + 256 j in order, then the fixed shuffle tree. +template +__device__ __forceinline__ void tile(const __nv_bfloat16* x, int64_t x_stride, const __nv_bfloat16* w, float* out, + int64_t out_stride, int rows, int n, int k) { + const int lane = threadIdx.x & 31; + const int col = blockIdx.x * WARPS + (threadIdx.x >> 5); + const int r0 = blockIdx.y * RT; + if (col >= n) return; + float acc[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) acc[r] = 0.f; + const __nv_bfloat16* wr = w + static_cast(col) * k; +#pragma unroll 4 // several 16-byte weight loads in flight a lane (memory-level parallelism) + for (int k0 = lane * VEC; k0 < k; k0 += 32 * VEC) { + float wf[VEC]; + unpack(*reinterpret_cast(wr + k0), wf); +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r0 + r >= rows) break; + float xf[VEC]; + unpack(*reinterpret_cast(x + (r0 + r) * x_stride + k0), xf); +#pragma unroll + for (int i = 0; i < VEC; ++i) acc[r] = __fmaf_rn(xf[i], wf[i], acc[r]); + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) { + float v = acc[r]; + for (int m = 16; m > 0; m >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xFFFFFFFFu, v, m)); + if (lane == 0 && r0 + r < rows) out[(r0 + r) * out_stride + col] = v; + } +} + +} // namespace tf_dense + +using namespace tf_dense; + +#define TF_DENSE_ROWS(RT) \ + extern "C" __global__ void __launch_bounds__(WARPS * 32) tf_dense_rows_##RT( \ + const __nv_bfloat16* x, long long x_stride, const __nv_bfloat16* w, float* out, long long out_stride, \ + int rows, int n, int k) { \ + tile(x, x_stride, w, out, out_stride, rows, n, k); \ + } + +TF_DENSE_ROWS(1) +TF_DENSE_ROWS(4) +TF_DENSE_ROWS(16) +TF_DENSE_ROWS(32) diff --git a/zig/kernels/cuda/moe_glue.cu b/zig/kernels/cuda/moe_glue.cu new file mode 100644 index 0000000000..f4beca0f82 --- /dev/null +++ b/zig/kernels/cuda/moe_glue.cu @@ -0,0 +1,184 @@ +// MoE glue for every CUDA family (float64 references: core/moe_ref.zig): sigmoid top-k and the weighted combine. + +// Every row's bits depend on that row alone; kinds 0 fp32, 1 fp16, 2 bf16 (as norms.cu). + +#include +#include +#include +#include + +namespace tf_moe { + +constexpr int NT = 256; +constexpr int MOST = 4096; // experts a row ranks + +__device__ __forceinline__ float load(const void* p, int kind, int64_t i) { + if (kind == 0) return static_cast(p)[i]; + if (kind == 1) return __half2float(static_cast(p)[i]); + return __bfloat162float(static_cast(p)[i]); +} + +__device__ __forceinline__ void store(void* p, int kind, int64_t i, float v) { + if (kind == 0) static_cast(p)[i] = v; + else if (kind == 1) static_cast<__half*>(p)[i] = __float2half_rn(v); + else static_cast<__nv_bfloat16*>(p)[i] = __float2bfloat16_rn(v); +} + +// (value, id) wins over (other value, other id): the larger value, the lower id among equals. +__device__ __forceinline__ bool beats(float v, int i, float w, int j) { + return v > w || (v == w && i < j); +} + +} // namespace tf_moe + +using namespace tf_moe; + +// Top `topk` by logit + bias (ties: lower id), weights sigmoid(logit) [normalized] * scale; `shared`: slot topk = ne. +extern "C" __global__ void __launch_bounds__(NT) tf_moe_topk_sigmoid(const float* logits, const float* bias, int* pick, + float* wts, int ne, int topk, int shared, + int normalize, float scale) { + __shared__ float v[MOST]; + __shared__ float best_v[NT / 32]; + __shared__ int best_i[NT / 32]; + __shared__ int picked[64]; + const int64_t r = blockIdx.x; + const float* lg = logits + r * ne; + for (int e = threadIdx.x; e < ne; e += NT) v[e] = __fadd_rn(lg[e], bias != nullptr ? bias[e] : 0.f); + __syncthreads(); + for (int k = 0; k < topk; ++k) { + float bv = -INFINITY; + int bi = ne; + for (int e = threadIdx.x; e < ne; e += NT) + if (beats(v[e], e, bv, bi)) bv = v[e], bi = e; + for (int m = 16; m > 0; m >>= 1) { + const float ov = __shfl_xor_sync(0xFFFFFFFFu, bv, m); + const int oi = __shfl_xor_sync(0xFFFFFFFFu, bi, m); + if (beats(ov, oi, bv, bi)) bv = ov, bi = oi; + } + if ((threadIdx.x & 31) == 0) best_v[threadIdx.x >> 5] = bv, best_i[threadIdx.x >> 5] = bi; + __syncthreads(); + if (threadIdx.x == 0) { + float tv = best_v[0]; + int ti = best_i[0]; + for (int w = 1; w < NT / 32; ++w) + if (beats(best_v[w], best_i[w], tv, ti)) tv = best_v[w], ti = best_i[w]; + picked[k] = ti; + v[ti] = -INFINITY; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + const int slots = topk + (shared ? 1 : 0); + float sum = 0.f; + float w[64]; + for (int k = 0; k < topk; ++k) { + w[k] = __fdiv_rn(1.f, __fadd_rn(1.f, expf(-lg[picked[k]]))); + sum = __fadd_rn(sum, w[k]); + } + for (int k = 0; k < topk; ++k) { + pick[r * slots + k] = picked[k]; + wts[r * slots + k] = __fmul_rn(normalize ? __fdiv_rn(w[k], sum) : w[k], scale); + } + if (shared) pick[r * slots + topk] = ne, wts[r * slots + topk] = 1.f; + } +} + +// A float's order as an unsigned (-0 as +0, so equal values stay equal). +__device__ __forceinline__ unsigned ordered(float v) { + const unsigned u = __float_as_uint(__fadd_rn(v, 0.f)); + return u & 0x80000000u ? ~u : u | 0x80000000u; +} + +// tf_moe_topk_sigmoid with a warp a row (up to 32 * PER experts): values in registers, no block barrier. +constexpr int PER = 16; + +extern "C" __global__ void __launch_bounds__(256) tf_moe_topk_sigmoid_warp(const float* logits, const float* bias, + int* pick, float* wts, int rows, int ne, + int topk, int shared, int normalize, + float scale) { + const int lane = threadIdx.x & 31; + const int64_t r = static_cast(blockIdx.x) * 8 + (threadIdx.x >> 5); + if (r >= rows) return; + const float* lg = logits + r * ne; + float v[PER], x[PER]; +#pragma unroll + for (int j = 0; j < PER; ++j) { + const int e = lane + 32 * j; + x[j] = e < ne ? lg[e] : 0.f; + v[j] = e < ne ? __fadd_rn(x[j], bias != nullptr ? bias[e] : 0.f) : -INFINITY; + } + const int slots = topk + (shared ? 1 : 0); + float xk = 0.f; // lane k's pick's logit + int pk = 0; // lane k's pick + for (int k = 0; k < topk; ++k) { + float bv = -INFINITY; + int bi = ne; +#pragma unroll + for (int j = 0; j < PER; ++j) { + const int e = lane + 32 * j; + if (e < ne && beats(v[j], e, bv, bi)) bv = v[j], bi = e; + } + { // the warp's best: the largest value (redux on an order-keeping key), then the lowest id among equals + const unsigned key = ordered(bv); + const unsigned top = __reduce_max_sync(0xFFFFFFFFu, key); + bi = static_cast(__reduce_min_sync(0xFFFFFFFFu, key == top ? static_cast(bi) : 0xFFFFFFFFu)); + } + float xb = 0.f; +#pragma unroll + for (int j = 0; j < PER; ++j) + if (lane + 32 * j == bi) v[j] = -INFINITY, xb = x[j]; + xb = __shfl_sync(0xFFFFFFFFu, xb, bi & 31); + if (lane == k) xk = xb, pk = bi; + } + // the weights after the picks, each in its lane; the sum in pick order + const float wk = __fdiv_rn(1.f, __fadd_rn(1.f, expf(-xk))); + float sum = 0.f; + for (int k = 0; k < topk; ++k) sum = __fadd_rn(sum, __shfl_sync(0xFFFFFFFFu, wk, k)); + if (lane < topk) { + pick[r * slots + lane] = pk; + wts[r * slots + lane] = __fmul_rn(normalize ? __fdiv_rn(wk, sum) : wk, scale); + } + if (shared && lane == 0) pick[r * slots + topk] = ne, wts[r * slots + topk] = 1.f; +} + +// out[r, d] = sum over the row's slots, in slot order, of w[r, k] * y[r, k, d] (fp32 fma), rounded once to out's kind. +template +__device__ __forceinline__ void combine_row(const void* y, const float* w, void* out, int slots, int dims, int64_t r, int d) { + float acc = 0.f; + int k = 0; + for (; k + 8 <= slots; k += 8) { // eight slots' loads in flight, then folded in order + float v[8], wk[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) v[i] = load(y, YK, (r * slots + k + i) * dims + d), wk[i] = w[r * slots + k + i]; +#pragma unroll + for (int i = 0; i < 8; ++i) acc = __fmaf_rn(wk[i], v[i], acc); + } + if (k < slots) { + float v[8], wk[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + v[i] = k + i < slots ? load(y, YK, (r * slots + k + i) * dims + d) : 0.f; + wk[i] = k + i < slots ? w[r * slots + k + i] : 0.f; + } +#pragma unroll + for (int i = 0; i < 8; ++i) + if (k + i < slots) acc = __fmaf_rn(wk[i], v[i], acc); + } + store(out, OK, r * dims + d, acc); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_moe_combine(const void* y, int y_kind, const float* w, void* out, + int out_kind, int slots, int dims) { + const int64_t r = blockIdx.x; + const int d = blockIdx.y * NT + threadIdx.x; + if (d >= dims) return; + // the kinds picked once a thread (a branch an element costs instruction fetch on small calls) + if (y_kind == 0 && out_kind == 2) combine_row<0, 2>(y, w, out, slots, dims, r, d); + else if (y_kind == 2 && out_kind == 2) combine_row<2, 2>(y, w, out, slots, dims, r, d); + else if (y_kind == 0 && out_kind == 0) combine_row<0, 0>(y, w, out, slots, dims, r, d); + else { + float acc = 0.f; + for (int k = 0; k < slots; ++k) acc = __fmaf_rn(w[r * slots + k], load(y, y_kind, (r * slots + k) * dims + d), acc); + store(out, out_kind, r * dims + d, acc); + } +} diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu new file mode 100644 index 0000000000..542e80a8ba --- /dev/null +++ b/zig/kernels/cuda/window_attention.cu @@ -0,0 +1,570 @@ +// Sliding-window attention for every CUDA family (float64 references: core/attention_ref.zig): decode and prompt rows. + +// Decode: a block a (row, key head); keys walk 64-key tiles by absolute position, so a row's bits are its own. + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_window { + +constexpr int NT = 256; +constexpr int BN = 64; // keys a tile +constexpr int GMAX = 16; // query heads a key head +constexpr int KT = BN + 2; // kt's row stride: a bank further each dim +constexpr int DMAX = 128; // head dim (static shared memory: q, a K and a V tile, scores within 48 KiB) + +} // namespace tf_window + +using namespace tf_window; + +// Decode split by key tile (block: row, key head, tile); tf_window_merge folds the partials in tile order. +extern "C" __global__ void __launch_bounds__(NT) tf_window_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 d, int win, long long ring, float scale) { + __shared__ float qs[GMAX * DMAX]; + __shared__ __nv_bfloat16 kt[DMAX * KT]; // K transposed, [d][key], padded: a warp reads 32 keys of one dim + __shared__ __nv_bfloat16 vs[BN * DMAX]; + __shared__ float sc[GMAX * BN]; + __shared__ float mx[GMAX], ls[GMAX]; + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, tiles = gridDim.z, j = blockIdx.z; + const int g = h / hk; + 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 + 1, 0) : 0; + const int t = lo / BN + j; + const int64_t part = (r * hk + kh) * tiles + j; // this block's partial, g heads of d values + if (t > p / BN) { // past the row's last tile: an empty partial + for (int i = threadIdx.x; i < g; i += NT) pm[part * g + i] = -INFINITY, pl[part * g + i] = 0.f; + return; + } + for (int i = threadIdx.x; i < g * d; i += NT) qs[i] = __bfloat162float(q[(r * h + kh * g + i / d) * d + i % d]); + for (int i = threadIdx.x; i < BN * d; i += NT) { + const int key = t * BN + i / d; + const bool in = key >= lo && key <= p; + const int64_t at = ((base + (in ? key : lo) % ring) * hk + kh) * d + i % d; + kt[(i % d) * KT + i / d] = kc[at]; + vs[i] = vc[at]; + } + __syncthreads(); + for (int i = threadIdx.x; i < g * BN; i += NT) { // score (head i / BN, key i % BN) + const int hd = i / BN, jj = i % BN, key = t * BN + jj; + float s = 0.f; +#pragma unroll 8 + for (int e = 0; e < d; ++e) s = __fmaf_rn(qs[hd * d + e], __bfloat162float(kt[e * KT + jj]), s); + sc[i] = (key >= lo && key <= p) ? __fmul_rn(s, scale) : -INFINITY; + } + __syncthreads(); + { // a warp a head: the tile's max, the weights exp(s - max) and their sum, lanes in a fixed tree + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + for (int hd = warp; hd < g; hd += NT / 32) { + float a = sc[hd * BN + lane], b = sc[hd * BN + 32 + lane]; + float m = fmaxf(a, b); + for (int k = 16; k > 0; k >>= 1) m = fmaxf(m, __shfl_xor_sync(0xFFFFFFFFu, m, k)); + a = a == -INFINITY ? 0.f : expf(__fsub_rn(a, m)); + b = b == -INFINITY ? 0.f : expf(__fsub_rn(b, m)); + sc[hd * BN + lane] = a; + sc[hd * BN + 32 + lane] = b; + float l = __fadd_rn(a, b); + for (int k = 16; k > 0; k >>= 1) l = __fadd_rn(l, __shfl_xor_sync(0xFFFFFFFFu, l, k)); + if (lane == 0) mx[hd] = m, ls[hd] = l; + } + } + __syncthreads(); + for (int o = threadIdx.x; o < g * d; o += NT) { + const int hd = o / d, e = o % d; + float a = 0.f; + for (int jj = 0; jj < BN; ++jj) a = __fmaf_rn(sc[hd * BN + jj], __bfloat162float(vs[jj * d + e]), a); + po[part * g * d + o] = a; + } + if (threadIdx.x < g) pm[part * g + threadIdx.x] = mx[threadIdx.x], pl[part * g + threadIdx.x] = ls[threadIdx.x]; +} + +// out [W, H, D] bf16: a row's partials folded in tile order, o = sum_j exp(m_j - M) o_j / sum_j exp(m_j - M) l_j. +extern "C" __global__ void __launch_bounds__(NT) tf_window_merge(const float* __restrict__ po, + const float* __restrict__ pm, + const float* __restrict__ pl, + __nv_bfloat16* __restrict__ out, int h, int hk, int d, + int tiles) { + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, g = h / hk; + const int64_t part0 = (r * hk + kh) * tiles; + for (int o = threadIdx.x; o < g * d; o += NT) { + const int hd = o / d; + float m = -INFINITY; + for (int j = 0; j < tiles; ++j) m = fmaxf(m, pm[(part0 + j) * g + hd]); + float acc = 0.f, l = 0.f; + for (int j = 0; j < tiles; ++j) { + const float mj = pm[(part0 + j) * g + hd]; + if (mj == -INFINITY) continue; + const float a = expf(__fsub_rn(mj, m)); + acc = __fmaf_rn(a, po[(part0 + j) * g * d + o], acc); + l = __fmaf_rn(a, pl[(part0 + j) * g + hd], l); + } + out[(r * h + kh * g + hd) * d + o % d] = __float2bfloat16_rn(__fdiv_rn(acc, l)); + } +} + +// Prompt rows: prefill_attention.cu's tensor-core kernel with a window (keys (pos - win, pos]) and a ring of keys. +namespace tf_window_prompt { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +constexpr int BN = 64, HALF = 32; // keys a tile (the fold's unit), keys a staging slot +constexpr float LOG2E = 1.4426950408889634f; + +template +__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__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))); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ float tdiv(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +template +__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } + +__device__ __forceinline__ float rowsum(const float (&c)[8]) { + float w[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); + w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); + } + return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), + __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); +} + +template +__global__ void __launch_bounds__(32 * WARPS, 1) +wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, + int G, float scale, int win, int ring) { + constexpr int C = D / 8, K16 = D / 16, NB = D / 8, RB = WARPS / HPC, THREADS = 32 * WARPS; + constexpr bool QREG = D <= 128; // the queries' fragments stay in registers, else in shared memory + extern __shared__ uint4 smem[]; + uint4* slots = smem; + uint4* qs = smem + NS * HALF * C; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g = lane / 4, qd = lane % 4, mt = lane / 8; + const int groups = G / HPC, hk = blockIdx.y / groups, hg = blockIdx.y % groups; + const int row0 = blockIdx.x * 16 * RB, last = min(row0 + 16 * RB, W) - 1; + const int head = hk * G + hg * HPC + warp % HPC, wrow = row0 + (warp / HPC) * 16; + // the first tile any row of the block sees: its first row's window start (win 0: key 0) + const int t0 = win > 0 ? max(p0 + row0 - win + 1, 0) / BN : 0; + const int tiles = (p0 + last) / BN + 1 - t0, items = 4 * tiles, limit = p0 + W; + auto stage = [&](int item) { + const __nv_bfloat16* src = item % 4 < 2 ? kc : vc; + const int key0 = (t0 + item / 4) * BN + (item % 2) * HALF; + uint4* dst = slots + (item % NS) * HALF * C; + for (int i = threadIdx.x; i < HALF * C; i += THREADS) { + const int key = key0 + i / C; + const bool in = key < limit; // past the cache's filled keys: zeros, never read + const int row = in ? (ring > 0 ? key % ring : key) : 0; // a ring keeps key p at row p % ring + cp16z(dst + sw(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in); + } + }; + uint32_t qa[QREG ? K16 : 1][4]; + if constexpr (QREG) { + const int r0 = wrow + g, r1 = r0 + 8; + const uint32_t* q0 = reinterpret_cast(q + ((int64_t)min(r0, W - 1) * H + head) * D); + const uint32_t* q1 = reinterpret_cast(q + ((int64_t)min(r1, W - 1) * H + head) * D); +#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; + } + } else { + for (int i = threadIdx.x; i < WARPS * 16 * C; i += THREADS) { + const int w = i / (16 * C), r = (i / C) % 16, row = row0 + (w / HPC) * 16 + r; + const int hd = hk * G + hg * HPC + w % HPC; + cp16z(qs + sw(i / C, i % C), q + ((int64_t)min(row, W - 1) * H + hd) * D + (i % C) * 8, row < W); + } + commit(); + } +#pragma unroll + for (int i = 0; i < NS; ++i) { + if (i < items) stage(i); + 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 sc[8][4]; + uint32_t pa[4][4]; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + const int pos0 = p0 + wrow + g, pos1 = pos0 + 8; + auto take = [&](int it) -> const uint4* { // item it's slot, once every thread's copy has landed + wait_group(); + __syncthreads(); + return slots + (it % NS) * HALF * C; + }; + auto done = [&](int it) { // the slot is free: stage the item NS ahead into it + __syncthreads(); + if (it + NS < items) stage(it + NS); + commit(); + }; + auto scores = [&](float (&s)[4][4], const uint4* sl) { +#pragma unroll + for (int k = 0; k < K16; ++k) { + uint32_t a[4]; + if constexpr (QREG) { + a[0] = qa[k][0]; a[1] = qa[k][1]; a[2] = qa[k][2]; a[3] = qa[k][3]; + } else { + ldmatrix4(a, qs + sw(warp * 16 + lane % 16, 2 * k + lane / 16)); + } +#pragma unroll + for (int n = 0; n < 4; n += 2) { + uint32_t b[4]; + ldmatrix4(b, sl + sw(8 * (n + mt / 2) + lane % 8, 2 * k + mt % 2)); + if (k == 0) { + mma0(s[n], a, b[0], b[1]); + mma0(s[n + 1], a, b[2], b[3]); + } else { + mma(s[n], a, b[0], b[1]); + mma(s[n + 1], a, b[2], b[3]); + } + } + } + }; + auto values = [&](const uint32_t (&pa0)[4], const uint32_t (&pa1)[4], const uint4* sl) { +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, sl + sw(8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa0, b[0], b[1]); + mma(o[i + 1], pa0, b[2], b[3]); + } +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, sl + sw(16 + 8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa1, b[0], b[1]); + mma(o[i + 1], pa1, b[2], b[3]); + } + }; + for (int t = 0; t < tiles; ++t) { + const int it = 4 * t; + scores(*reinterpret_cast(&sc[0]), take(it)); + done(it); + scores(*reinterpret_cast(&sc[4]), take(it + 1)); + done(it + 1); + { // _tile's fold of the tile's 64 scores + const int kb = (t0 + t) * BN + 2 * qd; + const int lo0 = win > 0 ? pos0 - win + 1 : 0, lo1 = win > 0 ? pos1 - win + 1 : 0; + float tm0 = -INFINITY, tm1 = -INFINITY; +#pragma unroll + for (int n = 0; n < 8; ++n) { + const int c = kb + 8 * n; + sc[n][0] = c >= lo0 && c <= pos0 ? __fmul_rn(sc[n][0], scale) : -INFINITY; + sc[n][1] = c + 1 >= lo0 && c + 1 <= pos0 ? __fmul_rn(sc[n][1], scale) : -INFINITY; + sc[n][2] = c >= lo1 && c <= pos1 ? __fmul_rn(sc[n][2], scale) : -INFINITY; + sc[n][3] = c + 1 >= lo1 && c + 1 <= pos1 ? __fmul_rn(sc[n][3], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1])); + tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, 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 : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + float c0[8], c1[8]; +#pragma unroll + for (int n = 0; n < 8; ++n) { + const int c = kb + 8 * n; + sc[n][0] = act0 && c >= lo0 && c <= pos0 ? texp(__fsub_rn(sc[n][0], n0)) : 0.0f; + sc[n][1] = act0 && c + 1 >= lo0 && c + 1 <= pos0 ? texp(__fsub_rn(sc[n][1], n0)) : 0.0f; + sc[n][2] = act1 && c >= lo1 && c <= pos1 ? texp(__fsub_rn(sc[n][2], n1)) : 0.0f; + sc[n][3] = act1 && c + 1 >= lo1 && c + 1 <= pos1 ? texp(__fsub_rn(sc[n][3], n1)) : 0.0f; + c0[n] = __fadd_rn(sc[n][0], sc[n][1]); + c1[n] = __fadd_rn(sc[n][2], sc[n][3]); + } + l0 = __fmaf_rn(l0, a0, rowsum(c0)); + l1 = __fmaf_rn(l1, a1, rowsum(c1)); + m0 = n0; + m1 = n1; +#pragma unroll + for (int i = 0; i < NB; ++i) { + o[i][0] = __fmul_rn(o[i][0], a0); + o[i][1] = __fmul_rn(o[i][1], a0); + o[i][2] = __fmul_rn(o[i][2], a1); + o[i][3] = __fmul_rn(o[i][3], a1); + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + pa[j][0] = pack(sc[2 * j][0], sc[2 * j][1]); + pa[j][1] = pack(sc[2 * j][2], sc[2 * j][3]); + pa[j][2] = pack(sc[2 * j + 1][0], sc[2 * j + 1][1]); + pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); + } + } + values(pa[0], pa[1], take(it + 2)); + done(it + 2); + values(pa[2], pa[3], take(it + 3)); + done(it + 3); + } + wait_group<0>(); + const int r0 = wrow + g, r1 = r0 + 8; +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + if (r0 < W) + *reinterpret_cast(out + ((int64_t)r0 * H + head) * D + d) = + pack(tdiv(o[i][0], l0), tdiv(o[i][1], l0)); + if (r1 < W) + *reinterpret_cast(out + ((int64_t)r1 * H + head) * D + d) = + pack(tdiv(o[i][2], l1), tdiv(o[i][3], l1)); + } +} + + +// Decode on tensor cores (d 128): a row's heads one mma tile; warps take every 4th 32-key tile, fold in warp order. +extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, + __nv_bfloat16* __restrict__ out, const int* __restrict__ pos, const int* __restrict__ slot, int h, int hk, int win, + long long ring, float scale) { + constexpr int D = 128, C = D / 8, K16 = D / 16, NB = D / 8, HB = 32; + extern __shared__ uint4 smem[]; + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g8 = lane / 4, qd = lane % 4, mt = lane / 8; + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, G = h / hk; + uint4* ks = smem + warp * 2 * HB * C; + uint4* vs = ks + HB * C; + 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 + 1, 0) : 0; + const int lor = static_cast(lo % ring); + uint32_t qa[K16][4]; + { + const int h0 = g8, h1 = g8 + 8; + const uint32_t* q0 = reinterpret_cast(q + ((int64_t)r * h + kh * G + min(h0, G - 1)) * D); + const uint32_t* q1 = reinterpret_cast(q + ((int64_t)r * h + kh * G + min(h1, G - 1)) * D); +#pragma unroll + for (int k = 0; k < K16; ++k) { + qa[k][0] = h0 < G ? q0[8 * k + qd] : 0u; + qa[k][1] = h1 < G ? q1[8 * k + qd] : 0u; + qa[k][2] = h0 < G ? q0[8 * k + qd + 4] : 0u; + qa[k][3] = h1 < G ? q1[8 * k + qd + 4] : 0u; + } + } + 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; + // K of tile t loads while tile t - 4's values fold in, V of tile t while its scores run: one commit group each + auto stage = [&](uint4* dst, const __nv_bfloat16* src, int t) { + if (t <= p / HB) { + const int r0 = static_cast((t * HB) % ring); +#pragma unroll 4 + for (int i = lane; i < HB * C; i += 32) { + const int j = i / C, key = t * HB + j; + const bool in = key >= lo && key <= p; + int row = r0 + j; + if (row >= ring) row -= static_cast(ring); + cp16z(dst + sw(j, i % C), src + ((base + (in ? row : lor)) * hk + kh) * D + (i % C) * 8, in); + } + } + commit(); + }; + stage(ks, kc, lo / HB + warp); + stage(vs, vc, lo / HB + warp); + for (int t = lo / HB + warp; t <= p / HB; t += 4) { + wait_group<1>(); // tile t's K (its V may still be in flight) + __syncwarp(); + float sc[4][4]; +#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, ks + sw(8 * (n + mt / 2) + lane % 8, 2 * k + mt % 2)); + if (k == 0) { + mma0(sc[n], qa[k], b[0], b[1]); + mma0(sc[n + 1], qa[k], b[2], b[3]); + } else { + mma(sc[n], qa[k], b[0], b[1]); + mma(sc[n + 1], qa[k], b[2], b[3]); + } + } + } + __syncwarp(); + stage(ks, kc, t + 4); + const int kb = t * HB + 2 * qd; + float tm0 = -INFINITY, tm1 = -INFINITY; +#pragma unroll + for (int n = 0; n < 4; ++n) { + const int c = kb + 8 * n; + const bool i0 = c >= lo && c <= p, i1 = c + 1 >= lo && c + 1 <= p; + sc[n][0] = i0 ? __fmul_rn(sc[n][0], scale) : -INFINITY; + sc[n][1] = i1 ? __fmul_rn(sc[n][1], scale) : -INFINITY; + sc[n][2] = i0 ? __fmul_rn(sc[n][2], scale) : -INFINITY; + sc[n][3] = i1 ? __fmul_rn(sc[n][3], scale) : -INFINITY; + tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1])); + tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); + tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); + } + const bool act = tm0 != -INFINITY; // the keys' mask is the same for every head + const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; + const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + float c0 = 0.0f, c1 = 0.0f; +#pragma unroll + for (int n = 0; n < 4; ++n) { + sc[n][0] = sc[n][0] != -INFINITY ? texp(__fsub_rn(sc[n][0], n0)) : 0.0f; + sc[n][1] = sc[n][1] != -INFINITY ? texp(__fsub_rn(sc[n][1], n0)) : 0.0f; + sc[n][2] = sc[n][2] != -INFINITY ? texp(__fsub_rn(sc[n][2], n1)) : 0.0f; + sc[n][3] = sc[n][3] != -INFINITY ? texp(__fsub_rn(sc[n][3], n1)) : 0.0f; + c0 = __fadd_rn(c0, __fadd_rn(sc[n][0], sc[n][1])); + c1 = __fadd_rn(c1, __fadd_rn(sc[n][2], sc[n][3])); + } + for (int x = 1; x <= 2; x *= 2) { + c0 = __fadd_rn(c0, __shfl_xor_sync(0xffffffffu, c0, x)); + c1 = __fadd_rn(c1, __shfl_xor_sync(0xffffffffu, c1, x)); + } + l0 = __fmaf_rn(l0, a0, c0); + l1 = __fmaf_rn(l1, a1, c1); + m0 = n0; + m1 = n1; +#pragma unroll + for (int i = 0; i < NB; ++i) { + o[i][0] = __fmul_rn(o[i][0], a0); + o[i][1] = __fmul_rn(o[i][1], a0); + o[i][2] = __fmul_rn(o[i][2], a1); + o[i][3] = __fmul_rn(o[i][3], a1); + } + wait_group<1>(); // tile t's V (the next K may still be in flight) + __syncwarp(); + uint32_t pa[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) { + pa[j][0] = pack(sc[2 * j][0], sc[2 * j][1]); + pa[j][1] = pack(sc[2 * j][2], sc[2 * j][3]); + pa[j][2] = pack(sc[2 * j + 1][0], sc[2 * j + 1][1]); + pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); + } +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, vs + sw(8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa[0], b[0], b[1]); + mma(o[i + 1], pa[0], b[2], b[3]); + } +#pragma unroll + for (int i = 0; i < NB; i += 2) { + uint32_t b[4]; + ldmatrix4t(b, vs + sw(16 + 8 * (mt % 2) + lane % 8, i + mt / 2)); + mma(o[i], pa[1], b[0], b[1]); + mma(o[i + 1], pa[1], b[2], b[3]); + } + __syncwarp(); + stage(vs, vc, t + 4); + } + wait_group<0>(); + const int na = min(4, p / HB - lo / HB + 1); // warps that saw keys + if (na == 1) { // one warp's fold is the result (the merge below would scale it by exp(0) = 1) + if (warp == 0) { +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + if (g8 < G) + *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8) * D + d) = + pack(tdiv(o[i][0], l0), tdiv(o[i][1], l0)); + if (g8 + 8 < G) + *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8 + 8) * D + d) = + pack(tdiv(o[i][2], l1), tdiv(o[i][3], l1)); + } + } + return; + } + __syncthreads(); + float* os = reinterpret_cast(smem); // [4 warps][16 heads][D] + float* ms = os + 4 * 16 * D; // [4][16] + float* ls = ms + 4 * 16; +#pragma unroll + for (int i = 0; i < NB; ++i) { + const int d = 8 * i + 2 * qd; + os[(warp * 16 + g8) * D + d] = o[i][0]; + os[(warp * 16 + g8) * D + d + 1] = o[i][1]; + os[(warp * 16 + g8 + 8) * D + d] = o[i][2]; + os[(warp * 16 + g8 + 8) * D + d + 1] = o[i][3]; + } + if (qd == 0) { + ms[warp * 16 + g8] = m0, ls[warp * 16 + g8] = l0; + ms[warp * 16 + g8 + 8] = m1, ls[warp * 16 + g8 + 8] = l1; + } + __syncthreads(); + float* as = ls + 4 * 16; // [4][16]: warp w's factor for head hd + float* il = as + 4 * 16; // [16]: the folded l + if (threadIdx.x < 16) { + const int hd = threadIdx.x; + float mm = -INFINITY; + for (int w = 0; w < 4; ++w) mm = fmaxf(mm, ms[w * 16 + hd]); + float l = 0.0f; + for (int w = 0; w < 4; ++w) { + const float mw = ms[w * 16 + hd]; + const float a = mw == -INFINITY ? 0.0f : texp(__fsub_rn(mw, mm)); + as[w * 16 + hd] = a; + if (mw != -INFINITY) l = __fmaf_rn(a, ls[w * 16 + hd], l); + } + il[hd] = l; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < G * D / 2; idx += 128) { + const int hd = idx / (D / 2), e = 2 * (idx % (D / 2)); + float acc0 = 0.0f, acc1 = 0.0f; +#pragma unroll + for (int w = 0; w < 4; ++w) { + if (ms[w * 16 + hd] == -INFINITY) continue; + const float a = as[w * 16 + hd]; + const float2 v = *reinterpret_cast(os + (w * 16 + hd) * D + e); + acc0 = __fmaf_rn(a, v.x, acc0); + acc1 = __fmaf_rn(a, v.y, acc1); + } + const float l = il[hd]; + *reinterpret_cast(out + ((int64_t)r * h + kh * G + hd) * D + e) = pack(tdiv(acc0, l), tdiv(acc1, l)); + } +} + +} // namespace tf_window_prompt + +// Head dim 128, eight warps, four or eight query heads a block, eight staging slots. +template __global__ void tf_window_prompt::wattn_kernel<128, 8, 4, 8>(const __nv_bfloat16*, const __nv_bfloat16*, + const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float, int, int); +template __global__ void tf_window_prompt::wattn_kernel<128, 8, 8, 8>(const __nv_bfloat16*, const __nv_bfloat16*, + const __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, float, int, int); diff --git a/zig/src/core/dense_ref.zig b/zig/src/core/dense_ref.zig new file mode 100644 index 0000000000..355de01daa --- /dev/null +++ b/zig/src/core/dense_ref.zig @@ -0,0 +1,24 @@ +//! bf16 rows times weight rows in float64, the reference for every backend's dense row kernels (cuda/dense_rows.zig). +const std = @import("std"); + +/// out[n] = x . w[n] for one row; tol: fp32 error of k / 32 sequential terms and a 5-step tree, plus the result's ulp. +pub fn row(x: []const f64, w: []const f64, k: usize, out: []f64, tol: []f64) void { + for (out, tol, 0..) |*o, *t, j| { + var s: f64 = 0; + var mag: f64 = 0; + for (x[0..k], w[j * k ..][0..k]) |a, b| { + s += a * b; + mag += @abs(a * b); + } + o.* = s; + t.* = (@as(f64, @floatFromInt(k / 32 + 6)) * mag + @abs(s)) * std.math.pow(f64, 2, -24) + 1e-30; + } +} + +test "a row against hand values" { + var out: [2]f64 = undefined; + var tol: [2]f64 = undefined; + row(&.{ 1, 2 }, &.{ 3, 4, -1, 0.5 }, 2, &out, &tol); + try std.testing.expectEqualSlices(f64, &.{ 11, 0 }, &out); + try std.testing.expect(tol[0] > 0 and tol[1] > 0); +} diff --git a/zig/src/core/moe_ref.zig b/zig/src/core/moe_ref.zig new file mode 100644 index 0000000000..3789639ec5 --- /dev/null +++ b/zig/src/core/moe_ref.zig @@ -0,0 +1,55 @@ +//! MoE routing and combine in float64, the references every backend's MoE glue is checked against (cuda/moe_glue.zig). +const std = @import("std"); + +/// Sigmoid routing, as cuda/moe_glue.zig's Route. +pub const Route = struct { experts: usize, topk: usize, shared: bool, normalize: bool = false, scale: f64 = 1 }; + +/// A row's picks by the fp32 logit + bias (the kernels' ranking exactly; ties: lower id), weights in float64. +pub fn topk(logits: []const f32, bias: ?[]const f32, r: Route, pick: []u32, wts: []f64) void { + var taken: [4096]bool = @splat(false); + var sum: f64 = 0; + for (0..r.topk) |k| { + var best: ?usize = null; + var bv: f32 = -std.math.inf(f32); + for (logits[0..r.experts], 0..) |l, e| { + if (taken[e]) continue; + const v = l + (if (bias) |b| b[e] else 0); // one fp32 rounding, as the kernel's __fadd_rn + if (best == null or v > bv) { + best = e; + bv = v; + } + } + taken[best.?] = true; + pick[k] = @intCast(best.?); + wts[k] = 1 / (1 + @exp(-@as(f64, logits[best.?]))); + sum += wts[k]; + } + for (wts[0..r.topk]) |*w| w.* = (if (r.normalize) w.* / sum else w.*) * r.scale; + if (r.shared) { + pick[r.topk] = @intCast(r.experts); + wts[r.topk] = 1; + } +} + +/// One row: out[d] = sum over slots of w[k] * y[k][d]. +pub fn combine(y: []const f64, w: []const f64, dims: usize, out: []f64) void { + @memset(out, 0); + for (w, 0..) |wk, k| for (out, y[k * dims ..][0..dims]) |*o, v| { + o.* += wk * v; + }; +} + +test "picks rank by logit plus bias, weigh by the logit, and ties go to the lower id" { + var pick: [3]u32 = undefined; + var wts: [3]f64 = undefined; + topk(&.{ 0, 1, 1, -1 }, &.{ 0, 0, 0, 5 }, .{ .experts = 4, .topk = 2, .shared = true }, &pick, &wts); + try std.testing.expectEqualSlices(u32, &.{ 3, 1, 4 }, &pick); + try std.testing.expectApproxEqRel(@as(f64, 1 / (1 + @exp(@as(f64, 1)))), wts[0], 1e-15); + try std.testing.expectEqual(@as(f64, 1), wts[2]); + topk(&.{ 2, 2 }, null, .{ .experts = 2, .topk = 2, .shared = false, .normalize = true, .scale = 2 }, pick[0..2], wts[0..2]); + try std.testing.expectEqualSlices(u32, &.{ 0, 1 }, pick[0..2]); + try std.testing.expectApproxEqRel(@as(f64, 1), wts[0], 1e-15); + var out: [2]f64 = undefined; + combine(&.{ 1, 2, 10, 20 }, &.{ 0.5, 0.25 }, 2, &out); + try std.testing.expectEqualSlices(f64, &.{ 3, 6 }, &out); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 26cdc72a62..f73958cf91 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -29,6 +29,8 @@ pub const picture = @import("picture/picture.zig"); /// Qwen-VL image input every Qwen3.5-VL checkpoint shares: settings, geometry, M-RoPE positions, the tower's tables. pub const qwen_vl = @import("qwen_vl/qwen_vl.zig"); pub const attention_ref = @import("attention_ref.zig"); +pub const moe_ref = @import("moe_ref.zig"); +pub const dense_ref = @import("dense_ref.zig"); test { _ = safetensors; @@ -43,6 +45,8 @@ test { _ = act_ref; _ = contract; _ = attention_ref; + _ = moe_ref; + _ = dense_ref; _ = exl3_format; _ = exlr8; _ = eaq_container; diff --git a/zig/src/cuda/dense_rows.zig b/zig/src/cuda/dense_rows.zig new file mode 100644 index 0000000000..9241a4596e --- /dev/null +++ b/zig/src/cuda/dense_rows.zig @@ -0,0 +1,51 @@ +//! dense_rows.cu's bf16 rows times bf16 weight rows into fp32 (routers, output heads): a row's bits at any row count. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +const warps = 8; +/// The launch tiles: rows a tile shares each weight row across. +const tiles = [_]usize{ 1, 4, 16, 32 }; + +pub const DenseRows = struct { + mod: Module, + fns: [tiles.len]Function, + + pub fn load(d: *const Driver) !DenseRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.dense_rows); + errdefer mod.unload(); + var fns: [tiles.len]Function = undefined; + inline for (tiles, 0..) |t, i| fns[i] = try mod.function(std.fmt.comptimePrint("tf_dense_rows_{d}", .{t})); + return .{ .mod = mod, .fns = fns }; + } + + pub fn unload(g: *DenseRows) void { + g.mod.unload(); + } + + /// out [m, n] fp32 = x [m, k] bf16 times w [n, k] bf16 transposed; strides in elements, k % 8 == 0, 16 B aligned. + pub fn rows(g: *const DenseRows, s: Stream, x: u64, x_stride: usize, w: u64, out: u64, out_stride: usize, m: usize, n: usize, k: usize) !void { + if (k % 8 != 0 or x % 16 != 0 or w % 16 != 0 or x_stride % 8 != 0) return error.Misaligned; + if (m == 0 or n == 0) return; + const i: usize = if (m == 1) 0 else if (m <= 4) 1 else if (m <= 16) 2 else 3; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(x_stride))); + a.add(w); + a.add(out); + a.add(@as(i64, @intCast(out_stride))); + a.add(@as(i32, @intCast(m))); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(k))); + const grid: launch_.Dim3 = .{ .x = @intCast((n + warps - 1) / warps), .y = @intCast((m + tiles[i] - 1) / tiles[i]), .z = 1 }; + try launch_.launch(g.fns[i], .{ .grid = grid, .block = .{ .x = warps * 32 } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 7e0dd2b10a..85d74d684f 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -37,6 +37,9 @@ pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{} pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; 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 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 &.{}; pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; diff --git a/zig/src/cuda/moe_glue.zig b/zig/src/cuda/moe_glue.zig new file mode 100644 index 0000000000..948f71634f --- /dev/null +++ b/zig/src/cuda/moe_glue.zig @@ -0,0 +1,76 @@ +//! moe_glue.cu's routing and combine for every CUDA family: sigmoid top-k with bias, the weighted sum of slots. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const norms = @import("norms.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; +const Kind = norms.Kind; + +const threads = 256; +/// The most experts a row ranks and the most it keeps (moe_glue.cu MOST, the picked array). +pub const most_experts = 4096; +pub const most_topk = 64; +/// Experts a lane ranks in the warp form (moe_glue.cu PER): up to 512 experts take it. +const per_lane = 16; + +/// Sigmoid routing: top `topk` by logit + bias, weights sigmoid(logit) [normalized] * scale; `shared` adds slot topk. +pub const Route = struct { experts: usize, topk: usize, shared: bool, normalize: bool = false, scale: f32 = 1 }; + +pub const MoeGlue = struct { + mod: Module, + topk_fn: Function, + topk_warp_fn: Function, + combine_fn: Function, + + pub fn load(d: *const Driver) !MoeGlue { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.moe_glue); + errdefer mod.unload(); + return .{ .mod = mod, .topk_fn = try mod.function("tf_moe_topk_sigmoid"), .topk_warp_fn = try mod.function("tf_moe_topk_sigmoid_warp"), .combine_fn = try mod.function("tf_moe_combine") }; + } + + pub fn unload(g: *MoeGlue) void { + g.mod.unload(); + } + + /// fp32 logits [rows, experts] (+ fp32 bias) into int32 picks and fp32 weights [rows, topk (+1 with shared)]. + pub fn topk(g: *const MoeGlue, s: Stream, logits: u64, bias: ?u64, r: Route, pick: u64, wts: u64, rows: usize) !void { + if (r.experts > most_experts or r.topk == 0 or r.topk > most_topk or r.topk > r.experts) return error.RouteOutOfRange; + if (rows == 0) return; + const warp = r.experts <= 32 * per_lane and r.topk <= 32; + var a: launch_.Args = .{}; + a.add(logits); + a.add(bias orelse 0); + a.add(pick); + a.add(wts); + if (warp) a.add(@as(i32, @intCast(rows))); + a.add(@as(i32, @intCast(r.experts))); + a.add(@as(i32, @intCast(r.topk))); + a.add(@as(i32, @intFromBool(r.shared))); + a.add(@as(i32, @intFromBool(r.normalize))); + a.add(r.scale); + if (warp) return launch_.launch(g.topk_warp_fn, .{ .grid = .{ .x = @intCast((rows + 7) / 8), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a); + try launch_.launch(g.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + + /// out [rows, dims] = the sum, in slot order, of w [rows, slots] (fp32) times y [rows, slots, dims]. + pub fn combine(g: *const MoeGlue, s: Stream, y: u64, y_kind: Kind, w: u64, out: u64, out_kind: Kind, rows: usize, slots: usize, dims: usize) !void { + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(y); + a.add(@as(i32, @backingInt(y_kind))); + a.add(w); + a.add(out); + a.add(@as(i32, @backingInt(out_kind))); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intCast(dims))); + const cols: u32 = @intCast((dims + threads - 1) / threads); + try launch_.launch(g.combine_fn, .{ .grid = .{ .x = @intCast(rows), .y = cols, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index f48dd67738..21710818c9 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -43,6 +43,9 @@ pub const ck = @import("ck.zig"); pub const modelopt = @import("modelopt.zig"); pub const tree_attention = @import("tree_attention.zig"); pub const block_attention = @import("block_attention.zig"); +pub const window_attention = @import("window_attention.zig"); +pub const moe_glue = @import("moe_glue.zig"); +pub const dense_rows = @import("dense_rows.zig"); test { _ = @import("memory.zig"); @@ -67,4 +70,7 @@ test { _ = @import("tp_watch.zig"); _ = @import("tree_attention.zig"); _ = @import("block_attention.zig"); + _ = @import("window_attention.zig"); + _ = @import("moe_glue.zig"); + _ = @import("dense_rows.zig"); } diff --git a/zig/src/cuda/window_attention.zig b/zig/src/cuda/window_attention.zig new file mode 100644 index 0000000000..40569b7ef9 --- /dev/null +++ b/zig/src/cuda/window_attention.zig @@ -0,0 +1,115 @@ +//! window_attention.cu's sliding-window attention for every CUDA family: decode rows over a ring of keys. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +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 }; + +/// 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"; +const prompt_smem: u32 = 8 * 32 * 16 * 16; // NS * HALF * (D / 8) uint4 +/// tf_window_decode_tc: four warps' K and V slots of 32 keys (64 KiB), reused for the warps' partials. +const tc_smem: u32 = 4 * 2 * 32 * 16 * 16; + +pub const WindowAttention = struct { + mod: Module, + part_fn: Function, + merge_fn: Function, + tc_fn: Function, + prompt_fns: [2]Function, // 4, then 8 query heads a block + + pub fn load(d: *const Driver) !WindowAttention { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.window_attention); + errdefer mod.unload(); + var p: [2]Function = undefined; + inline for (.{ 4, 8 }, 0..) |hpc, i| { + p[i] = try mod.function(std.fmt.comptimePrint(prompt_name, .{hpc})); + try p[i].allowDynamicShared(prompt_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 }; + } + + pub fn unload(g: *WindowAttention) void { + g.mod.unload(); + } + + /// 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; + } + + /// The fp32 partials `decode` writes for up to `rows` rows: o, then m and l. + pub fn scratchBytes(w: Window, rows: usize) usize { + 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; + } + + /// 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 (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; + if (w.d == 128) { // tensor cores; the heads of a row are one mma tile + var c: launch_.Args = .{}; + for ([_]u64{ q, kc, vc, out, pos, slot orelse 0 }) |v| c.add(v); + c.add(@as(i32, @intCast(w.h))); + c.add(@as(i32, @intCast(w.hk))); + c.add(@as(i32, @intCast(w.win))); + c.add(@as(i64, @intCast(w.ring))); + c.add(w.scale); + return launch_.launch(g.tc_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = 1 }, .block = .{ .x = 128 }, .shared = tc_smem }, s, &c); + } + const t = tiles(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.d))); + a.add(@as(i32, @intCast(w.win))); + a.add(@as(i64, @intCast(w.ring))); + a.add(w.scale); + try launch_.launch(g.part_fn, .{ .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(w.d))); + b.add(@as(i32, @intCast(t))); + try launch_.launch(g.merge_fn, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = 1 }, .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 { + 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; + const hpc: usize = if (grp % 8 == 0) 8 else 4; + var a: launch_.Args = .{}; + for ([_]u64{ q, kc, vc, out }) |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 rows = 16 * (8 / hpc); + 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); + } +}; diff --git a/zig/tests/cuda/dense_rows_tests.zig b/zig/tests/cuda/dense_rows_tests.zig new file mode 100644 index 0000000000..7ec56a424b --- /dev/null +++ b/zig/tests/cuda/dense_rows_tests.zig @@ -0,0 +1,79 @@ +//! dense_rows.cu against core/dense_ref.zig's float64 rows, and every row's bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const counts = [_]usize{ 1, 3, 17, 40, 300 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn case(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.dense_rows.DenseRows, rng: std.Random, n: usize, k: usize, worst: *f64) !void { + const xs = try a.alloc(u16, most * k); + const xv = try a.alloc(f64, most * k); + for (xs, xv, 0..) |*h, *v, i| { + h.* = bf16(rng.floatNorm(f32) * (if (i % 61 == 0) @as(f32, 12) else 1)); + v.* = @as(f32, @bitCast(@as(u32, h.*) << 16)); + } + const ws = try a.alloc(u16, n * k); + const wv = try a.alloc(f64, n * k); + for (ws, wv) |*h, *v| { + h.* = bf16(rng.floatNorm(f32) * 0.02); + v.* = @as(f32, @bitCast(@as(u32, h.*) << 16)); + } + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*b| b.free(); + bufs.deinit(gpu.gpa); + } + const dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xs)); + try bufs.append(gpu.gpa, dx); + const dw = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(ws)); + try bufs.append(gpu.gpa, dw); + const outs = try a.alloc(cuda.DeviceBuffer, counts.len); + for (counts, outs) |m, *o| { + o.* = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 4); + try bufs.append(gpu.gpa, o.*); + try o.fill8(0xff, s.handle); + try g.rows(s, dx.ptr, k, dw.ptr, o.ptr, n, m, n, k); + } + try s.synchronize(); + const wide = try a.alloc(u8, most * n * 4); + try outs[counts.len - 1].download(0, wide); + const want = try a.alloc(f64, n); + const tol = try a.alloc(f64, n); + var q_most: f64 = 0; + for (0..most) |r| { + core.dense_ref.row(xv[r * k ..][0..k], wv, k, want, tol); + for (want, tol, 0..) |w, t, j| { + const got: f32 = @bitCast(std.mem.readInt(u32, wide[4 * (r * n + j) ..][0..4], .little)); + q_most = @max(q_most, @abs(got - w) / t); + } + } + worst.* = @max(worst.*, q_most); + try check.expect(q_most <= 1, "dense rows [{d}, {d}]: within float64 tolerance (worst {d:.3})", .{ n, k, q_most }); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, o| { + const got = try a.alloc(u8, m * n * 4); + try o.download(0, got); + try check.sameBytes("dense rows: a row count's bits equal the widest call's", got, wide[0 .. m * n * 4]); + } +} + +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 cuda.dense_rows.DenseRows.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(29); + var worst: f64 = 0; + try case(gpu, arena.allocator(), s, &g, prng.random(), 384, 2560, &worst); // Kolibri 1's router + try case(gpu, arena.allocator(), s, &g, prng.random(), 4103, 2560, &worst); // a head's shape, odd columns + try case(gpu, arena.allocator(), s, &g, prng.random(), 257, 7168, &worst); + check.pass("dense rows: within float64 tolerance (worst {d:.3}), every row count's bits equal", .{worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f1e2682483..97f57fddbf 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -40,6 +40,9 @@ const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); const tree_tests = @import("tree_tests.zig"); const block_tests = @import("block_tests.zig"); +const window_tests = @import("window_tests.zig"); +const moe_glue_tests = @import("moe_glue_tests.zig"); +const dense_rows_tests = @import("dense_rows_tests.zig"); const usage = \\usage: tf-cuda-test @@ -91,6 +94,9 @@ const usage = \\ exlr8-dense [DIR] EXLR8 dense formats: decode bit-equal, FP64, row windows; DIR adds the real matrices \\ 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) + \\ 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) \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR @@ -194,6 +200,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); 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, "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); if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); diff --git a/zig/tests/cuda/moe_glue_tests.zig b/zig/tests/cuda/moe_glue_tests.zig new file mode 100644 index 0000000000..2445e6dedf --- /dev/null +++ b/zig/tests/cuda/moe_glue_tests.zig @@ -0,0 +1,161 @@ +//! moe_glue.cu against core/moe_ref.zig: picks exact, weights and sums in float64 tolerance, bits at any row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const ref = core.moe_ref; +const counts = [_]usize{ 1, 3, 17, 300 }; +const most = counts[counts.len - 1]; + +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + g: *const cuda.moe_glue.MoeGlue, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + worst: f64 = 0, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(bytes, 4)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); + return b.ptr; + } + + fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const o = try r.a.alloc(u8, len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, o); + return o; + } +}; + +fn f32At(raw: []const u8, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); +} + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn topkCase(r: *Rig, route: cuda.moe_glue.Route, with_bias: bool) !void { + const ne = route.experts; + const logits = try r.a.alloc(f32, most * ne); + for (logits, 0..) |*l, i| { + l.* = r.rng.floatNorm(f32) * 2; + if (i % 7 == 3) l.* = logits[i - 1]; // equal neighbours: the tie rule decides + } + const bias = try r.a.alloc(f32, ne); + for (bias, 0..) |*b, i| b.* = if (i % 5 == 0) 0.25 else 0; + const slots = route.topk + @intFromBool(route.shared); + const dl = try r.dev(std.mem.sliceAsBytes(logits)); + const db = try r.dev(std.mem.sliceAsBytes(bias)); + const want_p = try r.a.alloc(u32, most * slots); + const want_w = try r.a.alloc(f64, most * slots); + const rr: ref.Route = .{ .experts = ne, .topk = route.topk, .shared = route.shared, .normalize = route.normalize, .scale = route.scale }; + for (0..most) |row| ref.topk(logits[row * ne ..][0..ne], if (with_bias) bias else null, rr, want_p[row * slots ..][0..slots], want_w[row * slots ..][0..slots]); + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "topk {d} of {d}{s}{s}{s}", .{ route.topk, ne, if (route.shared) " + shared" else "", if (with_bias) ", bias" else "", if (route.normalize) ", normalized" else "" }); + const ps = try r.a.alloc(u64, counts.len); + const ws = try r.a.alloc(u64, counts.len); + for (counts, ps, ws) |m, *p, *w| { + p.* = try r.out(m * slots * 4); + w.* = try r.out(m * slots * 4); + try r.g.topk(r.s, dl, if (with_bias) db else null, route, p.*, w.*, m); + } + const gp = try r.back(ps[counts.len - 1], most * slots * 4); + const gw = try r.back(ws[counts.len - 1], most * slots * 4); + const picks = std.mem.bytesAsSlice(u32, @as([]align(1) u8, gp)); + var q_most: f64 = 0; + for (want_p, 0..) |p, i| try check.expect(picks[i] == p, "{s}: pick {d} is {d}, want {d}", .{ label, i, picks[i], p }); + for (want_w, 0..) |w, i| q_most = @max(q_most, @abs(f32At(gw, i) - w) / (@abs(w) * std.math.pow(f64, 2, -20) + 1e-30)); + r.worst = @max(r.worst, q_most); + try check.expect(q_most <= 1, "{s}: weights within 2^-20 (worst {d:.3})", .{ label, q_most }); + for (counts[0 .. counts.len - 1], ps[0 .. counts.len - 1], ws[0 .. counts.len - 1]) |m, p, w| { + try check.sameBytes(label, try r.back(p, m * slots * 4), gp[0 .. m * slots * 4]); + try check.sameBytes(label, try r.back(w, m * slots * 4), gw[0 .. m * slots * 4]); + } +} + +fn combineCase(r: *Rig, slots: usize, dims: usize, y_bf16: bool) !void { + const n = most * slots * dims; + const yv = try r.a.alloc(f64, n); + const ybytes = try r.a.alloc(u8, n * @as(usize, if (y_bf16) 2 else 4)); + for (yv, 0..) |*v, i| { + const x = r.rng.floatNorm(f32) * (if (i % 97 == 0) @as(f32, 30) else 1); + if (y_bf16) { + const h = bf16(x); + std.mem.writeInt(u16, ybytes[2 * i ..][0..2], h, .little); + v.* = @as(f32, @bitCast(@as(u32, h) << 16)); + } else { + std.mem.writeInt(u32, ybytes[4 * i ..][0..4], @bitCast(x), .little); + v.* = x; + } + } + const w = try r.a.alloc(f32, most * slots); + for (w) |*x| x.* = r.rng.float(f32); + const wv = try r.a.alloc(f64, w.len); + for (w, wv) |x, *v| v.* = x; + const dy = try r.dev(ybytes); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const want = try r.a.alloc(f64, most * dims); + const tol = try r.a.alloc(f64, most * dims); + for (0..most) |row| { + ref.combine(yv[row * slots * dims ..][0 .. slots * dims], wv[row * slots ..][0..slots], dims, want[row * dims ..][0..dims]); + for (0..dims) |d| { // bf16 half-ulp of the sum, plus fp32 fma error over the slots + var mag: f64 = 0; + for (0..slots) |k| mag += @abs(wv[row * slots + k] * yv[(row * slots + k) * dims + d]); + tol[row * dims + d] = @abs(want[row * dims + d]) * std.math.pow(f64, 2, -8) + mag * @as(f64, @floatFromInt(slots)) * std.math.pow(f64, 2, -24) + 1e-30; + } + } + var name: [96]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "combine {d} slots x {d} ({s} -> bf16)", .{ slots, dims, if (y_bf16) "bf16" else "fp32" }); + const outs = try r.a.alloc(u64, counts.len); + for (counts, outs) |m, *o| { + o.* = try r.out(m * dims * 2); + try r.g.combine(r.s, dy, if (y_bf16) .bf16 else .f32, dw, o.*, .bf16, m, slots, dims); + } + const wide = try r.back(outs[counts.len - 1], most * dims * 2); + var q_most: f64 = 0; + for (want, tol, 0..) |x, t, i| { + const got: f64 = @as(f32, @bitCast(@as(u32, std.mem.readInt(u16, wide[2 * i ..][0..2], .little)) << 16)); + q_most = @max(q_most, @abs(got - x) / t); + } + r.worst = @max(r.worst, q_most); + try check.expect(q_most <= 1, "{s}: within its float64 tolerance (worst {d:.3})", .{ label, q_most }); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, o| try check.sameBytes(label, try r.back(o, m * dims * 2), wide[0 .. m * dims * 2]); +} + +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 cuda.moe_glue.MoeGlue.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(23); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .g = &g, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + try topkCase(&r, .{ .experts = 384, .topk = 6, .shared = true }, true); // Kolibri 1 + try topkCase(&r, .{ .experts = 256, .topk = 8, .shared = false, .normalize = true, .scale = 2.5 }, true); // DeepSeek-style + try topkCase(&r, .{ .experts = 512, .topk = 32, .shared = true }, true); // the warp form's widest + try topkCase(&r, .{ .experts = 513, .topk = 6, .shared = true }, true); // the block form's narrowest + try topkCase(&r, .{ .experts = 4096, .topk = 64, .shared = false }, false); + try combineCase(&r, 7, 2560, false); + try combineCase(&r, 7, 2560, true); + try combineCase(&r, 9, 7168, false); + check.pass("moe glue: picks exact, weights and sums within float64 tolerance (worst {d:.3}), every row count's bits equal", .{r.worst}); +} diff --git a/zig/tests/cuda/window_tests.zig b/zig/tests/cuda/window_tests.zig new file mode 100644 index 0000000000..c3d8c68899 --- /dev/null +++ b/zig/tests/cuda/window_tests.zig @@ -0,0 +1,217 @@ +//! window_attention.cu against core/attention_ref.zig: decode rows over a ring of keys, bits at any row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const counts = [_]usize{ 1, 3, 17, 64 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Case = struct { h: usize, hk: usize, d: usize, win: usize, ring: usize, streams: usize, top: usize }; + +fn run1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.window_attention.WindowAttention, rng: std.Random, c: Case, worst: *f64) !void { + const gq = c.h / c.hk; + const qn = most * c.h * c.d; + const cn = c.streams * c.ring * c.hk * c.d; + const qs = try a.alloc(u16, qn); + const ks = try a.alloc(u16, cn); + const vs = try a.alloc(u16, cn); + for (qs) |*x| x.* = bf16(rng.floatNorm(f32)); + for (ks) |*x| x.* = bf16(rng.floatNorm(f32)); + for (vs) |*x| x.* = bf16(rng.floatNorm(f32)); + 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 % 4 == 0) i % 40 else (i * 7919) % c.top); // short histories and wrapped rings both + sl.* = @intCast(i % c.streams); + } + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*b| b.free(); + bufs.deinit(gpu.gpa); + } + const up = struct { + fn f(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(gp.d, bytes); + try l.append(gp.gpa, b); + return b.ptr; + } + }; + const dq = try up.f(gpu, &bufs, std.mem.sliceAsBytes(qs)); + const dk = try up.f(gpu, &bufs, std.mem.sliceAsBytes(ks)); + const dv = try up.f(gpu, &bufs, std.mem.sliceAsBytes(vs)); + const dp = try up.f(gpu, &bufs, std.mem.sliceAsBytes(pos)); + const dsl = try up.f(gpu, &bufs, std.mem.sliceAsBytes(slot)); + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(c.d)))); + const w: cuda.window_attention.Window = .{ .h = c.h, .hk = c.hk, .d = c.d, .win = c.win, .ring = c.ring, .scale = scale }; + const scratch = try cuda.DeviceBuffer.alloc(gpu.d, cuda.window_attention.WindowAttention.scratchBytes(w, most)); + try bufs.append(gpu.gpa, scratch); + const outs = try a.alloc(cuda.DeviceBuffer, counts.len); + for (counts, outs) |m, *o| { + o.* = try cuda.DeviceBuffer.alloc(gpu.d, m * c.h * c.d * 2); + try bufs.append(gpu.gpa, o.*); + try o.fill8(0xff, s.handle); + try g.decode(s, dq, dk, dv, o.ptr, dp, dsl, w, scratch.ptr, m); + } + try s.synchronize(); + const wide = try a.alloc(u16, qn); + try outs[counts.len - 1].download(0, std.mem.sliceAsBytes(wide)); + // per (row, key head): the window's keys gathered from the ring into dense rows, then each query head + const kd = try a.alloc(f64, c.top * c.d); + const vd = try a.alloc(f64, c.top * c.d); + const keys = try a.alloc(usize, c.top); + const q = try a.alloc(f64, c.d); + const want = try a.alloc(f64, c.d); + const tol = try a.alloc(f64, c.d); + var q_most: f64 = 0; + for (0..most) |r| { + const p: usize = @intCast(pos[r]); + const lo = if (c.win > 0 and p + 1 > c.win) p + 1 - c.win else 0; + const nk = p + 1 - lo; + for (0..c.hk) |kh| { + for (0..nk) |j| { + const row = @as(usize, @intCast(slot[r])) * c.ring + (lo + j) % c.ring; + for (0..c.d) |e| { + kd[j * c.d + e] = bf(ks[(row * c.hk + kh) * c.d + e]); + vd[j * c.d + e] = bf(vs[(row * c.hk + kh) * c.d + e]); + } + keys[j] = j; + } + for (0..gq) |hh| { + const hd = kh * gq + hh; + for (q, 0..) |*x, e| x.* = bf(qs[(r * c.h + hd) * c.d + e]); + core.attention_ref.head(q, kd, vd, keys[0..nk], c.d, scale, want, tol); + for (want, tol, 0..) |x, t, e| q_most = @max(q_most, @abs(bf(wide[(r * c.h + hd) * c.d + e]) - x) / t); + } + } + } + worst.* = @max(worst.*, q_most); + try check.expect(q_most <= 1, "window decode h {d}/{d} d {d} win {d}: within float64 tolerance (worst {d:.3})", .{ c.h, c.hk, c.d, c.win, q_most }); + 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 o.download(0, std.mem.sliceAsBytes(got)); + try check.sameBytes("window 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])); + } +} + +const PCase = struct { h: usize, hk: usize, win: usize, ring: usize, p0: usize, n: usize }; + +/// Prompt rows p0 .. p0 + n: float64 tolerance, a shorter call's bits, a later chunk's bits (from p0 + n / 2). +fn prompt1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.window_attention.WindowAttention, rng: std.Random, c: PCase, worst: *f64) !void { + const d = 128; + const qn = c.n * c.h * d; + const crow = if (c.ring > 0) c.ring else c.p0 + c.n; // ring 0: a row a position + const cn = crow * c.hk * d; + const qs = try a.alloc(u16, qn); + const ks = try a.alloc(u16, cn); + const vs = try a.alloc(u16, cn); + for (qs) |*x| x.* = bf16(rng.floatNorm(f32)); + for (ks) |*x| x.* = bf16(rng.floatNorm(f32)); + for (vs) |*x| x.* = bf16(rng.floatNorm(f32)); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*b| b.free(); + bufs.deinit(gpu.gpa); + } + const up = struct { + fn f(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(gp.d, bytes); + try l.append(gp.gpa, b); + return b.ptr; + } + fn z(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), st: cuda.Stream, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(gp.d, n); + try l.append(gp.gpa, b); + try b.fill8(0xff, st.handle); + return b.ptr; + } + }; + const dq = try up.f(gpu, &bufs, std.mem.sliceAsBytes(qs)); + const dk = try up.f(gpu, &bufs, std.mem.sliceAsBytes(ks)); + const dv = try up.f(gpu, &bufs, std.mem.sliceAsBytes(vs)); + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, d))); + const w: cuda.window_attention.Window = .{ .h = c.h, .hk = c.hk, .d = d, .win = c.win, .ring = c.ring, .scale = scale }; + const full = try up.z(gpu, &bufs, s, qn * 2); + try g.prompt(s, dq, dk, dv, full, c.p0, c.n, w); + const pre_n = 17; + const pre = try up.z(gpu, &bufs, s, pre_n * c.h * d * 2); + try g.prompt(s, dq, dk, dv, pre, c.p0, pre_n, w); + const half = c.n / 2; + const late = try up.z(gpu, &bufs, s, (c.n - half) * c.h * d * 2); + try g.prompt(s, dq + half * c.h * d * 2, dk, dv, late, c.p0 + half, c.n - half, w); + try s.synchronize(); + const got = try a.alloc(u16, qn); + const fb: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = full, .len = qn * 2 }; + try fb.download(0, std.mem.sliceAsBytes(got)); + const gq = c.h / c.hk; + const top = c.p0 + c.n; + const kd = try a.alloc(f64, @min(top, c.win + 0) * d + top * d); + const vd = try a.alloc(f64, kd.len); + const keys = try a.alloc(usize, top); + const q = try a.alloc(f64, d); + const want = try a.alloc(f64, d); + const tol = try a.alloc(f64, d); + var q_most: f64 = 0; + var r: usize = 0; + while (r < c.n) : (r += 7) { // every seventh row, and the last + const p = c.p0 + r; + const lo = if (c.win > 0 and p + 1 > c.win) p + 1 - c.win else 0; + const nk = p + 1 - lo; + for (0..c.hk) |kh| { + for (0..nk) |j| { + const row = if (c.ring > 0) (lo + j) % c.ring else lo + j; + for (0..d) |e| { + kd[j * d + e] = bf(ks[(row * c.hk + kh) * d + e]); + vd[j * d + e] = bf(vs[(row * c.hk + kh) * d + e]); + } + keys[j] = j; + } + for (0..gq) |hh| { + const hd = kh * gq + hh; + for (q, 0..) |*x, e| x.* = bf(qs[(r * c.h + hd) * d + e]); + core.attention_ref.head(q, kd, vd, keys[0..nk], d, scale, want, tol); + for (want, tol, 0..) |x, t, e| q_most = @max(q_most, @abs(bf(got[(r * c.h + hd) * d + e]) - x) / t); + } + } + if (r + 7 >= c.n and r != c.n - 1) r = c.n - 1 - 7; + } + worst.* = @max(worst.*, q_most); + try check.expect(q_most <= 1, "window prompt h {d}/{d} win {d} ring {d} p0 {d} n {d}: within float64 tolerance (worst {d:.3})", .{ c.h, c.hk, c.win, c.ring, c.p0, c.n, q_most }); + const pb: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = pre, .len = pre_n * c.h * d * 2 }; + const pg = try a.alloc(u16, pre_n * c.h * d); + try pb.download(0, std.mem.sliceAsBytes(pg)); + try check.sameBytes("window prompt: a shorter call's rows", std.mem.sliceAsBytes(pg), std.mem.sliceAsBytes(got[0 .. pre_n * c.h * d])); + const lb: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = late, .len = (c.n - half) * c.h * d * 2 }; + const lg = try a.alloc(u16, (c.n - half) * c.h * d); + try lb.download(0, std.mem.sliceAsBytes(lg)); + try check.sameBytes("window prompt: a later chunk's rows", std.mem.sliceAsBytes(lg), std.mem.sliceAsBytes(got[half * c.h * d ..])); +} + +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 cuda.window_attention.WindowAttention.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(37); + var worst: f64 = 0; + const a = arena.allocator(); + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 640, .streams = 2, .top = 3000 }, &worst); // Kolibri + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 600, .streams = 2, .top = 3000 }, &worst); // the ring's end inside a key tile + try run1(gpu, a, s, &g, prng.random(), .{ .h = 8, .hk = 2, .d = 64, .win = 100, .ring = 128, .streams = 3, .top = 1000 }, &worst); + try run1(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 1, .d = 128, .win = 0, .ring = 300, .streams = 1, .top = 300 }, &worst); // no window + try prompt1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .win = 513, .ring = 2624, .p0 = 0, .n = 600 }, &worst); // Kolibri, first chunk + try prompt1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .win = 513, .ring = 2624, .p0 = 2300, .n = 900 }, &worst); // the ring wraps + try prompt1(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 2, .win = 0, .ring = 0, .p0 = 100, .n = 333 }, &worst); // no window, 8 heads a block + check.pass("window attention: decode and prompt rows within float64 tolerance (worst {d:.3}), bits at any row count and chunk", .{worst}); +} From 31d3baf48020b992e2e0c3a445475adad94535e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 07:40:22 +0200 Subject: [PATCH 13/41] cuda/act_rows, gdn_prep, qk_rows: registered with the kernel contract --- zig/src/core/contract/registry.zig | 15 ++ zig/tests/cuda/contract_act.zig | 352 +++++++++++++++++++++++++++++ zig/tests/cuda/contract_gdn.zig | 305 +++++++++++++++++++++++++ zig/tests/cuda/contract_qk.zig | 271 ++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 23 +- 5 files changed, 965 insertions(+), 1 deletion(-) create mode 100644 zig/tests/cuda/contract_act.zig create mode 100644 zig/tests/cuda/contract_gdn.zig create mode 100644 zig/tests/cuda/contract_qk.zig diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..565030a2c8 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,21 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows", .symbol = "tf_swiglu_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows, 16 rows or fewer", .symbol = "tf_swiglu_rows", .image = "act_rows", .threads = 128 }, .usage = .{ .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_gate_rows", .symbol = "tf_gate_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_gate_rows, 16 rows or fewer", .symbol = "tf_gate_rows", .image = "act_rows", .threads = 64 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_group_sums_rows", .symbol = "tf_group_sums_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_dconv_rows", .symbol = "tf_dconv_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_64", .symbol = "tf_gdn_prep_64", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_64, 16 rows or fewer", .symbol = "tf_gdn_prep_64", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_128", .symbol = "tf_gdn_prep_128", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_128, 16 rows or fewer", .symbol = "tf_gdn_prep_128", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_256", .symbol = "tf_gdn_prep_256", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_256, 16 rows or fewer", .symbol = "tf_gdn_prep_256", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows at dim 64", .symbol = "tf_qk_rows", .image = "qk_rows", .threads = 64 }, .usage = .{ .{ .regs = 30, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows at dim 256", .symbol = "tf_qk_rows", .image = "qk_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows_w", .symbol = "tf_qk_rows_w", .image = "qk_rows", .threads = 128 }, .usage = .{ .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_act.zig b/zig/tests/cuda/contract_act.zig new file mode 100644 index 0000000000..1f4804f713 --- /dev/null +++ b/zig/tests/cuda/contract_act.zig @@ -0,0 +1,352 @@ +//! act_rows.cu (#797) under the contract on CUDA: SwiGLU, the output gate, 64-value sums and DFlash2's two-tap conv. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const aref = core.act_ref; + +pub const Which = enum { swiglu, swiglu_sums, gate, gate_sums, group_sums, dconv, dconv_res }; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// A thread's 8 values in order, then 8 lanes' butterfly. +fn groupDepth(_: c.Shape) usize { + return 11; +} + +/// x k0 rounded, then the fma that adds the previous row's tap. +fn convDepth(_: c.Shape) usize { + return 2; +} + +const gs = 16; +const tiles_rows: c.Tile = .{ .size = 16 }; // launch.rowThreads: 16 rows or fewer take the small blocks +const tiles_n: c.Tile = .{ .size = 2048, .step = 64, .min = 64 }; // 256 threads of 8 values +const swiglu_served = &[_]c.Shape{ .{ .rows = 1, .n = 17408, .k = 1 }, .{ .rows = 64, .n = 17408, .k = 1 } }; // the 27B's MLP +const gate_served = &[_]c.Shape{ .{ .rows = 1, .n = 24, .k = 256 }, .{ .rows = 64, .n = 24, .k = 256 } }; // its output gate +const conv_served = &[_]c.Shape{ .{ .rows = 16, .n = 5120, .k = 16 }, .{ .rows = 64, .n = 5120, .k = 16 } }; // DFlash2's block + +/// silu and sigmoid: expf (2 ulp, 4 units of 2^-24), the add, the division and the products. +const act_ops = 8; +const res_rule = "own check, core/act_ref's bound: the conv is rounded to bf16 before the residual add, an error of the conv's size"; + +/// swiglu, gate: rows x outputs (gate: heads x dim); sums: rows x 64-value groups; dconv: rows x channels x segment. +pub fn spec(which: Which) c.Spec { + return switch (which) { + .swiglu => .{ .name = "cuda act_rows tf_swiglu_rows", .accum = .none, .depth = noSums, .ops = act_ops, .out = .bf16, .tiles = .{ tiles_rows, tiles_n, .{} }, .served = swiglu_served }, + .swiglu_sums => .{ .name = "cuda act_rows tf_swiglu_rows group sums (of its stored bf16 outputs)", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ tiles_rows, tiles_n, .{} }, .served = swiglu_served }, + .gate => .{ .name = "cuda act_rows tf_gate_rows", .accum = .none, .depth = noSums, .ops = act_ops, .out = .bf16, .tiles = .{ tiles_rows, .{}, .{ .size = 256, .step = 64, .min = 64 } }, .served = gate_served }, + .gate_sums => .{ .name = "cuda act_rows tf_gate_rows group sums (of its stored bf16 outputs)", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ tiles_rows, .{}, .{ .size = 256, .step = 64, .min = 64 } }, .served = gate_served }, + .group_sums => .{ .name = "cuda act_rows tf_group_sums_rows", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ .{}, .{}, tiles_n }, .served = &.{ .{ .rows = 1, .n = 1, .k = 5120 }, .{ .rows = 64, .n = 1, .k = 17408 } } }, + .dconv => .{ .name = "cuda act_rows tf_dconv_rows", .accum = .f32, .depth = convDepth, .out = .bf16, .tiles = .{ tiles_rows, .{ .size = 2048, .step = gs, .min = gs }, .{} }, .served = conv_served }, + .dconv_res => .{ .name = "cuda act_rows tf_dconv_rows with residual", .accum = .f32, .depth = convDepth, .out = .bf16, .nan = .{ .stated = res_rule }, .tiles = .{ tiles_rows, .{ .size = 2048, .step = gs, .min = gs }, .{} }, .served = conv_served }, + }; +} + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +fn word(v: f32) u16 { + return c.toBf16(v); +} + +/// Gate 1's places: NaN, +inf, -inf, then both infinities in row 3 (rows that exist), in `x` [rows, w]. +fn plant(x: []u16, rows_: usize, w: usize, first: bool) void { + const at = [_][2]usize{ .{ 0, w / 3 }, .{ 1, w / 2 }, .{ 2, w - 1 }, .{ 3, 0 }, .{ 3, w - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put, 0..) |p, v, i| if (p[0] < rows_ and (i < 2) == first) { + x[p[0] * w + p[1]] = word(v); + }; +} + +/// Gate 4's 64-value groups: 16 integers, then their negatives (four nudged by one), twice: big partials in each order. +fn mirrored(rng: std.Random, x: []u16, scale: f32, top: u32) void { + var i: usize = 0; + while (i + 64 <= x.len) : (i += 64) for (0..2) |h| for (0..16) |j| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, top)); + x[i + 32 * h + j] = word(v * scale); + x[i + 32 * h + 16 + j] = word((-v + @as(f32, if (j < 4) 1 else 0)) * scale); + }; +} + +pub const Act = struct { + rig: rig_.Rig, + which: Which, + k: *const cuda.act_rows.ActRows, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + y: []u16 = &.{}, + base: []u16 = &.{}, + dyn: []u16 = &.{}, + res: []u16 = &.{}, + in: [5]u64 = @splat(0), + out: u64 = 0, + xs: u64 = 0, + fill: g.Fill = .random, + ran: bool = false, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Act) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Act) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = if (self.which == .dconv_res) &vt_res else &vt }; + } + + const vt = g.hooks(Act); + const vt_res = blk: { + var v = g.hooks(Act); + v.check = checkRes; + break :blk v; + }; + + /// Outputs a row's call writes: values, or (sums) one a 64-value group. + fn width(self: *const Act) usize { + return switch (self.which) { + .gate, .gate_sums => self.shape.n * self.shape.k, + .group_sums => self.shape.k, + else => self.shape.n, + }; + } + + fn sums(self: *const Act) bool { + return self.which == .swiglu_sums or self.which == .gate_sums or self.which == .group_sums; + } + + fn bfs(self: *Act, n: usize, rng: std.Random, scale: f32) ![]u16 { + const v = try self.arena.allocator().alloc(u16, n); + for (v, 0..) |*x, i| x.* = word((rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); + return v; + } + + pub fn prepare(self: *Act, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + self.fill = fill; + self.ran = false; + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const w = self.width(); + if (w % 64 != 0 and self.which != .dconv and self.which != .dconv_res) return false; + switch (self.which) { + .swiglu, .swiglu_sums, .gate, .gate_sums => { + // x: the gates ([gate | up] rows, or [q | gate] heads), y: up or o + const gated = self.which == .gate or self.which == .gate_sums; + self.x = try self.bfs(r * 2 * w, rng, 3); + self.y = try self.bfs(r * w, rng, 1); + if (fill == .nonfinite) { + plant(self.y, r, w, false); + for (0..2) |i| if (i < r) { + self.x[i * 2 * w + (if (gated) shape.k else 0) + (if (i == 0) w / 3 else w / 2) % shape.k] = word(if (i == 0) std.math.nan(f32) else std.math.inf(f32)); + }; + } + if (fill == .stress) { + @memset(self.x, word(17)); // silu(17) and sigmoid(17) are 17 and 1 in fp32 + mirrored(rng, self.y, if (gated) 17 else 1, 15); // outputs 17 u: odd, up to 255 + } + const both = try self.arena.allocator().alloc(u16, r * 2 * w); + @memcpy(both, self.x); + // [gate | up]: up after a row's gates; [q | gate] heads: the gates are the second half of each head + if (!gated) for (0..r) |i| @memcpy(both[i * 2 * w + w ..][0..w], self.y[i * w ..][0..w]); + self.in[0] = try self.rig.input(std.mem.sliceAsBytes(both)); + if (gated) self.in[1] = try self.rig.input(std.mem.sliceAsBytes(self.y)); + if (self.sums()) { + self.out = try self.rig.scratch(r * w * 2); + self.xs = try self.rig.output(r * w / 64 * 4); + } else { + self.out = try self.rig.output(r * w * 2); + self.xs = try self.rig.scratch(r * w / 64 * 4); + } + }, + .group_sums => { + self.x = try self.bfs(r * (w + 64), rng, 2); + if (fill == .nonfinite) plant(self.x, r, w + 64, true); + if (fill == .nonfinite) plant(self.x, r, w + 64, false); + if (fill == .stress) for (0..r) |i| mirrored(rng, self.x[i * (w + 64) ..][0..w], 1, 255); + self.in[0] = try self.rig.input(std.mem.sliceAsBytes(self.x)); + self.xs = try self.rig.output(r * w / 64 * 4); + }, + .dconv, .dconv_res => { + if (w % gs != 0) return false; + self.x = try self.bfs(r * w, rng, 1); + self.dyn = try self.bfs(r * 4 * (w / gs), rng, 0.3); + self.base = try self.bfs(4 * w, rng, 0.7); + self.res = try self.bfs(r * w, rng, 1); + if (fill == .nonfinite) { + plant(self.x, r, w, true); + plant(self.x, r, w, false); + } + if (fill == .stress) { + // rows alternate 255 and 254, taps 255 and -256: x k0 + p k1 = 1, which fp16 and bf16 lose + @memset(self.base, 0); + for (0..r) |i| { + @memset(self.x[i * w ..][0..w], word(if (i % 2 == 0) 255 else 254)); + @memset(self.res[i * w ..][0..w], word(0.25)); + for (0..2) |b| for (0..2) |t| { + const v: f32 = if ((i % 2 == 0) == (t == 0)) 255 else -256; + @memset(self.dyn[((i * 2 + b) * 2 + t) * (w / gs) ..][0 .. w / gs], word(v)); + }; + } + } + for ([_][]u16{ self.x, self.dyn, self.base, self.res }, 0..) |v, i| self.in[i] = try self.rig.input(std.mem.sliceAsBytes(v)); + self.out = try self.rig.output(r * w * 2); + }, + } + return true; + } + + fn branch(self: *const Act) usize { + return if (self.which == .dconv_res) 1 else 0; + } + + pub fn run(self: *Act) !void { + const s = self.rig.s; + const r = self.shape.rows; + const w = self.width(); + switch (self.which) { + .swiglu, .swiglu_sums => try self.k.swiglu(s, self.in[0], 2 * w, self.in[0] + 2 * w, 2 * w, self.out, self.xs, r, w), + .gate, .gate_sums => try self.k.outputGate(s, self.in[1], self.in[0] + 2 * self.shape.k, 2 * w, 2 * self.shape.k, self.out, self.xs, r, self.shape.n, self.shape.k), + .group_sums => try self.k.groupSums(s, self.in[0], w + 64, self.xs, r, w), + .dconv, .dconv_res => try self.k.dconv(s, self.in[0], self.in[1], self.in[2], if (self.which == .dconv_res) self.in[3] else null, self.out, r, w, gs, self.branch(), self.shape.k), + } + try s.synchronize(); + self.ran = true; + } + + pub fn count(self: *Act) usize { + const n = self.shape.rows * self.width(); + return if (self.sums()) n / 64 else n; + } + + /// The values a group sums: x, or the bf16 outputs stored (before a run, gate 4's: silu(17) and sigmoid(17) exact). + fn summed(self: *Act) ![]const u16 { + if (self.which == .group_sums) return self.x; + const n = self.shape.rows * self.width(); + const v = try self.arena.allocator().alloc(u16, n); + if (self.ran) { + @memcpy(std.mem.sliceAsBytes(v), try self.rig.bytes(self.out, n * 2)); + return v; + } + if (self.fill != .stress) return error.NoRunYet; + const k = if (self.which == .swiglu_sums) aref.swiglu(17, 1) else aref.sigmoid(17); + for (v, self.y) |*o, y| o.* = word(@floatCast(k * bf(y))); + return v; + } + + fn group(self: *Act, vals: []const u16, o: usize, t: *[64]f64) void { + const w = self.width(); + const ld = if (self.which == .group_sums) w + 64 else w; + const row = o / (w / 64); + const j = o % (w / 64); + for (t, 0..) |*v, e| v.* = bf(vals[row * ld + 64 * j + e]); + } + + /// The conv's two terms at output o, and the residual (0 without). + fn conv(self: *const Act, o: usize, t: *[2]f64) f64 { + const w = self.shape.n; + const row = o / w; + const i = o % w; + const b = self.branch(); + const grp = i / gs; + const tap = struct { + fn f(base: f64, dyn: f64) f64 { + return bf(word(@as(f32, @floatCast(base)) + @as(f32, @floatCast(dyn)))); // k = bf16(base + dyn) in fp32 + } + }.f; + const k0 = tap(bf(self.base[(b * 2) * w + i]), bf(self.dyn[((row * 2 + b) * 2) * (w / gs) + grp])); + const k1 = tap(bf(self.base[(b * 2 + 1) * w + i]), bf(self.dyn[((row * 2 + b) * 2 + 1) * (w / gs) + grp])); + const prev = if (row % self.shape.k > 0) bf(self.x[(row - 1) * w + i]) else 0; + t.* = .{ bf(self.x[row * w + i]) * k0, prev * k1 }; + return if (self.which == .dconv_res) bf(self.res[row * w + i]) else 0; + } + + pub fn reference(self: *Act, want: []f64, mag: []f64) !void { + const w = self.width(); + if (self.sums()) { + const vals = try self.summed(); + var t: [64]f64 = undefined; + for (want, mag, 0..) |*o, *m, i| { + self.group(vals, i, &t); + o.* = 0; + m.* = 0; + for (t) |v| { + o.* += v; + m.* += @abs(v); + } + } + return; + } + for (want, mag, 0..) |*o, *m, i| { + const row = i / w; + const e = i % w; + switch (self.which) { + .swiglu => o.* = aref.swiglu(bf(self.x[row * 2 * w + e]), bf(self.y[i])), + .gate => o.* = bf(self.y[i]) * aref.sigmoid(bf(self.x[row * 2 * w + (e / self.shape.k) * 2 * self.shape.k + self.shape.k + e % self.shape.k])), + .dconv, .dconv_res => { + var t: [2]f64 = undefined; + const res = self.conv(i, &t); + o.* = res + aref.dconv(1, 1, t[0], t[1]); + m.* = @abs(t[0]) + @abs(t[1]); + continue; + }, + else => unreachable, + } + m.* = @abs(o.*); + } + } + + pub fn simulate(self: *Act, acc: c.Accum, order: g.Order, out: []f64) !void { + if (self.sums()) { + const vals = try self.summed(); + var t: [64]f64 = undefined; + for (out, 0..) |*o, i| { + self.group(vals, i, &t); + o.* = rows.sumIn(acc, order, &t); + } + return; + } + if (self.which != .dconv and self.which != .dconv_res) return error.NoSums; + var t: [2]f64 = undefined; + for (out, 0..) |*o, i| { + const res = self.conv(i, &t); + o.* = res + rows.sumIn(acc, order, &t); + } + } + + /// core/act_ref's dconv bound with the residual: half a bf16 ulp of the output and of the conv, 2^-22 of its terms. + fn checkRes(ctx: *anyopaque, at: g.Gate) anyerror!?[]const u8 { + const self: *Act = @ptrCast(@alignCast(ctx)); + const out = try self.rig.outputBytes(); + const half = std.math.pow(f64, 2, -8); + var planted: usize = 0; + const n = self.count(); + var t: [2]f64 = undefined; + for (0..n) |i| { + const res = self.conv(i, &t); + const y = t[0] + t[1]; + const want = res + y; + const got = c.read(.bf16, out, i); + if (c.classOf(want) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(want)) return try std.fmt.bufPrint(&self.msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(got), c.classOf(want) }); + if (at == .nonfinite or c.classOf(want) != .finite) continue; + const tol = @abs(want) * half + @abs(y) * half + (@abs(t[0]) + @abs(t[1])) * std.math.pow(f64, 2, -22) + 1e-30; + if (!(@abs(got - want) <= tol)) return try std.fmt.bufPrint(&self.msg, "output {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of its bound", .{ i, got, want, @abs(got - want) / tol }); + } + if (at == .nonfinite and (planted == 0 or planted == n)) return try std.fmt.bufPrint(&self.msg, "the planted values make {d} of {d} outputs non-finite: no proof", .{ planted, n }); + return null; + } +}; diff --git a/zig/tests/cuda/contract_gdn.zig b/zig/tests/cuda/contract_gdn.zig new file mode 100644 index 0000000000..7c9653463e --- /dev/null +++ b/zig/tests/cuda/contract_gdn.zig @@ -0,0 +1,305 @@ +//! gdn_prep.cu (#797) under the contract on CUDA: a row's conv window, SiLU, the q/k norms, v, g and beta. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const ref = core.gdn_prep_ref; +const gp = cuda.gdn_prep; + +/// Which outputs an entry reads (the others are working memory): q and k, v, or g and beta. +pub const Part = enum { qk, v, gb }; + +/// The conv's taps (one fma chain). +fn taps(s: c.Shape) usize { + return s.k; +} + +fn noSums(_: c.Shape) usize { + return 0; +} + +const qk_why = "q and k are normed rows: a bf16 conv or square sum moves them by about an output ulp (fp16 overflows)"; +const v_why = "v is the conv's bf16 rounding: a bf16 sum of its few taps is within an output ulp or two (fp16 overflows)"; +const gb_rule = "own check, the PR's float64 bound: g = exp(-y) carries y times y's relative error, beyond a fixed op count"; + +/// Shape: rows x key heads (n; 3 v heads a key head) x conv taps (k). +pub fn spec(dk: usize, part: Part) c.Spec { + const served: []const c.Shape = switch (dk) { + 128 => &.{ .{ .rows = 1, .n = 16, .k = 4 }, .{ .rows = 64, .n = 16, .k = 4 } }, // Qwen3.8-27B: 16 key heads, 4 taps + 64 => &.{ .{ .rows = 1, .n = 2, .k = 4 }, .{ .rows = 40, .n = 2, .k = 4 } }, + else => &.{ .{ .rows = 1, .n = 1, .k = 2 }, .{ .rows = 40, .n = 1, .k = 2 } }, + }; + const tiles: [3]c.Tile = .{ .{ .size = 16 }, .{ .size = 4 }, .{ .size = 8, .min = 2 } }; // small blocks to 16 rows; 4 heads a block + const name: []const u8 = switch (part) { + inline else => |p| switch (dk) { + 64 => "cuda gdn_prep tf_gdn_prep_64 " ++ @tagName(p), + 128 => "cuda gdn_prep tf_gdn_prep_128 " ++ @tagName(p), + else => "cuda gdn_prep tf_gdn_prep_256 " ++ @tagName(p), + }, + }; + return switch (part) { + // silu 8, the norm's sum and rsqrt 12, the scale and products 4; c rounded to bf16, and the norm through it + .qk => .{ .name = name, .accum = .f32, .depth = taps, .ops = 24, .mids = 2, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .bf16 = 0, .why = qk_why } }, + .v => .{ .name = name, .accum = .f32, .depth = taps, .ops = 8, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .bf16 = 0, .why = v_why } }, + .gb => .{ .name = name, .accum = .none, .depth = noSums, .out = .f32, .nan = .{ .stated = gb_rule }, .tiles = tiles, .served = served }, + }; +} + +const vh_per = 3; + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +pub const Gdn = struct { + rig: rig_.Rig, + dk: usize, + part: Part, + k: *const gp.GdnPrep, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + sh: gp.Shape = .{ .kh = 0, .vh = 0, .dk = 0, .keep = 0 }, + qkv: []u16 = &.{}, + state: []u16 = &.{}, + conv: []u16 = &.{}, + av: []u16 = &.{}, + bv: []u16 = &.{}, + a_log: []f32 = &.{}, + dtb: []f32 = &.{}, + sid: []i32 = &.{}, + win: []i32 = &.{}, + io: gp.Io = undefined, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Gdn) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Gdn) g.Kernel { + return .{ .spec = spec(self.dk, self.part), .ctx = self, .vt = if (self.part == .gb) &vt_gb else &vt }; + } + + const vt = g.hooks(Gdn); + const vt_gb = blk: { + var v = g.hooks(Gdn); + v.check = checkGb; + break :blk v; + }; + + fn chans(self: *const Gdn) usize { + return (2 * self.sh.kh + self.sh.vh) * self.dk; + } + + fn ld(self: *const Gdn) usize { + return self.chans() + 64; + } + + pub fn prepare(self: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + if (shape.rows == 0 or shape.n == 0 or shape.k < 2 or shape.k > 8) return false; + self.shape = shape; + self.sh = .{ .kh = shape.n, .vh = vh_per * shape.n, .dk = self.dk, .keep = shape.k - 1 }; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const cc = self.chans(); + const keep = self.sh.keep; + const vh = self.sh.vh; + self.qkv = try a.alloc(u16, r * self.ld()); + self.state = try a.alloc(u16, 2 * keep * cc); + self.conv = try a.alloc(u16, cc * shape.k); + self.av = try a.alloc(u16, r * (vh + 3)); + self.bv = try a.alloc(u16, r * (vh + 3)); + for ([_][]u16{ self.qkv, self.state, self.conv, self.av, self.bv }, [_]f32{ 1, 1, 0.4, 2, 2 }) |v, sc| for (v) |*x| { + x.* = c.toBf16(rng.floatNorm(f32) * sc); + }; + self.a_log = try a.alloc(f32, vh); + self.dtb = try a.alloc(f32, vh); + for (self.a_log, self.dtb) |*x, *y| { + x.* = rng.float(f32) * 2 - 1; + y.* = rng.floatNorm(f32); + } + self.sid = try a.alloc(i32, r); + self.win = try a.alloc(i32, r * shape.k); + for (self.sid, 0..) |*x, i| x.* = @intCast(i % 2); + for (self.win, 0..) |*w, i| { // kept rows or any qkv row (a tree's window is any rows), the row itself last + w.* = if (rng.uintLessThan(u32, 3) == 0) @intCast(rng.uintLessThan(usize, keep)) else @intCast(keep + rng.uintLessThan(usize, r)); + if (i % shape.k == shape.k - 1) w.* = @intCast(keep + i / shape.k); + } + if (fill == .nonfinite) { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + const at = [_][2]usize{ .{ 0, cc / 3 }, .{ 1, cc / 2 }, .{ 2, cc - 1 }, .{ 3, 0 }, .{ 3, cc - 1 } }; + for (at, put) |p, v| if (p[0] < r) { + self.qkv[p[0] * self.ld() + p[1]] = c.toBf16(v); + }; + for ([_][]u16{ self.av, self.bv }, 0..) |v, i| if (i < r) { + v[i * (vh + 3) + i % vh] = c.toBf16(put[i]); + }; + } + if (fill == .stress) { + // every input 256, taps 255 but the last one 250 to 254: big integer sums, fp16 overflows past 65504 + @memset(self.qkv, c.toBf16(256)); + @memset(self.state, c.toBf16(256)); + for (0..cc) |i| for (0..shape.k) |t| { + self.conv[i * shape.k + t] = c.toBf16(if (t + 1 < shape.k) 255 else @floatFromInt(250 + i % 5)); + }; + } + self.io = .{ .qkv = try self.rig.input(std.mem.sliceAsBytes(self.qkv)), .qkv_ld = @intCast(self.ld()), .state = try self.rig.input(std.mem.sliceAsBytes(self.state)), .conv = try self.rig.input(std.mem.sliceAsBytes(self.conv)), .win = try self.rig.input(std.mem.sliceAsBytes(self.win)), .sid = try self.rig.input(std.mem.sliceAsBytes(self.sid)), .a = try self.rig.input(std.mem.sliceAsBytes(self.av)), .a_ld = @intCast(vh + 3), .b = try self.rig.input(std.mem.sliceAsBytes(self.bv)), .b_ld = @intCast(vh + 3), .a_log = try self.rig.input(std.mem.sliceAsBytes(self.a_log)), .dt_bias = try self.rig.input(std.mem.sliceAsBytes(self.dtb)), .q = 0, .k = 0, .v = 0, .g = 0, .beta = 0 }; + const qk_bytes = r * self.sh.kh * self.dk * 2; + const kinds = [5]bool{ self.part == .qk, self.part == .qk, self.part == .v, self.part == .gb, self.part == .gb }; + const sizes = [5]usize{ qk_bytes, qk_bytes, r * vh * self.dk * 2, r * vh * 4, r * vh * 4 }; + var ptrs: [5]u64 = undefined; + for (kinds, sizes, &ptrs) |out, n, *p| p.* = if (out) try self.rig.output(n) else try self.rig.scratch(n); + self.io.q = ptrs[0]; + self.io.k = ptrs[1]; + self.io.v = ptrs[2]; + self.io.g = ptrs[3]; + self.io.beta = ptrs[4]; + return true; + } + + pub fn run(self: *Gdn) !void { + try self.k.run(self.rig.s, self.io, self.shape.rows, self.sh); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Gdn) usize { + const r = self.shape.rows; + return switch (self.part) { + .qk => 2 * r * self.sh.kh * self.dk, + .v => r * self.sh.vh * self.dk, + .gb => 2 * r * self.sh.vh, + }; + } + + /// Output o's row, head (q heads, k heads, then v heads) and channel. + fn locate(self: *const Gdn, o: usize) [3]usize { + const r = self.shape.rows; + const kh = self.sh.kh; + const dk = self.dk; + if (self.part == .v) return .{ o / (self.sh.vh * dk), 2 * kh + o / dk % self.sh.vh, o % dk }; + const k_half = o >= r * kh * dk; + const i = if (k_half) o - r * kh * dk else o; + return .{ i / (kh * dk), (if (k_half) kh else 0) + i / dk % kh, i % dk }; + } + + /// Head hd of row r: each channel's tap terms into t [dk][taps]. + fn terms(self: *const Gdn, r: usize, hd: usize, t: []f64) void { + const n = self.shape.k; + const keep = self.sh.keep; + for (0..n) |tp| { + const src: usize = @intCast(self.win[r * n + tp]); + for (0..self.dk) |i| { + const ch = hd * self.dk + i; + const x = if (src < keep) bf(self.state[(@as(usize, @intCast(self.sid[r])) * keep + src) * self.chans() + ch]) else bf(self.qkv[(src - keep) * self.ld() + ch]); + t[i * n + tp] = x * bf(self.conv[ch * n + tp]); + } + } + } + + /// Head hd of row r's outputs and magnitudes (dk each), its conv sums in `acc` and `order` (f32 none: float64). + fn head(self: *Gdn, r: usize, hd: usize, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const a = self.arena.allocator(); + const n = self.shape.k; + const t = try a.alloc(f64, self.dk * n); + const cv = try a.alloc(f64, self.dk); + self.terms(r, hd, t); + for (cv, mag, 0..) |*o, *m, i| { + const row = t[i * n ..][0..n]; + var s: f64 = 0; + m.* = 0; + for (row) |v| { + s += v; + m.* += 1.1 * @abs(v); // silu' is at most 1.1 + } + if (acc) |ac| s = rows.sumIn(ac, order, row); + o.* = s / (1 + @exp(-s)); + } + if (hd >= 2 * self.sh.kh) { + @memcpy(want, cv); + return; + } + const is_q = hd < self.sh.kh; + const dkf: f64 = @floatFromInt(self.dk); + var f = ref.norm(cv, is_q, self.sh.eps, want); + if (acc) |ac| { + const sq = try a.alloc(f64, self.dk); + for (sq, cv) |*q, v| q.* = v * v; + f = 1 / @sqrt(rows.sumIn(ac, order, sq) / dkf + self.sh.eps) * (if (is_q) 1 / dkf else 1 / @sqrt(dkf)); + for (want, cv) |*o, v| o.* = v * f; + } + for (mag) |*m| m.* *= @abs(f); + } + + fn outputs(self: *Gdn, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const a = self.arena.allocator(); + const hw = try a.alloc(f64, self.dk); + const hm = try a.alloc(f64, self.dk); + var o: usize = 0; + while (o < want.len) : (o += self.dk) { + const at = self.locate(o); + try self.head(at[0], at[1], acc, order, hw, hm); + @memcpy(want[o..][0..self.dk], hw); + @memcpy(mag[o..][0..self.dk], hm); + } + } + + pub fn reference(self: *Gdn, want: []f64, mag: []f64) !void { + if (self.part == .gb) { + for (want, mag, 0..) |*w, *m, i| { + w.* = self.gbWant(i); + m.* = @abs(w.*); + } + return; + } + return self.outputs(null, .sequential, want, mag); + } + + pub fn simulate(self: *Gdn, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.arena.allocator().alloc(f64, out.len); + return self.outputs(acc, order, out, mag); + } + + /// g (first rows x vh) then beta, in float64. + fn gbWant(self: *const Gdn, i: usize) f64 { + const vh = self.sh.vh; + const half = self.shape.rows * vh; + const j = i % half; + const r = j / vh; + const hv = j % vh; + if (i < half) return ref.gate(bf(self.av[r * (vh + 3) + hv]), self.dtb[hv], self.a_log[hv]); + return 1 / (1 + @exp(-bf(self.bv[r * (vh + 3) + hv]))); + } + + /// gdn_prep_tests' bounds: g within 2^-20 (4 + 4 exp(a_log) softplus) of itself, beta within 2^-19. + fn checkGb(ctx: *anyopaque, at: g.Gate) anyerror!?[]const u8 { + const self: *Gdn = @ptrCast(@alignCast(ctx)); + const out = try self.rig.outputBytes(); + const vh = self.sh.vh; + const half = self.shape.rows * vh; + var planted: usize = 0; + const n = self.count(); + for (0..n) |i| { + const want = self.gbWant(i); + const got = c.read(.f32, out, i); + if (c.classOf(want) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(want)) return try std.fmt.bufPrint(&self.msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(got), c.classOf(want) }); + if (at == .nonfinite or c.classOf(want) != .finite) continue; + const j = i % half; + const x = bf(self.av[j / vh * (vh + 3) + j % vh]) + self.dtb[j % vh]; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + const tol = if (i < half) @abs(want) * std.math.pow(f64, 2, -20) * (4 + 4 * @exp(@as(f64, self.a_log[j % vh])) * sp) + 1e-30 else want * std.math.pow(f64, 2, -19) + 1e-30; + if (!(@abs(got - want) <= tol)) return try std.fmt.bufPrint(&self.msg, "output {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of its bound", .{ i, got, want, @abs(got - want) / tol }); + } + if (at == .nonfinite and (planted == 0 or planted == n)) return try std.fmt.bufPrint(&self.msg, "the planted values make {d} of {d} outputs non-finite: no proof", .{ planted, n }); + return null; + } +}; diff --git a/zig/tests/cuda/contract_qk.zig b/zig/tests/cuda/contract_qk.zig new file mode 100644 index 0000000000..137d24a2f4 --- /dev/null +++ b/zig/tests/cuda/contract_qk.zig @@ -0,0 +1,271 @@ +//! qk_rows.cu (#797) under the contract on CUDA: QK norm, rotate-half RoPE and v, against core/qk_ref.zig. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const qk = cuda.qk_rows; + +pub const Norm = enum { none, f32w, bf16w }; +pub const Rope = enum { none, table, inv }; + +/// One registered form: head dim, query heads a key head, rotated dims, norm weights and where cos and sin come from. +pub const Form = struct { name: []const u8, dim: usize, group: usize, rot: usize, norm: Norm, rope: Rope, h: usize }; + +/// Kolibri (48/4 x 128, sliding and full), a warp a head without the norm, the 27B (24/4 x 256, rotary 64), dim 64. +pub const forms = [_]Form{ + .{ .name = "cuda qk_rows tf_qk_rows_w (norm, rotation from a table)", .dim = 128, .group = 12, .rot = 128, .norm = .f32w, .rope = .table, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (norm, no rotation)", .dim = 128, .group = 12, .rot = 0, .norm = .f32w, .rope = .none, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (bf16 norm, rotation from inv_freq)", .dim = 128, .group = 12, .rot = 128, .norm = .bf16w, .rope = .inv, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (no norm, rotation from a table)", .dim = 128, .group = 8, .rot = 128, .norm = .none, .rope = .table, .h = 16 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 256, norm, rotary 64 from a table)", .dim = 256, .group = 6, .rot = 64, .norm = .f32w, .rope = .table, .h = 24 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 256, bf16 norm, rotary 64 from inv_freq)", .dim = 256, .group = 6, .rot = 64, .norm = .bf16w, .rope = .inv, .h = 24 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 64, no norm, rotary 32 from a table)", .dim = 64, .group = 4, .rot = 32, .norm = .none, .rope = .table, .h = 8 }, +}; + +/// A rotated value: x rstd, times w, times cos (cosf and sinf: 2 ulp, 4 units more), then the add; unrotated fewer. +fn depthTable(_: c.Shape) usize { + return 4; +} + +fn depthInv(_: c.Shape) usize { + return 8; +} + +fn depthBare(_: c.Shape) usize { + return 2; +} + +const norm_why = "a normed row follows its sum of squares as ss^-1/2: a bf16 sum moves it by about an output ulp (fp16 overflows)"; + +/// Shape: rows x query heads (n) x head dim (k). +pub fn spec(f: Form) c.Spec { + const served = switch (f.dim) { + 128 => if (f.h == 48) &[_]c.Shape{ .{ .rows = 1, .n = 48, .k = 128 }, .{ .rows = 300, .n = 48, .k = 128 } } else &[_]c.Shape{ .{ .rows = 1, .n = 16, .k = 128 }, .{ .rows = 300, .n = 16, .k = 128 } }, + 256 => &[_]c.Shape{ .{ .rows = 1, .n = 24, .k = 256 }, .{ .rows = 300, .n = 24, .k = 256 } }, + else => &[_]c.Shape{ .{ .rows = 1, .n = 8, .k = 64 }, .{ .rows = 300, .n = 8, .k = 64 } }, + }; + const depth: *const fn (c.Shape) usize = if (f.norm == .none) &depthBare else if (f.rope == .inv) &depthInv else &depthTable; + // the norm's sum of squares (14 deep, half of it in rsqrt), the division, eps, the root and the reciprocal + const ops: f64 = if (f.norm == .none) 0 else 12; + const sums: c.Sums = if (f.norm == .none) .{} else .{ .bf16 = 0, .why = norm_why }; + return .{ .name = f.name, .accum = .f32, .depth = depth, .ops = ops, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = served, .sums = sums }; +} + +const positions = 4096; +const len = 600; + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +pub const Qk = struct { + rig: rig_.Rig, + form: Form, + k: *const qk.QkRows, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + w: []f64 = &.{}, + cos: []f32 = &.{}, + sin: []f32 = &.{}, + inv: []f32 = &.{}, + pos: []i32 = &.{}, + dx: u64 = 0, + sh: qk.Shape = .{}, + dpos: u64 = 0, + dslot: u64 = 0, + out: qk.Out = undefined, + + pub fn deinit(self: *Qk) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Qk) g.Kernel { + return .{ .spec = spec(self.form), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Qk); + + fn hk(self: *const Qk) usize { + return self.shape.n / self.form.group; + } + + fn heads(self: *const Qk) usize { + return self.shape.n + 2 * self.hk(); + } + + pub fn prepare(self: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const f = self.form; + if (shape.rows == 0 or shape.k != f.dim or shape.n % f.group != 0 or shape.n == 0) return false; + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const width = self.heads() * f.dim; + self.x = try a.alloc(u16, r * width); + for (self.x, 0..) |*h, i| h.* = c.toBf16(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const rot = @max(f.rot, 2); + self.cos = try a.alloc(f32, positions * rot); + self.sin = try a.alloc(f32, positions * rot); + for (0..positions) |p| for (0..rot) |i| { + const ang = @as(f64, @floatFromInt(p)) / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * (i % (rot / 2)))) / @as(f64, @floatFromInt(rot))); + self.cos[p * rot + i] = @floatCast(@cos(ang)); + self.sin[p * rot + i] = @floatCast(@sin(ang)); + }; + self.inv = try a.alloc(f32, rot / 2); + for (self.inv, 0..) |*v, j| v.* = @floatCast(1 / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(rot)))); + const wf = try a.alloc(f32, 2 * f.dim); + for (wf) |*v| v.* = 0.5 + rng.float(f32); + if (fill == .nonfinite) { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + const at = [_][2]usize{ .{ 0, width / 3 }, .{ 1, width / 2 }, .{ 2, width - 1 }, .{ 3, 0 }, .{ 3, f.dim - 1 } }; + for (at, put) |p, v| if (p[0] < r) { + self.x[p[0] * width + p[1]] = c.toBf16(v); + }; + } + if (fill == .stress) { + if (f.norm == .none) { + // x c - x' s with c, s = 65535, 65534, x = 255: products near 2^24, exact in fp32, a difference of 255 + @memset(self.x, c.toBf16(255)); + @memset(self.cos, 65535); + @memset(self.sin, 65534); + } else { + // integers 32 to 64 of either sign: squares sum past fp16's range, exactly in fp32; cos 1, sin 0 + for (self.x) |*h| h.* = c.toBf16(@as(f32, @floatFromInt(32 + rng.uintLessThan(u32, 33))) * (if (rng.boolean()) @as(f32, 1) else -1)); + @memset(self.cos, 1); + @memset(self.sin, 0); + @memset(self.inv, 0); + } + } + self.w = try a.alloc(f64, 2 * f.dim); + const wb = try a.alloc(u16, 2 * f.dim); + for (self.w, wf, wb) |*v, x, *b| { + b.* = c.toBf16(x); + v.* = if (f.norm == .bf16w) bf(b.*) else x; + } + self.pos = try a.alloc(i32, r); + const slot = try a.alloc(i32, r); + for (self.pos, slot, 0..) |*p, *sl, i| { // distinct (slot, pos % len) a row: no row's cache entry overwrites another's + p.* = @intCast((i / 2) % len + len * (i % 5)); + sl.* = @intCast(i % 2); + } + self.dx = try self.rig.input(std.mem.sliceAsBytes(self.x)); + const nw: ?u64 = switch (f.norm) { + .none => null, + .f32w => try self.rig.input(std.mem.sliceAsBytes(wf)), + .bf16w => try self.rig.input(std.mem.sliceAsBytes(wb)), + }; + const tc = try self.rig.input(std.mem.sliceAsBytes(self.cos)); + const ts = try self.rig.input(std.mem.sliceAsBytes(self.sin)); + const ti = try self.rig.input(std.mem.sliceAsBytes(self.inv)); + self.sh = .{ .norms = nw, .norms_bf16 = f.norm == .bf16w, .cos = if (f.rope == .table) tc else null, .sin = if (f.rope == .table) ts else null, .inv = if (f.rope == .inv) ti else null, .rot = f.rot }; + self.dpos = try self.rig.input(std.mem.sliceAsBytes(self.pos)); + self.dslot = try self.rig.input(std.mem.sliceAsBytes(slot)); + const kv = r * self.hk() * f.dim * 2; + const cache = 2 * len * self.hk() * f.dim * 2; + self.out = .{ .q = try self.rig.output(r * shape.n * f.dim * 2), .k = try self.rig.output(kv), .v = try self.rig.output(kv), .kc = try self.rig.scratch(cache), .vc = try self.rig.scratch(cache), .len = len }; + return true; + } + + pub fn run(self: *Qk) !void { + try self.k.rows(self.rig.s, self.dx, .{ .h = self.shape.n, .hk = self.hk(), .dim = self.form.dim }, self.sh, self.dpos, self.dslot, self.out, self.shape.rows); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Qk) usize { + return self.shape.rows * self.heads() * self.form.dim; + } + + /// Output o's row and head (q heads, k heads, v heads as the outputs lie: all q, then all k, then all v). + fn locate(self: *const Qk, o: usize) [2]usize { + const r = self.shape.rows; + const h = self.shape.n; + const d = self.form.dim; + const i = o / d; + if (i < r * h) return .{ i / h, i % h }; + const j = i - r * h; + const hk_ = self.hk(); + if (j < r * hk_) return .{ j / hk_, h + j % hk_ }; + return .{ (j - r * hk_) / hk_, h + hk_ + (j - r * hk_) % hk_ }; + } + + /// One head's outputs and term magnitudes; its sums in `acc` and `order` (null: float64, core/qk_ref.head). + fn head(self: *Qk, r: usize, hd: usize, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const f = self.form; + const d = f.dim; + const a = self.arena.allocator(); + const x = try a.alloc(f64, d); + for (x, 0..) |*v, i| v.* = bf(self.x[(r * self.heads() + hd) * d + i]); + if (hd >= self.shape.n + self.hk()) { + @memcpy(want, x); + for (mag, x) |*m, v| m.* = @abs(v); + return; + } + const w: ?[]const f64 = if (f.norm == .none) null else self.w[(if (hd >= self.shape.n) d else 0)..][0..d]; + const rot = @max(f.rot, 2); + const p: usize = @intCast(self.pos[r]); + const cv = try a.alloc(f64, rot); + const sv = try a.alloc(f64, rot); + for (0..rot) |i| { + if (f.rope == .inv) { + const ang: f64 = @as(f32, @floatFromInt(p)) * self.inv[i % (rot / 2)]; // the angle in fp32, as served + cv[i] = @cos(ang); + sv[i] = @sin(ang); + } else { + cv[i] = self.cos[p * rot + i]; + sv[i] = self.sin[p * rot + i]; + } + } + const roped = f.rope != .none; + var ss: f64 = 0; + const sq = try a.alloc(f64, d); + for (sq, x) |*q, v| { + q.* = v * v; + ss += q.*; + } + if (acc) |ac| ss = rows.sumIn(ac, order, sq); + const inv = if (w != null) 1 / @sqrt(ss / @as(f64, @floatFromInt(d)) + 1e-6) else 1; + const z = try a.alloc(f64, d); + for (z, x, 0..) |*o, v, i| o.* = v * inv * (if (w) |ww| ww[i] else 1); + if (acc == null) core.qk_ref.head(x, w, 1e-6, if (roped) cv else null, if (roped) sv else null, f.rot, want); + const half = f.rot / 2; + for (want, mag, 0..) |*o, *m, i| { + if (!roped or i >= f.rot) { + if (acc != null) o.* = z[i]; + m.* = @abs(z[i]); + continue; + } + const pi = if (i < half) i + half else i - half; + const t = [2]f64{ z[i] * cv[i], (if (i < half) -z[pi] else z[pi]) * sv[i] }; + if (acc) |ac| o.* = rows.sumIn(ac, order, &t); + m.* = @abs(t[0]) + @abs(t[1]); + } + } + + fn outputs(self: *Qk, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const d = self.form.dim; + var o: usize = 0; + while (o < want.len) : (o += d) { + const at = self.locate(o); + try self.head(at[0], at[1], acc, order, want[o..][0..d], mag[o..][0..d]); + } + } + + pub fn reference(self: *Qk, want: []f64, mag: []f64) !void { + return self.outputs(null, .sequential, want, mag); + } + + pub fn simulate(self: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.arena.allocator().alloc(f64, out.len); + return self.outputs(acc, order, out, mag); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e87a2e90e5..bfb343528c 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -8,6 +8,9 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const rank_k = @import("contract_rank.zig"); +const act_k = @import("contract_act.zig"); +const gdn_k = @import("contract_gdn.zig"); +const qk_k = @import("contract_qk.zig"); const contract = core.contract; const g = contract.gates; @@ -21,7 +24,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "act_rows", "gdn_prep", "qk_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -89,6 +92,24 @@ pub fn run(gpu: check.Gpu) !void { var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var act_rows = try cuda.act_rows.ActRows.load(gpu.d); + defer act_rows.unload(); + var acts: [std.enums.values(act_k.Which).len]act_k.Act = undefined; + for (std.enums.values(act_k.Which), &acts) |w, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .k = &act_rows, .arena = .init(a) }; + defer for (&acts) |*x| x.deinit(); + for (&acts) |*x| try kernels.append(a, x.kernel()); + var gdn_prep = try cuda.gdn_prep.GdnPrep.load(gpu.d); + defer gdn_prep.unload(); + var gdns: [9]gdn_k.Gdn = undefined; + for (&gdns, 0..) |*x, i| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .dk = @as(usize, 64) << @intCast(i / 3), .part = @fromBackingInt(@intCast(i % 3)), .k = &gdn_prep, .arena = .init(a) }; + defer for (&gdns) |*x| x.deinit(); + for (&gdns) |*x| try kernels.append(a, x.kernel()); + var qk_rows = try cuda.qk_rows.QkRows.load(gpu.d); + defer qk_rows.unload(); + var qks: [qk_k.forms.len]qk_k.Qk = undefined; + 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 t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From 90aebfcb1bcdfef21feaa4760a84d8b5282d99b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 14/41] cuda/tree_attention: each row's own max flag, so a NaN head leaves other rows' keys --- zig/kernels/cuda/tree_attention.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index eea5d2645d..4e50f442e9 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -190,10 +190,10 @@ __device__ __forceinline__ void tree_group( tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); } - const bool act = tm0 != -INFINITY; - const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; - const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; - const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; // each row's own: a NaN head moves only its row + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; float c0 = 0.0f, c1 = 0.0f; #pragma unroll for (int n = 0; n < NT8; ++n) { From fc48c41c42cae20882996ecb48a62ef9ff7e7055 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 15/41] cuda/block_attention: keys a row cannot see add nothing, so an infinite value there stays out of its sum A block whose output holds a NaN or infinity runs again with the values of the tiles some row cannot see all of taken out of the mma; each row adds back the ones it sees. On finite values the first pass is the kernel as it was, byte for byte. --- zig/kernels/cuda/attn_odd.cuh | 63 +++++++++++++++++++++++++++++ zig/kernels/cuda/block_attention.cu | 60 ++++++++++++++++++--------- 2 files changed, 105 insertions(+), 18 deletions(-) create mode 100644 zig/kernels/cuda/attn_odd.cuh diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh new file mode 100644 index 0000000000..596ee3a95e --- /dev/null +++ b/zig/kernels/cuda/attn_odd.cuh @@ -0,0 +1,63 @@ +// NaN and infinite values in a tensor-core attention tile's V: a key a row cannot see adds nothing to that row's sum. +#pragma once + +#include +#include +#include + +namespace attn_odd { + +// Each bf16 half of x that is NaN or an infinity (all-ones exponent) as +0, the finite halves as they are. +__device__ __forceinline__ uint32_t finite2(uint32_t x) { + return x & (((x & 0x7f80u) == 0x7f80u ? 0u : 0xffffu) | ((x & 0x7f800000u) == 0x7f800000u ? 0u : 0xffff0000u)); +} + +// Some NaN or infinite value in a warp's mma rows g (r0: a real row) and g + 8 (r1) of its fold. +template +__device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, bool r1) { + float a = 0.0f, b = 0.0f; // x - x is 0 for a finite x, NaN for NaN and the infinities +#pragma unroll + for (int i = 0; i < NB; ++i) { + a = __fadd_rn(a, __fadd_rn(__fsub_rn(o[i][0], o[i][0]), __fsub_rn(o[i][1], o[i][1]))); + b = __fadd_rn(b, __fadd_rn(__fsub_rn(o[i][2], o[i][2]), __fsub_rn(o[i][3], o[i][3]))); + } + return (r0 && a != a) || (r1 && b != b); +} + +// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. +template +__device__ __forceinline__ void clean(uint4* vs, int threads) { + for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) { + uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7)); + const uint4 x = *p; + *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w)); + } +} + +// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them. +template +__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0, + float l1) { + const int g = lane / 4, qd = lane % 4; + for (int k = k0; k < k1; ++k) { + const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN + if (!v0 && !v1) continue; + const uint32_t* v = reinterpret_cast(vrow(k)); +#pragma unroll + for (int i = 0; i < D / 8; ++i) { + const uint32_t x = v[4 * i + qd]; + if (finite2(x) == x) continue; + const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u); + if ((x & 0x7f80u) == 0x7f80u) { + if (v0) o[i][0] = __fadd_rn(o[i][0], lo); + if (v1) o[i][2] = __fadd_rn(o[i][2], lo); + } + if ((x & 0x7f800000u) == 0x7f800000u) { + if (v0) o[i][1] = __fadd_rn(o[i][1], hi); + if (v1) o[i][3] = __fadd_rn(o[i][3], hi); + } + } + } +} + +} // namespace attn_odd diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index e3cd06b121..f5b1e2b18a 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -7,6 +7,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_block { @@ -166,30 +167,53 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( } Fold f; const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) - for (int k0 = lo / KT * KT; k0 < s; k0 += KT) { + const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; + // the context's edge tiles, where some row cannot see some key: e0..e1 + const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT); + auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below) + for (int k0 = e0; k0 < s; k0 += KT) { + for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { + const int key = k0 + i / C; + const bool in = key < s; + cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + } + commit(); + wait_all(); + if (odd && k0 < e1) attn_odd::clean(vs, blockDim.x); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, k0, vis); + __syncthreads(); + } + }; + auto block_keys = [&] { // the block's own keys for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { - const int key = k0 + i / C; - const bool in = key < s; - cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); - cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + const int c = i / C; + const bool in = c < L; + const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; + cp16z(ks + sw(c, i % C), kb + at, in); + cp16z(vs + sw(c, i % C), vb + at, in); } commit(); wait_all(); + }; + const auto own = [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }; + context(false); + block_keys(); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, 0, own); + // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(f.o, g8 < L, g8 + 8 < L))) { + f = Fold(); + context(true); + attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1); + block_keys(); + if (causal) attn_odd::clean(vs, blockDim.x); // row r cannot see keys past r __syncthreads(); - f.tile(qa, ks, vs, scale, lane, k0, [&](int r, int key) { return r < L && key < s && s + r - key <= window; }); - __syncthreads(); - } - for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { // the block's own keys - const int c = i / C; - const bool in = c < L; - const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; - cp16z(ks + sw(c, i % C), kb + at, in); - cp16z(vs + sw(c, i % C), vb + at, in); + f.tile(qa, ks, vs, scale, lane, 0, own); + const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; + if (causal) attn_odd::add_back(f.o, row, 0, L, lane, own, f.l0, f.l1); } - commit(); - wait_all(); - __syncthreads(); - f.tile(qa, ks, vs, scale, lane, 0, [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }); if (g8 < L) #pragma unroll for (int i = 0; i < NB; ++i) From ee81000d8ace5925c4dca89cfb3d75eff92d8f1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 16/41] cuda/tree_attention, block_attention: registered with the kernel contract --- zig/src/core/contract/registry.zig | 15 ++ zig/tests/cuda/contract_attn.zig | 124 +++++++++ zig/tests/cuda/contract_block_attn.zig | 224 +++++++++++++++++ zig/tests/cuda/contract_tests.zig | 25 +- zig/tests/cuda/contract_tree_attn.zig | 333 +++++++++++++++++++++++++ 5 files changed, 720 insertions(+), 1 deletion(-) create mode 100644 zig/tests/cuda/contract_attn.zig create mode 100644 zig/tests/cuda/contract_block_attn.zig create mode 100644 zig/tests/cuda/contract_tree_attn.zig diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..b2cfbc0e18 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,21 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_paths", .symbol = "tf_tree_paths", .image = "tree_attention", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1s3", .symbol = "tf_tree_shared_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part", .symbol = "tf_tree_part", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1", .symbol = "tf_tree_part_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1s3", .symbol = "tf_tree_part_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256", .symbol = "tf_tree_shared_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1", .symbol = "tf_tree_shared_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1s3", .symbol = "tf_tree_shared_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256", .symbol = "tf_tree_part_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_attn.zig b/zig/tests/cuda/contract_attn.zig new file mode 100644 index 0000000000..cec73374a3 --- /dev/null +++ b/zig/tests/cuda/contract_attn.zig @@ -0,0 +1,124 @@ +//! Attention kernels under the contract on CUDA: their outputs against core/attention_ref.zig, head by head. +const std = @import("std"); +const core = @import("core"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const ref = core.attention_ref; + +/// Why gates 1 and 3 run attention_ref's bound: P enters the P x V mma as bf16, and the spec's bound takes fp32 terms. +pub const why = "P is rounded to bf16 before the P x V mma: gates 1 and 3 hold each output to attention_ref's float64 bound"; + +/// One output head's float64 inputs: its query, and its visible keys' K and V rows ([keys][d]). +pub const Head = struct { q: []f64, k: []f64, v: []f64, n: usize = 0 }; + +/// The float64 side of an attention test: `t.load(i, *Head)` fills output head i (outputs [heads][d] bf16). +pub const Ref = struct { + a: std.mem.Allocator, + d: usize, + head: Head, + keys: []usize, + p: []f64, + want: []f64, + tol: []f64, + msg: [240]u8 = undefined, + + pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref { + const keys = try a.alloc(usize, most); + for (keys, 0..) |*k, i| k.* = i; + return .{ .a = a, .d = d, .head = .{ .q = try a.alloc(f64, d), .k = try a.alloc(f64, most * d), .v = try a.alloc(f64, most * d) }, .keys = keys, .p = try a.alloc(f64, most), .want = try a.alloc(f64, d), .tol = try a.alloc(f64, d) }; + } + + pub fn deinit(r: *Ref) void { + for ([_][]f64{ r.head.q, r.head.k, r.head.v, r.p, r.want, r.tol }) |x| r.a.free(x); + r.a.free(r.keys); + } + + /// Gates 1 and 3: every output's class equals float64's, and a finite one is within attention_ref's bound. + pub fn check(r: *Ref, t: anytype, heads: usize, scale: f64, out: []const u8, at: g.Gate) !?[]const u8 { + var planted: usize = 0; + for (0..heads) |i| { + try t.load(i, &r.head); + ref.head(r.head.q, r.head.k, r.head.v, r.keys[0..r.head.n], r.d, scale, r.want, r.tol); + for (r.want, r.tol, 0..) |w, tl, e| { + const got = c.read(.bf16, out, i * r.d + e); + if (c.classOf(w) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(w)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d}: {t} where float64 gives {t}", .{ i, e, c.classOf(got), c.classOf(w) }); + if (at != .nonfinite and std.math.isFinite(w) and !(@abs(got - w) <= tl)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of attention_ref's bound", .{ i, e, got, w, @abs(got - w) / tl }); + } + } + if (at == .nonfinite and (planted == 0 or planted == heads * r.d)) return try std.fmt.bufPrint(&r.msg, "the planted values make {d} of {d} reference outputs non-finite: no proof", .{ planted, heads * r.d }); + return null; + } + + fn weights(r: *Ref, scale: f64) f64 { + const h = r.head; + var m: f64 = -std.math.inf(f64); + for (r.p[0..h.n], 0..) |*p, j| { + var s: f64 = 0; + for (0..r.d) |e| s += h.q[e] * h.k[j * r.d + e]; + p.* = s * scale; + m = @max(m, p.*); + } + var l: f64 = 0; + for (r.p[0..h.n]) |*p| { + p.* = @exp(p.* - m); + l += p.*; + } + return l; + } + + /// Float64 outputs and each one's term magnitude, sum p_j |v_j| / sum p_j (gate 4 reads them). + pub fn reference(r: *Ref, t: anytype, heads: usize, scale: f64, want: []f64, mag: []f64) !void { + for (0..heads) |i| { + try t.load(i, &r.head); + const l = r.weights(scale); + for (0..r.d) |e| { + var s: f64 = 0; + var abs: f64 = 0; + for (r.p[0..r.head.n], 0..) |p, j| { + s += p * r.head.v[j * r.d + e]; + abs += p * @abs(r.head.v[j * r.d + e]); + } + want[i * r.d + e] = s / l; + mag[i * r.d + e] = abs / l; + } + } + } + + /// Gate 4's sums in an accumulator: its keys score alike, so every p is 1 and the sums are V's terms and l's ones. + pub fn simulate(r: *Ref, t: anytype, heads: usize, scale: f64, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try r.a.alloc(f64, r.p.len); + defer r.a.free(terms); + const ones = try r.a.alloc(f64, r.p.len); + defer r.a.free(ones); + @memset(ones, 1); + for (0..heads) |i| { + try t.load(i, &r.head); + _ = r.weights(scale); + for (r.p[0..r.head.n]) |p| if (p != 1) return error.KeysScoreApart; + const l = rows.sumIn(acc, order, ones[0..r.head.n]); + for (0..r.d) |e| { + for (terms[0..r.head.n], 0..) |*x, j| x.* = r.head.v[j * r.d + e]; + out[i * r.d + e] = rows.sumIn(acc, order, terms[0..r.head.n]) / l; + } + } + } +}; + +/// Gate 4's values: odd integers, every other one times 64 (bf16-exact), the second half cancelling the first. +pub 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; +} + +pub fn normal(rng: std.Random) u16 { + return c.toBf16(rng.floatNorm(f32)); +} diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig new file mode 100644 index 0000000000..155374b86d --- /dev/null +++ b/zig/tests/cuda/contract_block_attn.zig @@ -0,0 +1,224 @@ +//! block_attention.cu (#746) under the contract on CUDA: DFlash block rows over their windowed context, head by head. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const block = cuda.block_attention; + +const D = 128; +const h = 32; +const hk = 8; +const window = 2048; + +/// The P x V chain of a row: its window's keys and the block's, a 32-key tile at a time. +fn keyChain(s: c.Shape) usize { + return @min(s.k, window + 1) + 32; +} + +/// Shape: streams (rows), query heads (n) and stream 0's context keys (k); each stream's block of L rows. +pub fn spec(causal: bool) c.Spec { + // Qwen3.8-27B's DFlash2 drafter: blocks of 16 over 32 heads (8 key heads) and a 2,048-key window + return .{ + .name = if (causal) "cuda block attention tf_block_attention, causal blocks of 7" else "cuda block attention tf_block_attention (DFlash2 blocks of 16)", + .accum = .f32, + .depth = keyChain, + .ops = 1, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 1, .n = h, .k = 2100 }, .{ .rows = 4, .n = h, .k = 700 } }, + .sums = .{ .why = attn.why }, + }; +} + +pub const Block = struct { + rig: rig_.Rig, + ba: *const block.BlockAttention, + causal: bool, + r: ?attn.Ref = null, + ns: usize = 0, + lens: std.ArrayList(i32) = .empty, + ctx: std.ArrayList([]u16) = .empty, + q: std.ArrayList(u16) = .empty, + kb: std.ArrayList(u16) = .empty, + vb: std.ArrayList(u16) = .empty, + dev: [5]u64 = @splat(0), + out: u64 = 0, + + fn len(self: *const Block) usize { + return if (self.causal) 7 else 16; + } + + fn scale() f32 { + return @floatCast(1 / @sqrt(@as(f64, D))); + } + + pub fn deinit(self: *Block) void { + const a = self.rig.gpu.gpa; + self.drop(); + if (self.r) |*r| r.deinit(); + inline for (.{ "lens", "ctx", "q", "kb", "vb" }) |f| @field(self, f).deinit(a); + self.rig.deinit(); + } + + fn drop(self: *Block) void { + for (self.ctx.items) |x| self.rig.gpu.gpa.free(x); + self.ctx.clearRetainingCapacity(); + } + + pub fn kernel(self: *Block) g.Kernel { + return .{ .spec = spec(self.causal), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Block); + + pub fn prepare(self: *Block, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.drop(); + if (shape.n != h or shape.rows == 0 or shape.k == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const L = self.len(); + self.ns = shape.rows; + const rows = self.ns * L; + try self.lens.resize(a, self.ns); + for (self.lens.items, 0..) |*l, j| l.* = @intCast(@max(1, shape.k - j * shape.k / self.ns)); + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, D, @min(shape.k, window + 1) + L); + for (self.lens.items) |l| { + const x = try a.alloc(u16, 2 * hk * @as(usize, @intCast(l)) * D); // K [hk, len, D], then V + try self.ctx.append(a, x); + for (x) |*e| e.* = attn.normal(rng); + } + try self.q.resize(a, h * rows * D); + try self.kb.resize(a, hk * rows * D); + try self.vb.resize(a, hk * rows * D); + for ([_][]u16{ self.q.items, self.kb.items, self.vb.items }) |x| for (x) |*e| { + e.* = attn.normal(rng); + }; + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng), + } + const table = try a.alloc(i64, 2 * self.ns); + defer a.free(table); + for (self.ctx.items, 0..) |x, j| { + const base = try self.rig.input(std.mem.sliceAsBytes(x)); + table[2 * j] = @intCast(base); + table[2 * j + 1] = @intCast(base + x.len); // V follows K: half the u16s, all the bytes + } + for ([_][]const u8{ std.mem.sliceAsBytes(self.q.items), std.mem.sliceAsBytes(self.kb.items), std.mem.sliceAsBytes(self.vb.items), std.mem.sliceAsBytes(table), std.mem.sliceAsBytes(self.lens.items) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + self.out = try self.rig.output(rows * h * D * 2); + return true; + } + + fn ctxLen(self: *const Block, j: usize) usize { + return @intCast(self.lens.items[j]); + } + + /// NaN in a query, ±inf in values (context, block, a window's first key), both in one column, NaN in a key. + fn plant(self: *Block) void { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + const v0 = self.ctx.items[0][hk * s0 * D ..]; + self.q.items[3] = nan; // head 0, stream 0's row 0 + v0[(0 * s0 + s0 - 1) * D + 7] = inf; + self.vb.items[(1 * rows + 2) * D + 9] = ninf; // key head 1, stream 0's row 2 + if (s0 >= 3) { + v0[(2 * s0 + s0 - 2) * D + 11] = inf; + v0[(2 * s0 + s0 - 3) * D + 11] = ninf; + } + if (s0 > window) v0[(4 * s0 + s0 - window) * D + 13] = inf; // a key block row 0 alone sees + self.ctx.items[self.ns - 1][(3 * self.ctxLen(self.ns - 1) + self.ctxLen(self.ns - 1) - 1) * D] = nan; + } + + /// Gate 4: every key alike (each p is 1); values mirrored over the context keys every row sees, else 0. + fn stress(self: *Block, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const key = try a.alloc(u16, D); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + const L = self.len(); + @memset(self.vb.items, 0); + for (0..self.kb.items.len / D) |i| @memcpy(self.kb.items[i * D ..][0..D], key); + for (self.ctx.items, 0..) |x, j| { + const s = self.ctxLen(j); + for (0..hk * s) |i| @memcpy(x[i * D ..][0..D], key); + const v = x[hk * s * D ..]; + @memset(v, 0); + const lo = @min(s, (s + L) -| window); + const col = try a.alloc(f32, s - lo); + defer a.free(col); + for (0..hk) |kh| for (0..D) |e| { + attn.mirrored(rng, col); + for (col, lo..) |y, k| v[(kh * s + k) * D + e] = c.toBf16(y); + }; + } + } + + pub fn run(self: *Block) !void { + const L = self.len(); + try self.ba.run(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.dev[3], self.dev[4], self.out, .{ .h = h, .hk = hk, .rows = self.ns * L, .block = L, .window = window, .scale = scale(), .causal = self.causal }); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Block) usize { + return self.ns * self.len() * h * D; + } + + /// Output head i (block row i / h, head i % h): context keys k with len + r - k <= window, then the block's. + pub fn load(self: *Block, i: usize, hd: *attn.Head) !void { + const L = self.len(); + const rows = self.ns * L; + const row = i / h; + const head = i % h; + const kh = head / (h / hk); + const j = row / L; + const r = row % L; + const s = self.ctxLen(j); + const x = self.ctx.items[j]; + for (hd.q, 0..) |*y, e| y.* = c.fromBf16(self.q.items[(head * rows + row) * D + e]); + var n: usize = 0; + for (0..s) |k| if (s + r - k <= window) { + for (0..D) |e| { + hd.k[n * D + e] = c.fromBf16(x[(kh * s + k) * D + e]); + hd.v[n * D + e] = c.fromBf16(x[hk * s * D + (kh * s + k) * D + e]); + } + n += 1; + }; + for (0..L) |col| if (!self.causal or col <= r) { + for (0..D) |e| { + hd.k[n * D + e] = c.fromBf16(self.kb.items[(kh * rows + j * L + col) * D + e]); + hd.v[n * D + e] = c.fromBf16(self.vb.items[(kh * rows + j * L + col) * D + e]); + } + n += 1; + }; + hd.n = n; + } + + fn heads(self: *Block) usize { + return self.ns * self.len() * h; + } + + pub fn check(self: *Block, at: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), scale(), try self.rig.outputBytes(), at); + } + + pub fn reference(self: *Block, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), scale(), want, mag); + } + + pub fn simulate(self: *Block, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), scale(), acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e87a2e90e5..8114049a57 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -8,6 +8,8 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const rank_k = @import("contract_rank.zig"); +const tree_k = @import("contract_tree_attn.zig"); +const block_k = @import("contract_block_attn.zig"); const contract = core.contract; const g = contract.gates; @@ -21,7 +23,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "tree_attention", "block_attention" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -89,6 +91,27 @@ pub fn run(gpu: check.Gpu) !void { var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var tree_paths = try cuda.tree_attention.TreeAttention.load(gpu.d); + defer tree_paths.unload(); + var paths = tree_k.Paths{ .rig = .{ .gpu = gpu, .s = s }, .ta = &tree_paths }; + defer paths.deinit(); + try kernels.append(a, paths.kernel()); + var trees: [tree_k.variants.len]tree_k.Tree = undefined; + var ntrees: usize = 0; + defer for (trees[0..ntrees]) |*x| { + x.ta.unload(); + x.deinit(); + }; + for (tree_k.variants) |v| { + trees[ntrees] = .{ .rig = .{ .gpu = gpu, .s = s }, .v = v, .ta = try cuda.tree_attention.TreeAttention.load(gpu.d) }; + ntrees += 1; + } + for (&trees) |*x| try kernels.append(a, x.kernel()); + var block_attn = try cuda.block_attention.BlockAttention.load(gpu.d); + defer block_attn.unload(); + var blocks: [2]block_k.Block = .{ .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = false }, .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = true } }; + defer for (&blocks) |*x| x.deinit(); + for (&blocks) |*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_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig new file mode 100644 index 0000000000..3dd57f47e6 --- /dev/null +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -0,0 +1,333 @@ +//! tree_attention.cu (#746) under the contract on CUDA: paths byte for byte, attention head by head against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const tree = cuda.tree_attention; + +/// A tree per stream: its rows (at most 16, the 27B's tree), each a path through earlier rows of its stream. +const width = 16; + +/// The P x V chain of the longest row: every committed key, its path, a few combines. +fn keyChain(s: c.Shape) usize { + return s.k + tree.max_depth + 16; +} + +/// A tree attention call's blocks: four warps (sms 0), or one warp a slice with 2 (sms 1) or 3 tile stages. +pub const Variant = struct { d: usize, h: usize, hk: usize, sms: usize, name: []const u8 }; + +pub const variants = [_]Variant{ + .{ .d = 128, .h = 48, .hk = 4, .sms = 0, .name = "cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge)" }, + .{ .d = 128, .h = 48, .hk = 4, .sms = 1, .name = "cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge)" }, + .{ .d = 128, .h = 48, .hk = 4, .sms = 1 << 20, .name = "cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 0, .name = "cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 1, .name = "cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 1 << 20, .name = "cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3)" }, +}; + +/// Shape: drafted rows (trees of 16 over the streams), query heads (n) and stream 0's committed keys (k). +pub fn treeSpec(v: Variant) c.Spec { + // Kolibri (dim 128, 48 heads over 4) and Qwen3.8-27B (dim 256, 24 over 4): one stream's tree, then four streams' + const served: []const c.Shape = if (v.d == 128) &.{ .{ .rows = 16, .n = 48, .k = 2100 }, .{ .rows = 64, .n = 48, .k = 700 } } else &.{ .{ .rows = 16, .n = 24, .k = 2100 }, .{ .rows = 64, .n = 24, .k = 700 } }; + return .{ .name = v.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = width }, .{}, .{ .size = tree.span } }, .served = served, .sums = .{ .why = attn.why } }; +} + +pub const Tree = struct { + rig: rig_.Rig, + v: Variant, + ta: tree.TreeAttention, + r: ?attn.Ref = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + scale: f32 = 0, + ng: usize = 0, + committed: std.ArrayList(usize) = .empty, + row_stream: std.ArrayList(i32) = .empty, + paths: std.ArrayList(i32) = .empty, + depths: std.ArrayList(i32) = .empty, + offs: std.ArrayList(i64) = .empty, + q: std.ArrayList(u16) = .empty, + kn: std.ArrayList(u16) = .empty, + vn: std.ArrayList(u16) = .empty, + cache: std.ArrayList(u16) = .empty, + dev: [9]u64 = @splat(0), + sc: tree.Scratch = undefined, + out: u64 = 0, + + pub fn deinit(self: *Tree) void { + const a = self.rig.gpu.gpa; + if (self.r) |*r| r.deinit(); + inline for (.{ "committed", "row_stream", "paths", "depths", "offs", "q", "kn", "vn", "cache" }) |f| @field(self, f).deinit(a); + self.rig.deinit(); + } + + pub fn kernel(self: *Tree) g.Kernel { + return .{ .spec = treeSpec(self.v), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Tree); + + pub fn prepare(self: *Tree, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const v = self.v; + if (shape.n != v.h or shape.rows == 0 or shape.k == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + const n = shape.rows; + const ns = (n + width - 1) / width; + try self.committed.resize(a, ns); + for (self.committed.items, 0..) |*p, si| p.* = @max(1, shape.k - si * shape.k / ns); + try self.row_stream.resize(a, n); + try self.depths.resize(a, n); + try self.paths.resize(a, n * tree.max_depth); + @memset(self.paths.items, -1); + var streams = try a.alloc(i32, 4 * ns); + defer a.free(streams); + const parents = try a.alloc(i32, n); + defer a.free(parents); + for (0..ns) |si| { + const start = si * width; + const w = @min(width, n - start); + for (0..w) |j| { + self.row_stream.items[start + j] = @intCast(si); + parents[start + j] = if (j == 0 or rng.uintLessThan(u32, 6) == 0) -1 else @intCast(start + rng.uintLessThan(usize, j)); + } + for ([_]usize{ start, w, self.committed.items[si], 0 }, 0..) |x, k| streams[4 * si + k] = @intCast(x); + } + self.ng = 0; + var most: usize = 0; + for (0..n) |r| { + var depth: usize = 0; + var cur: i32 = @intCast(r); + while (cur >= 0) : (cur = parents[@intCast(cur)]) depth += 1; + self.depths.items[r] = @intCast(depth); + cur = @intCast(r); + var slot = depth; + while (slot > 0) : (cur = parents[@intCast(cur)]) { + slot -= 1; + self.paths.items[r * tree.max_depth + slot] = cur; + } + const keys = self.committed.items[@intCast(self.row_stream.items[r])] + depth; + self.ng = @max(self.ng, tree.TreeAttention.groups(keys)); + most = @max(most, keys); + } + // one-warp blocks take only grids that fit their slice area; anything larger runs four-warp blocks + if (v.sms > 0 and n * self.ng * v.hk > tree.slice_items) return false; + self.ta.sms = v.sms; + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, v.d, most); + const kv = v.hk * v.d; + var total: usize = 0; + for (self.committed.items) |p| total += p; + try self.offs.resize(a, 2 * ns); + var at: usize = 0; + for (self.committed.items, 0..) |p, si| { + self.offs.items[2 * si] = @intCast(at * kv); + self.offs.items[2 * si + 1] = @intCast((total + at) * kv); + at += p; + } + try self.cache.resize(a, 2 * total * kv); + try self.q.resize(a, n * v.h * v.d); + try self.kn.resize(a, n * kv); + try self.vn.resize(a, n * kv); + for ([_][]u16{ self.cache.items, self.q.items, self.kn.items, self.vn.items }) |x| for (x) |*e| { + e.* = attn.normal(rng); + }; + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng, total), + } + self.scale = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(v.d)))); + for ([_][]const u8{ std.mem.sliceAsBytes(self.q.items), std.mem.sliceAsBytes(self.kn.items), std.mem.sliceAsBytes(self.vn.items), std.mem.sliceAsBytes(self.cache.items), std.mem.sliceAsBytes(self.offs.items), std.mem.sliceAsBytes(self.row_stream.items), std.mem.sliceAsBytes(streams), std.mem.sliceAsBytes(self.paths.items), std.mem.sliceAsBytes(self.depths.items) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + const base = try self.rig.scratch(tree.TreeAttention.scratchBytes(n, v.h, v.hk, self.ng, v.d)); + self.sc = tree.TreeAttention.scratch(base, n, v.h, v.hk, self.ng, v.d); + self.out = try self.rig.output(n * v.h * v.d * 2); + return true; + } + + /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key. + fn plant(self: *Tree) void { + const d = self.v.d; + const kv = self.v.hk * d; + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + const p0 = self.committed.items[0]; + const ko: usize = @intCast(self.offs.items[0]); + const vo: usize = @intCast(self.offs.items[1]); + self.q.items[3] = nan; // row 0, head 0 + self.cache.items[vo + (p0 - 1) * kv + 0 * d + 7] = inf; + self.vn.items[1 * kv + 1 * d + 9] = ninf; // path row 1 + if (p0 >= 3) { + self.cache.items[vo + (p0 - 2) * kv + 2 * d + 11] = inf; + self.cache.items[vo + (p0 - 3) * kv + 2 * d + 11] = ninf; + } + self.cache.items[ko + (p0 - 1) * kv + 3 * d + 0] = nan; + } + + /// Gate 4: every key alike (each p is 1), committed values mirrored per column, path values 0. + fn stress(self: *Tree, rng: std.Random, total: usize) !void { + const a = self.rig.gpu.gpa; + const v = self.v; + const kv = v.hk * v.d; + const key = try a.alloc(u16, kv); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + for (0..total) |j| @memcpy(self.cache.items[j * kv ..][0..kv], key); + for (0..self.shape.rows) |j| @memcpy(self.kn.items[j * kv ..][0..kv], key); + @memset(self.vn.items, 0); + const col = try a.alloc(f32, total); + defer a.free(col); + for (self.committed.items, 0..) |p, si| for (0..kv) |e| { + attn.mirrored(rng, col[0..p]); + const vo: usize = @intCast(self.offs.items[2 * si + 1]); + for (col[0..p], 0..) |x, j| self.cache.items[vo + j * kv + e] = c.toBf16(x); + }; + } + + pub fn run(self: *Tree) !void { + const v = self.v; + const rws: tree.Rows = .{ .row_stream = self.dev[5], .streams = self.dev[6], .offs = self.dev[4], .origin = self.dev[3], .paths = self.dev[7], .depths = self.dev[8] }; + try self.ta.run(self.rig.s, self.dev[0], self.dev[1], self.dev[2], rws, self.sc, self.out, self.shape.rows, v.h, v.hk, v.d, self.ng, self.scale); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Tree) usize { + return self.shape.rows * self.v.h * self.v.d; + } + + /// Output head i (row i / h, head i % h): its stream's committed keys, then its path root first. + pub fn load(self: *Tree, i: usize, hd: *attn.Head) !void { + const v = self.v; + const d = v.d; + const kv = v.hk * d; + const r = i / v.h; + const head = i % v.h; + const kh = head / (v.h / v.hk); + const si: usize = @intCast(self.row_stream.items[r]); + const p = self.committed.items[si]; + const ko: usize = @intCast(self.offs.items[2 * si]); + const vo: usize = @intCast(self.offs.items[2 * si + 1]); + for (hd.q, 0..) |*x, e| x.* = c.fromBf16(self.q.items[(r * v.h + head) * d + e]); + for (0..p) |j| for (0..d) |e| { + hd.k[j * d + e] = c.fromBf16(self.cache.items[ko + j * kv + kh * d + e]); + hd.v[j * d + e] = c.fromBf16(self.cache.items[vo + j * kv + kh * d + e]); + }; + const depth: usize = @intCast(self.depths.items[r]); + for (0..depth) |j| { + const node: usize = @intCast(self.paths.items[r * tree.max_depth + j]); + for (0..d) |e| { + hd.k[(p + j) * d + e] = c.fromBf16(self.kn.items[node * kv + kh * d + e]); + hd.v[(p + j) * d + e] = c.fromBf16(self.vn.items[node * kv + kh * d + e]); + } + } + hd.n = p + depth; + } + + fn heads(self: *Tree) usize { + return self.shape.rows * self.v.h; + } + + pub fn check(self: *Tree, at: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), self.scale, try self.rig.outputBytes(), at); + } + + pub fn reference(self: *Tree, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), self.scale, want, mag); + } + + pub fn simulate(self: *Tree, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), self.scale, acc, order, out); + } +}; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const paths_rule = "integers only: each row's depth (at most 128) and its nearest 128 ancestors root first; other slots untouched"; + +pub fn pathsSpec() c.Spec { + return .{ .name = "cuda tree attention tf_tree_paths", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = paths_rule }, .tiles = .{ .{ .size = 128 }, .{}, .{} }, .served = &.{ .{ .rows = 16, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 }, .{ .rows = 600, .n = 1, .k = 1 } } }; +} + +/// tf_tree_paths over a forest of `rows` rows: random trees, and from 300 rows a 200-deep chain past the 128 cap. +pub const Paths = struct { + rig: rig_.Rig, + ta: *const tree.TreeAttention, + n: usize = 0, + parents: std.ArrayList(i32) = .empty, + dev: [3]u64 = @splat(0), + msg: [160]u8 = undefined, + + pub fn deinit(self: *Paths) void { + self.parents.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Paths) g.Kernel { + return .{ .spec = pathsSpec(), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Paths); + + pub fn prepare(self: *Paths, shape: c.Shape, _: g.Fill, seed: u64) !bool { + self.rig.reset(); + if (shape.rows == 0) return false; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.n = shape.rows; + try self.parents.resize(self.rig.gpu.gpa, self.n); + for (self.parents.items, 0..) |*p, r| p.* = if (r == 0) -1 else if (self.n >= 300 and r < 200) @intCast(r - 1) else if (rng.uintLessThan(u32, 5) == 0) -1 else @intCast(rng.uintLessThan(usize, r)); + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.parents.items)); + self.dev[1] = try self.rig.output(self.n * tree.max_depth * 4); + self.dev[2] = try self.rig.output(self.n * 4); + return true; + } + + pub fn run(self: *Paths) !void { + try self.ta.paths(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.n); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Paths) usize { + return self.n * (tree.max_depth + 1) * 4; + } + + pub fn reference(_: *Paths, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + pub fn check(self: *Paths, _: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const got = try self.rig.outputBytes(); + const want = try a.alloc(u8, got.len); + defer a.free(want); + @memset(want, 0xff); + const paths = std.mem.bytesAsSlice(i32, want[0 .. self.n * tree.max_depth * 4]); + const depths = std.mem.bytesAsSlice(i32, want[self.n * tree.max_depth * 4 ..]); + for (0..self.n) |r| { + var depth: usize = 0; + var cur: i32 = @intCast(r); + while (cur >= 0 and depth < tree.max_depth) : (cur = self.parents.items[@intCast(cur)]) depth += 1; + depths[r] = @intCast(depth); + cur = @intCast(r); + var slot = depth; + while (slot > 0) : (cur = self.parents.items[@intCast(cur)]) { + slot -= 1; + paths[r * tree.max_depth + slot] = cur; + } + } + if (std.mem.indexOfDiff(u8, got, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} is 0x{x:0>2} where the host gives 0x{x:0>2}", .{ at, got.len, got[at], want[at] }); + return null; + } +}; From 492716baa56c1d14fb16710b1307913e78b1cfd0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 17/41] cuda/tree_attention: each row's own max flag, so a NaN head leaves other rows' keys --- zig/kernels/cuda/tree_attention.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index eea5d2645d..4e50f442e9 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -190,10 +190,10 @@ __device__ __forceinline__ void tree_group( tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); } - const bool act = tm0 != -INFINITY; - const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; - const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; - const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; // each row's own: a NaN head moves only its row + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; float c0 = 0.0f, c1 = 0.0f; #pragma unroll for (int n = 0; n < NT8; ++n) { From f812fbe7b1eacfca39b99e7bbf22fc056a54a1d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 18/41] cuda/block_attention: keys a row cannot see add nothing, so an infinite value there stays out of its sum A block whose output holds a NaN or infinity runs again with the values of the tiles some row cannot see all of taken out of the mma; each row adds back the ones it sees. On finite values the first pass is the kernel as it was, byte for byte. --- zig/kernels/cuda/attn_odd.cuh | 63 +++++++++++++++++++++++++++++ zig/kernels/cuda/block_attention.cu | 60 ++++++++++++++++++--------- 2 files changed, 105 insertions(+), 18 deletions(-) create mode 100644 zig/kernels/cuda/attn_odd.cuh diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh new file mode 100644 index 0000000000..596ee3a95e --- /dev/null +++ b/zig/kernels/cuda/attn_odd.cuh @@ -0,0 +1,63 @@ +// NaN and infinite values in a tensor-core attention tile's V: a key a row cannot see adds nothing to that row's sum. +#pragma once + +#include +#include +#include + +namespace attn_odd { + +// Each bf16 half of x that is NaN or an infinity (all-ones exponent) as +0, the finite halves as they are. +__device__ __forceinline__ uint32_t finite2(uint32_t x) { + return x & (((x & 0x7f80u) == 0x7f80u ? 0u : 0xffffu) | ((x & 0x7f800000u) == 0x7f800000u ? 0u : 0xffff0000u)); +} + +// Some NaN or infinite value in a warp's mma rows g (r0: a real row) and g + 8 (r1) of its fold. +template +__device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, bool r1) { + float a = 0.0f, b = 0.0f; // x - x is 0 for a finite x, NaN for NaN and the infinities +#pragma unroll + for (int i = 0; i < NB; ++i) { + a = __fadd_rn(a, __fadd_rn(__fsub_rn(o[i][0], o[i][0]), __fsub_rn(o[i][1], o[i][1]))); + b = __fadd_rn(b, __fadd_rn(__fsub_rn(o[i][2], o[i][2]), __fsub_rn(o[i][3], o[i][3]))); + } + return (r0 && a != a) || (r1 && b != b); +} + +// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. +template +__device__ __forceinline__ void clean(uint4* vs, int threads) { + for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) { + uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7)); + const uint4 x = *p; + *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w)); + } +} + +// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them. +template +__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0, + float l1) { + const int g = lane / 4, qd = lane % 4; + for (int k = k0; k < k1; ++k) { + const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN + if (!v0 && !v1) continue; + const uint32_t* v = reinterpret_cast(vrow(k)); +#pragma unroll + for (int i = 0; i < D / 8; ++i) { + const uint32_t x = v[4 * i + qd]; + if (finite2(x) == x) continue; + const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u); + if ((x & 0x7f80u) == 0x7f80u) { + if (v0) o[i][0] = __fadd_rn(o[i][0], lo); + if (v1) o[i][2] = __fadd_rn(o[i][2], lo); + } + if ((x & 0x7f800000u) == 0x7f800000u) { + if (v0) o[i][1] = __fadd_rn(o[i][1], hi); + if (v1) o[i][3] = __fadd_rn(o[i][3], hi); + } + } + } +} + +} // namespace attn_odd diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index e3cd06b121..f5b1e2b18a 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -7,6 +7,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_block { @@ -166,30 +167,53 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( } Fold f; const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) - for (int k0 = lo / KT * KT; k0 < s; k0 += KT) { + const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; + // the context's edge tiles, where some row cannot see some key: e0..e1 + const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT); + auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below) + for (int k0 = e0; k0 < s; k0 += KT) { + for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { + const int key = k0 + i / C; + const bool in = key < s; + cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + } + commit(); + wait_all(); + if (odd && k0 < e1) attn_odd::clean(vs, blockDim.x); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, k0, vis); + __syncthreads(); + } + }; + auto block_keys = [&] { // the block's own keys for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { - const int key = k0 + i / C; - const bool in = key < s; - cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); - cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + const int c = i / C; + const bool in = c < L; + const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; + cp16z(ks + sw(c, i % C), kb + at, in); + cp16z(vs + sw(c, i % C), vb + at, in); } commit(); wait_all(); + }; + const auto own = [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }; + context(false); + block_keys(); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, 0, own); + // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(f.o, g8 < L, g8 + 8 < L))) { + f = Fold(); + context(true); + attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1); + block_keys(); + if (causal) attn_odd::clean(vs, blockDim.x); // row r cannot see keys past r __syncthreads(); - f.tile(qa, ks, vs, scale, lane, k0, [&](int r, int key) { return r < L && key < s && s + r - key <= window; }); - __syncthreads(); - } - for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { // the block's own keys - const int c = i / C; - const bool in = c < L; - const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; - cp16z(ks + sw(c, i % C), kb + at, in); - cp16z(vs + sw(c, i % C), vb + at, in); + f.tile(qa, ks, vs, scale, lane, 0, own); + const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; + if (causal) attn_odd::add_back(f.o, row, 0, L, lane, own, f.l0, f.l1); } - commit(); - wait_all(); - __syncthreads(); - f.tile(qa, ks, vs, scale, lane, 0, [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }); if (g8 < L) #pragma unroll for (int i = 0; i < NB; ++i) From d18a7205bb785dfc9fbc3af886c1f755f7df2de3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 19/41] cuda/tree_attention, block_attention: registered with the kernel contract --- zig/src/core/contract/registry.zig | 15 ++ zig/tests/cuda/contract_attn.zig | 124 +++++++++ zig/tests/cuda/contract_block_attn.zig | 224 +++++++++++++++++ zig/tests/cuda/contract_tests.zig | 25 +- zig/tests/cuda/contract_tree_attn.zig | 333 +++++++++++++++++++++++++ 5 files changed, 720 insertions(+), 1 deletion(-) create mode 100644 zig/tests/cuda/contract_attn.zig create mode 100644 zig/tests/cuda/contract_block_attn.zig create mode 100644 zig/tests/cuda/contract_tree_attn.zig diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..b2cfbc0e18 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,21 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_paths", .symbol = "tf_tree_paths", .image = "tree_attention", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1s3", .symbol = "tf_tree_shared_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part", .symbol = "tf_tree_part", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1", .symbol = "tf_tree_part_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1s3", .symbol = "tf_tree_part_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256", .symbol = "tf_tree_shared_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1", .symbol = "tf_tree_shared_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1s3", .symbol = "tf_tree_shared_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256", .symbol = "tf_tree_part_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_attn.zig b/zig/tests/cuda/contract_attn.zig new file mode 100644 index 0000000000..cec73374a3 --- /dev/null +++ b/zig/tests/cuda/contract_attn.zig @@ -0,0 +1,124 @@ +//! Attention kernels under the contract on CUDA: their outputs against core/attention_ref.zig, head by head. +const std = @import("std"); +const core = @import("core"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const ref = core.attention_ref; + +/// Why gates 1 and 3 run attention_ref's bound: P enters the P x V mma as bf16, and the spec's bound takes fp32 terms. +pub const why = "P is rounded to bf16 before the P x V mma: gates 1 and 3 hold each output to attention_ref's float64 bound"; + +/// One output head's float64 inputs: its query, and its visible keys' K and V rows ([keys][d]). +pub const Head = struct { q: []f64, k: []f64, v: []f64, n: usize = 0 }; + +/// The float64 side of an attention test: `t.load(i, *Head)` fills output head i (outputs [heads][d] bf16). +pub const Ref = struct { + a: std.mem.Allocator, + d: usize, + head: Head, + keys: []usize, + p: []f64, + want: []f64, + tol: []f64, + msg: [240]u8 = undefined, + + pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref { + const keys = try a.alloc(usize, most); + for (keys, 0..) |*k, i| k.* = i; + return .{ .a = a, .d = d, .head = .{ .q = try a.alloc(f64, d), .k = try a.alloc(f64, most * d), .v = try a.alloc(f64, most * d) }, .keys = keys, .p = try a.alloc(f64, most), .want = try a.alloc(f64, d), .tol = try a.alloc(f64, d) }; + } + + pub fn deinit(r: *Ref) void { + for ([_][]f64{ r.head.q, r.head.k, r.head.v, r.p, r.want, r.tol }) |x| r.a.free(x); + r.a.free(r.keys); + } + + /// Gates 1 and 3: every output's class equals float64's, and a finite one is within attention_ref's bound. + pub fn check(r: *Ref, t: anytype, heads: usize, scale: f64, out: []const u8, at: g.Gate) !?[]const u8 { + var planted: usize = 0; + for (0..heads) |i| { + try t.load(i, &r.head); + ref.head(r.head.q, r.head.k, r.head.v, r.keys[0..r.head.n], r.d, scale, r.want, r.tol); + for (r.want, r.tol, 0..) |w, tl, e| { + const got = c.read(.bf16, out, i * r.d + e); + if (c.classOf(w) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(w)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d}: {t} where float64 gives {t}", .{ i, e, c.classOf(got), c.classOf(w) }); + if (at != .nonfinite and std.math.isFinite(w) and !(@abs(got - w) <= tl)) return try std.fmt.bufPrint(&r.msg, "head {d} element {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of attention_ref's bound", .{ i, e, got, w, @abs(got - w) / tl }); + } + } + if (at == .nonfinite and (planted == 0 or planted == heads * r.d)) return try std.fmt.bufPrint(&r.msg, "the planted values make {d} of {d} reference outputs non-finite: no proof", .{ planted, heads * r.d }); + return null; + } + + fn weights(r: *Ref, scale: f64) f64 { + const h = r.head; + var m: f64 = -std.math.inf(f64); + for (r.p[0..h.n], 0..) |*p, j| { + var s: f64 = 0; + for (0..r.d) |e| s += h.q[e] * h.k[j * r.d + e]; + p.* = s * scale; + m = @max(m, p.*); + } + var l: f64 = 0; + for (r.p[0..h.n]) |*p| { + p.* = @exp(p.* - m); + l += p.*; + } + return l; + } + + /// Float64 outputs and each one's term magnitude, sum p_j |v_j| / sum p_j (gate 4 reads them). + pub fn reference(r: *Ref, t: anytype, heads: usize, scale: f64, want: []f64, mag: []f64) !void { + for (0..heads) |i| { + try t.load(i, &r.head); + const l = r.weights(scale); + for (0..r.d) |e| { + var s: f64 = 0; + var abs: f64 = 0; + for (r.p[0..r.head.n], 0..) |p, j| { + s += p * r.head.v[j * r.d + e]; + abs += p * @abs(r.head.v[j * r.d + e]); + } + want[i * r.d + e] = s / l; + mag[i * r.d + e] = abs / l; + } + } + } + + /// Gate 4's sums in an accumulator: its keys score alike, so every p is 1 and the sums are V's terms and l's ones. + pub fn simulate(r: *Ref, t: anytype, heads: usize, scale: f64, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try r.a.alloc(f64, r.p.len); + defer r.a.free(terms); + const ones = try r.a.alloc(f64, r.p.len); + defer r.a.free(ones); + @memset(ones, 1); + for (0..heads) |i| { + try t.load(i, &r.head); + _ = r.weights(scale); + for (r.p[0..r.head.n]) |p| if (p != 1) return error.KeysScoreApart; + const l = rows.sumIn(acc, order, ones[0..r.head.n]); + for (0..r.d) |e| { + for (terms[0..r.head.n], 0..) |*x, j| x.* = r.head.v[j * r.d + e]; + out[i * r.d + e] = rows.sumIn(acc, order, terms[0..r.head.n]) / l; + } + } + } +}; + +/// Gate 4's values: odd integers, every other one times 64 (bf16-exact), the second half cancelling the first. +pub 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; +} + +pub fn normal(rng: std.Random) u16 { + return c.toBf16(rng.floatNorm(f32)); +} diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig new file mode 100644 index 0000000000..155374b86d --- /dev/null +++ b/zig/tests/cuda/contract_block_attn.zig @@ -0,0 +1,224 @@ +//! block_attention.cu (#746) under the contract on CUDA: DFlash block rows over their windowed context, head by head. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const block = cuda.block_attention; + +const D = 128; +const h = 32; +const hk = 8; +const window = 2048; + +/// The P x V chain of a row: its window's keys and the block's, a 32-key tile at a time. +fn keyChain(s: c.Shape) usize { + return @min(s.k, window + 1) + 32; +} + +/// Shape: streams (rows), query heads (n) and stream 0's context keys (k); each stream's block of L rows. +pub fn spec(causal: bool) c.Spec { + // Qwen3.8-27B's DFlash2 drafter: blocks of 16 over 32 heads (8 key heads) and a 2,048-key window + return .{ + .name = if (causal) "cuda block attention tf_block_attention, causal blocks of 7" else "cuda block attention tf_block_attention (DFlash2 blocks of 16)", + .accum = .f32, + .depth = keyChain, + .ops = 1, + .out = .bf16, + .tiles = .{ .{}, .{}, .{ .size = 32 } }, + .served = &.{ .{ .rows = 1, .n = h, .k = 2100 }, .{ .rows = 4, .n = h, .k = 700 } }, + .sums = .{ .why = attn.why }, + }; +} + +pub const Block = struct { + rig: rig_.Rig, + ba: *const block.BlockAttention, + causal: bool, + r: ?attn.Ref = null, + ns: usize = 0, + lens: std.ArrayList(i32) = .empty, + ctx: std.ArrayList([]u16) = .empty, + q: std.ArrayList(u16) = .empty, + kb: std.ArrayList(u16) = .empty, + vb: std.ArrayList(u16) = .empty, + dev: [5]u64 = @splat(0), + out: u64 = 0, + + fn len(self: *const Block) usize { + return if (self.causal) 7 else 16; + } + + fn scale() f32 { + return @floatCast(1 / @sqrt(@as(f64, D))); + } + + pub fn deinit(self: *Block) void { + const a = self.rig.gpu.gpa; + self.drop(); + if (self.r) |*r| r.deinit(); + inline for (.{ "lens", "ctx", "q", "kb", "vb" }) |f| @field(self, f).deinit(a); + self.rig.deinit(); + } + + fn drop(self: *Block) void { + for (self.ctx.items) |x| self.rig.gpu.gpa.free(x); + self.ctx.clearRetainingCapacity(); + } + + pub fn kernel(self: *Block) g.Kernel { + return .{ .spec = spec(self.causal), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Block); + + pub fn prepare(self: *Block, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.drop(); + if (shape.n != h or shape.rows == 0 or shape.k == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const L = self.len(); + self.ns = shape.rows; + const rows = self.ns * L; + try self.lens.resize(a, self.ns); + for (self.lens.items, 0..) |*l, j| l.* = @intCast(@max(1, shape.k - j * shape.k / self.ns)); + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, D, @min(shape.k, window + 1) + L); + for (self.lens.items) |l| { + const x = try a.alloc(u16, 2 * hk * @as(usize, @intCast(l)) * D); // K [hk, len, D], then V + try self.ctx.append(a, x); + for (x) |*e| e.* = attn.normal(rng); + } + try self.q.resize(a, h * rows * D); + try self.kb.resize(a, hk * rows * D); + try self.vb.resize(a, hk * rows * D); + for ([_][]u16{ self.q.items, self.kb.items, self.vb.items }) |x| for (x) |*e| { + e.* = attn.normal(rng); + }; + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng), + } + const table = try a.alloc(i64, 2 * self.ns); + defer a.free(table); + for (self.ctx.items, 0..) |x, j| { + const base = try self.rig.input(std.mem.sliceAsBytes(x)); + table[2 * j] = @intCast(base); + table[2 * j + 1] = @intCast(base + x.len); // V follows K: half the u16s, all the bytes + } + for ([_][]const u8{ std.mem.sliceAsBytes(self.q.items), std.mem.sliceAsBytes(self.kb.items), std.mem.sliceAsBytes(self.vb.items), std.mem.sliceAsBytes(table), std.mem.sliceAsBytes(self.lens.items) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + self.out = try self.rig.output(rows * h * D * 2); + return true; + } + + fn ctxLen(self: *const Block, j: usize) usize { + return @intCast(self.lens.items[j]); + } + + /// NaN in a query, ±inf in values (context, block, a window's first key), both in one column, NaN in a key. + fn plant(self: *Block) void { + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + const v0 = self.ctx.items[0][hk * s0 * D ..]; + self.q.items[3] = nan; // head 0, stream 0's row 0 + v0[(0 * s0 + s0 - 1) * D + 7] = inf; + self.vb.items[(1 * rows + 2) * D + 9] = ninf; // key head 1, stream 0's row 2 + if (s0 >= 3) { + v0[(2 * s0 + s0 - 2) * D + 11] = inf; + v0[(2 * s0 + s0 - 3) * D + 11] = ninf; + } + if (s0 > window) v0[(4 * s0 + s0 - window) * D + 13] = inf; // a key block row 0 alone sees + self.ctx.items[self.ns - 1][(3 * self.ctxLen(self.ns - 1) + self.ctxLen(self.ns - 1) - 1) * D] = nan; + } + + /// Gate 4: every key alike (each p is 1); values mirrored over the context keys every row sees, else 0. + fn stress(self: *Block, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const key = try a.alloc(u16, D); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + const L = self.len(); + @memset(self.vb.items, 0); + for (0..self.kb.items.len / D) |i| @memcpy(self.kb.items[i * D ..][0..D], key); + for (self.ctx.items, 0..) |x, j| { + const s = self.ctxLen(j); + for (0..hk * s) |i| @memcpy(x[i * D ..][0..D], key); + const v = x[hk * s * D ..]; + @memset(v, 0); + const lo = @min(s, (s + L) -| window); + const col = try a.alloc(f32, s - lo); + defer a.free(col); + for (0..hk) |kh| for (0..D) |e| { + attn.mirrored(rng, col); + for (col, lo..) |y, k| v[(kh * s + k) * D + e] = c.toBf16(y); + }; + } + } + + pub fn run(self: *Block) !void { + const L = self.len(); + try self.ba.run(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.dev[3], self.dev[4], self.out, .{ .h = h, .hk = hk, .rows = self.ns * L, .block = L, .window = window, .scale = scale(), .causal = self.causal }); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Block) usize { + return self.ns * self.len() * h * D; + } + + /// Output head i (block row i / h, head i % h): context keys k with len + r - k <= window, then the block's. + pub fn load(self: *Block, i: usize, hd: *attn.Head) !void { + const L = self.len(); + const rows = self.ns * L; + const row = i / h; + const head = i % h; + const kh = head / (h / hk); + const j = row / L; + const r = row % L; + const s = self.ctxLen(j); + const x = self.ctx.items[j]; + for (hd.q, 0..) |*y, e| y.* = c.fromBf16(self.q.items[(head * rows + row) * D + e]); + var n: usize = 0; + for (0..s) |k| if (s + r - k <= window) { + for (0..D) |e| { + hd.k[n * D + e] = c.fromBf16(x[(kh * s + k) * D + e]); + hd.v[n * D + e] = c.fromBf16(x[hk * s * D + (kh * s + k) * D + e]); + } + n += 1; + }; + for (0..L) |col| if (!self.causal or col <= r) { + for (0..D) |e| { + hd.k[n * D + e] = c.fromBf16(self.kb.items[(kh * rows + j * L + col) * D + e]); + hd.v[n * D + e] = c.fromBf16(self.vb.items[(kh * rows + j * L + col) * D + e]); + } + n += 1; + }; + hd.n = n; + } + + fn heads(self: *Block) usize { + return self.ns * self.len() * h; + } + + pub fn check(self: *Block, at: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), scale(), try self.rig.outputBytes(), at); + } + + pub fn reference(self: *Block, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), scale(), want, mag); + } + + pub fn simulate(self: *Block, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), scale(), acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e87a2e90e5..8114049a57 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -8,6 +8,8 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const rank_k = @import("contract_rank.zig"); +const tree_k = @import("contract_tree_attn.zig"); +const block_k = @import("contract_block_attn.zig"); const contract = core.contract; const g = contract.gates; @@ -21,7 +23,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "tree_attention", "block_attention" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -89,6 +91,27 @@ pub fn run(gpu: check.Gpu) !void { var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var tree_paths = try cuda.tree_attention.TreeAttention.load(gpu.d); + defer tree_paths.unload(); + var paths = tree_k.Paths{ .rig = .{ .gpu = gpu, .s = s }, .ta = &tree_paths }; + defer paths.deinit(); + try kernels.append(a, paths.kernel()); + var trees: [tree_k.variants.len]tree_k.Tree = undefined; + var ntrees: usize = 0; + defer for (trees[0..ntrees]) |*x| { + x.ta.unload(); + x.deinit(); + }; + for (tree_k.variants) |v| { + trees[ntrees] = .{ .rig = .{ .gpu = gpu, .s = s }, .v = v, .ta = try cuda.tree_attention.TreeAttention.load(gpu.d) }; + ntrees += 1; + } + for (&trees) |*x| try kernels.append(a, x.kernel()); + var block_attn = try cuda.block_attention.BlockAttention.load(gpu.d); + defer block_attn.unload(); + var blocks: [2]block_k.Block = .{ .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = false }, .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = true } }; + defer for (&blocks) |*x| x.deinit(); + for (&blocks) |*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_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig new file mode 100644 index 0000000000..3dd57f47e6 --- /dev/null +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -0,0 +1,333 @@ +//! tree_attention.cu (#746) under the contract on CUDA: paths byte for byte, attention head by head against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const tree = cuda.tree_attention; + +/// A tree per stream: its rows (at most 16, the 27B's tree), each a path through earlier rows of its stream. +const width = 16; + +/// The P x V chain of the longest row: every committed key, its path, a few combines. +fn keyChain(s: c.Shape) usize { + return s.k + tree.max_depth + 16; +} + +/// A tree attention call's blocks: four warps (sms 0), or one warp a slice with 2 (sms 1) or 3 tile stages. +pub const Variant = struct { d: usize, h: usize, hk: usize, sms: usize, name: []const u8 }; + +pub const variants = [_]Variant{ + .{ .d = 128, .h = 48, .hk = 4, .sms = 0, .name = "cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge)" }, + .{ .d = 128, .h = 48, .hk = 4, .sms = 1, .name = "cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge)" }, + .{ .d = 128, .h = 48, .hk = 4, .sms = 1 << 20, .name = "cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 0, .name = "cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 1, .name = "cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1)" }, + .{ .d = 256, .h = 24, .hk = 4, .sms = 1 << 20, .name = "cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3)" }, +}; + +/// Shape: drafted rows (trees of 16 over the streams), query heads (n) and stream 0's committed keys (k). +pub fn treeSpec(v: Variant) c.Spec { + // Kolibri (dim 128, 48 heads over 4) and Qwen3.8-27B (dim 256, 24 over 4): one stream's tree, then four streams' + const served: []const c.Shape = if (v.d == 128) &.{ .{ .rows = 16, .n = 48, .k = 2100 }, .{ .rows = 64, .n = 48, .k = 700 } } else &.{ .{ .rows = 16, .n = 24, .k = 2100 }, .{ .rows = 64, .n = 24, .k = 700 } }; + return .{ .name = v.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = width }, .{}, .{ .size = tree.span } }, .served = served, .sums = .{ .why = attn.why } }; +} + +pub const Tree = struct { + rig: rig_.Rig, + v: Variant, + ta: tree.TreeAttention, + r: ?attn.Ref = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + scale: f32 = 0, + ng: usize = 0, + committed: std.ArrayList(usize) = .empty, + row_stream: std.ArrayList(i32) = .empty, + paths: std.ArrayList(i32) = .empty, + depths: std.ArrayList(i32) = .empty, + offs: std.ArrayList(i64) = .empty, + q: std.ArrayList(u16) = .empty, + kn: std.ArrayList(u16) = .empty, + vn: std.ArrayList(u16) = .empty, + cache: std.ArrayList(u16) = .empty, + dev: [9]u64 = @splat(0), + sc: tree.Scratch = undefined, + out: u64 = 0, + + pub fn deinit(self: *Tree) void { + const a = self.rig.gpu.gpa; + if (self.r) |*r| r.deinit(); + inline for (.{ "committed", "row_stream", "paths", "depths", "offs", "q", "kn", "vn", "cache" }) |f| @field(self, f).deinit(a); + self.rig.deinit(); + } + + pub fn kernel(self: *Tree) g.Kernel { + return .{ .spec = treeSpec(self.v), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Tree); + + pub fn prepare(self: *Tree, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const v = self.v; + if (shape.n != v.h or shape.rows == 0 or shape.k == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + const n = shape.rows; + const ns = (n + width - 1) / width; + try self.committed.resize(a, ns); + for (self.committed.items, 0..) |*p, si| p.* = @max(1, shape.k - si * shape.k / ns); + try self.row_stream.resize(a, n); + try self.depths.resize(a, n); + try self.paths.resize(a, n * tree.max_depth); + @memset(self.paths.items, -1); + var streams = try a.alloc(i32, 4 * ns); + defer a.free(streams); + const parents = try a.alloc(i32, n); + defer a.free(parents); + for (0..ns) |si| { + const start = si * width; + const w = @min(width, n - start); + for (0..w) |j| { + self.row_stream.items[start + j] = @intCast(si); + parents[start + j] = if (j == 0 or rng.uintLessThan(u32, 6) == 0) -1 else @intCast(start + rng.uintLessThan(usize, j)); + } + for ([_]usize{ start, w, self.committed.items[si], 0 }, 0..) |x, k| streams[4 * si + k] = @intCast(x); + } + self.ng = 0; + var most: usize = 0; + for (0..n) |r| { + var depth: usize = 0; + var cur: i32 = @intCast(r); + while (cur >= 0) : (cur = parents[@intCast(cur)]) depth += 1; + self.depths.items[r] = @intCast(depth); + cur = @intCast(r); + var slot = depth; + while (slot > 0) : (cur = parents[@intCast(cur)]) { + slot -= 1; + self.paths.items[r * tree.max_depth + slot] = cur; + } + const keys = self.committed.items[@intCast(self.row_stream.items[r])] + depth; + self.ng = @max(self.ng, tree.TreeAttention.groups(keys)); + most = @max(most, keys); + } + // one-warp blocks take only grids that fit their slice area; anything larger runs four-warp blocks + if (v.sms > 0 and n * self.ng * v.hk > tree.slice_items) return false; + self.ta.sms = v.sms; + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, v.d, most); + const kv = v.hk * v.d; + var total: usize = 0; + for (self.committed.items) |p| total += p; + try self.offs.resize(a, 2 * ns); + var at: usize = 0; + for (self.committed.items, 0..) |p, si| { + self.offs.items[2 * si] = @intCast(at * kv); + self.offs.items[2 * si + 1] = @intCast((total + at) * kv); + at += p; + } + try self.cache.resize(a, 2 * total * kv); + try self.q.resize(a, n * v.h * v.d); + try self.kn.resize(a, n * kv); + try self.vn.resize(a, n * kv); + for ([_][]u16{ self.cache.items, self.q.items, self.kn.items, self.vn.items }) |x| for (x) |*e| { + e.* = attn.normal(rng); + }; + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng, total), + } + self.scale = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(v.d)))); + for ([_][]const u8{ std.mem.sliceAsBytes(self.q.items), std.mem.sliceAsBytes(self.kn.items), std.mem.sliceAsBytes(self.vn.items), std.mem.sliceAsBytes(self.cache.items), std.mem.sliceAsBytes(self.offs.items), std.mem.sliceAsBytes(self.row_stream.items), std.mem.sliceAsBytes(streams), std.mem.sliceAsBytes(self.paths.items), std.mem.sliceAsBytes(self.depths.items) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + const base = try self.rig.scratch(tree.TreeAttention.scratchBytes(n, v.h, v.hk, self.ng, v.d)); + self.sc = tree.TreeAttention.scratch(base, n, v.h, v.hk, self.ng, v.d); + self.out = try self.rig.output(n * v.h * v.d * 2); + return true; + } + + /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key. + fn plant(self: *Tree) void { + const d = self.v.d; + const kv = self.v.hk * d; + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + const p0 = self.committed.items[0]; + const ko: usize = @intCast(self.offs.items[0]); + const vo: usize = @intCast(self.offs.items[1]); + self.q.items[3] = nan; // row 0, head 0 + self.cache.items[vo + (p0 - 1) * kv + 0 * d + 7] = inf; + self.vn.items[1 * kv + 1 * d + 9] = ninf; // path row 1 + if (p0 >= 3) { + self.cache.items[vo + (p0 - 2) * kv + 2 * d + 11] = inf; + self.cache.items[vo + (p0 - 3) * kv + 2 * d + 11] = ninf; + } + self.cache.items[ko + (p0 - 1) * kv + 3 * d + 0] = nan; + } + + /// Gate 4: every key alike (each p is 1), committed values mirrored per column, path values 0. + fn stress(self: *Tree, rng: std.Random, total: usize) !void { + const a = self.rig.gpu.gpa; + const v = self.v; + const kv = v.hk * v.d; + const key = try a.alloc(u16, kv); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + for (0..total) |j| @memcpy(self.cache.items[j * kv ..][0..kv], key); + for (0..self.shape.rows) |j| @memcpy(self.kn.items[j * kv ..][0..kv], key); + @memset(self.vn.items, 0); + const col = try a.alloc(f32, total); + defer a.free(col); + for (self.committed.items, 0..) |p, si| for (0..kv) |e| { + attn.mirrored(rng, col[0..p]); + const vo: usize = @intCast(self.offs.items[2 * si + 1]); + for (col[0..p], 0..) |x, j| self.cache.items[vo + j * kv + e] = c.toBf16(x); + }; + } + + pub fn run(self: *Tree) !void { + const v = self.v; + const rws: tree.Rows = .{ .row_stream = self.dev[5], .streams = self.dev[6], .offs = self.dev[4], .origin = self.dev[3], .paths = self.dev[7], .depths = self.dev[8] }; + try self.ta.run(self.rig.s, self.dev[0], self.dev[1], self.dev[2], rws, self.sc, self.out, self.shape.rows, v.h, v.hk, v.d, self.ng, self.scale); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Tree) usize { + return self.shape.rows * self.v.h * self.v.d; + } + + /// Output head i (row i / h, head i % h): its stream's committed keys, then its path root first. + pub fn load(self: *Tree, i: usize, hd: *attn.Head) !void { + const v = self.v; + const d = v.d; + const kv = v.hk * d; + const r = i / v.h; + const head = i % v.h; + const kh = head / (v.h / v.hk); + const si: usize = @intCast(self.row_stream.items[r]); + const p = self.committed.items[si]; + const ko: usize = @intCast(self.offs.items[2 * si]); + const vo: usize = @intCast(self.offs.items[2 * si + 1]); + for (hd.q, 0..) |*x, e| x.* = c.fromBf16(self.q.items[(r * v.h + head) * d + e]); + for (0..p) |j| for (0..d) |e| { + hd.k[j * d + e] = c.fromBf16(self.cache.items[ko + j * kv + kh * d + e]); + hd.v[j * d + e] = c.fromBf16(self.cache.items[vo + j * kv + kh * d + e]); + }; + const depth: usize = @intCast(self.depths.items[r]); + for (0..depth) |j| { + const node: usize = @intCast(self.paths.items[r * tree.max_depth + j]); + for (0..d) |e| { + hd.k[(p + j) * d + e] = c.fromBf16(self.kn.items[node * kv + kh * d + e]); + hd.v[(p + j) * d + e] = c.fromBf16(self.vn.items[node * kv + kh * d + e]); + } + } + hd.n = p + depth; + } + + fn heads(self: *Tree) usize { + return self.shape.rows * self.v.h; + } + + pub fn check(self: *Tree, at: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), self.scale, try self.rig.outputBytes(), at); + } + + pub fn reference(self: *Tree, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), self.scale, want, mag); + } + + pub fn simulate(self: *Tree, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), self.scale, acc, order, out); + } +}; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const paths_rule = "integers only: each row's depth (at most 128) and its nearest 128 ancestors root first; other slots untouched"; + +pub fn pathsSpec() c.Spec { + return .{ .name = "cuda tree attention tf_tree_paths", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = paths_rule }, .tiles = .{ .{ .size = 128 }, .{}, .{} }, .served = &.{ .{ .rows = 16, .n = 1, .k = 1 }, .{ .rows = 128, .n = 1, .k = 1 }, .{ .rows = 600, .n = 1, .k = 1 } } }; +} + +/// tf_tree_paths over a forest of `rows` rows: random trees, and from 300 rows a 200-deep chain past the 128 cap. +pub const Paths = struct { + rig: rig_.Rig, + ta: *const tree.TreeAttention, + n: usize = 0, + parents: std.ArrayList(i32) = .empty, + dev: [3]u64 = @splat(0), + msg: [160]u8 = undefined, + + pub fn deinit(self: *Paths) void { + self.parents.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Paths) g.Kernel { + return .{ .spec = pathsSpec(), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Paths); + + pub fn prepare(self: *Paths, shape: c.Shape, _: g.Fill, seed: u64) !bool { + self.rig.reset(); + if (shape.rows == 0) return false; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.n = shape.rows; + try self.parents.resize(self.rig.gpu.gpa, self.n); + for (self.parents.items, 0..) |*p, r| p.* = if (r == 0) -1 else if (self.n >= 300 and r < 200) @intCast(r - 1) else if (rng.uintLessThan(u32, 5) == 0) -1 else @intCast(rng.uintLessThan(usize, r)); + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.parents.items)); + self.dev[1] = try self.rig.output(self.n * tree.max_depth * 4); + self.dev[2] = try self.rig.output(self.n * 4); + return true; + } + + pub fn run(self: *Paths) !void { + try self.ta.paths(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.n); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Paths) usize { + return self.n * (tree.max_depth + 1) * 4; + } + + pub fn reference(_: *Paths, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + pub fn check(self: *Paths, _: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const got = try self.rig.outputBytes(); + const want = try a.alloc(u8, got.len); + defer a.free(want); + @memset(want, 0xff); + const paths = std.mem.bytesAsSlice(i32, want[0 .. self.n * tree.max_depth * 4]); + const depths = std.mem.bytesAsSlice(i32, want[self.n * tree.max_depth * 4 ..]); + for (0..self.n) |r| { + var depth: usize = 0; + var cur: i32 = @intCast(r); + while (cur >= 0 and depth < tree.max_depth) : (cur = self.parents.items[@intCast(cur)]) depth += 1; + depths[r] = @intCast(depth); + cur = @intCast(r); + var slot = depth; + while (slot > 0) : (cur = self.parents.items[@intCast(cur)]) { + slot -= 1; + paths[r * tree.max_depth + slot] = cur; + } + } + if (std.mem.indexOfDiff(u8, got, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} is 0x{x:0>2} where the host gives 0x{x:0>2}", .{ at, got.len, got[at], want[at] }); + return null; + } +}; From ce160986b93c0371f79e233e69ec7098a40e1ef6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:52:32 +0200 Subject: [PATCH 20/41] cuda/window_attention: the tensor-core decode keeps each head's own max flag, so a NaN head leaves the other heads' keys --- zig/kernels/cuda/window_attention.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index 542e80a8ba..4e009fc669 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -439,10 +439,10 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); } - const bool act = tm0 != -INFINITY; // the keys' mask is the same for every head - const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; - const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; - const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; // each head's own: a NaN head moves only its row + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; float c0 = 0.0f, c1 = 0.0f; #pragma unroll for (int n = 0; n < 4; ++n) { From 63d15071286409828826ca2db338ac1e549736de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:52:32 +0200 Subject: [PATCH 21/41] cuda/window_attention: the split decode's slots outside the window hold zeros, not a copy of the window's first key, so an infinite value there reaches no sum --- zig/kernels/cuda/window_attention.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index 4e009fc669..695ed8efd5 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -51,8 +51,8 @@ extern "C" __global__ void __launch_bounds__(NT) tf_window_part(const __nv_bfloa const int key = t * BN + i / d; const bool in = key >= lo && key <= p; const int64_t at = ((base + (in ? key : lo) % ring) * hk + kh) * d + i % d; - kt[(i % d) * KT + i / d] = kc[at]; - vs[i] = vc[at]; + kt[(i % d) * KT + i / d] = in ? kc[at] : __float2bfloat16_rn(0.0f); // outside the window: zeros + vs[i] = in ? vc[at] : __float2bfloat16_rn(0.0f); } __syncthreads(); for (int i = threadIdx.x; i < g * BN; i += NT) { // score (head i / BN, key i % BN) From 40de67ddac1e8b2fba4b8e901296abc3d88c08eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:52:32 +0200 Subject: [PATCH 22/41] cuda/window_attention: prompt rows' keys a row cannot see add nothing, so a later row's infinite value stays out of earlier rows' sums A block whose output holds a NaN or infinity runs again, out of line, with the values of the tiles some row cannot see all of taken out of the mma; each row adds back the ones it sees. On finite values the first pass is the kernel as it was, byte for byte. --- zig/kernels/cuda/window_attention.cu | 58 +++++++++++++++++++++------- 1 file changed, 44 insertions(+), 14 deletions(-) diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index 695ed8efd5..e8619c872d 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -7,6 +7,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_window { @@ -164,16 +165,16 @@ __device__ __forceinline__ float rowsum(const float (&c)[8]) { __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); } -template -__global__ void __launch_bounds__(32 * WARPS, 1) -wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, - const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, - int G, float scale, int win, int ring) { +// ODD: the edge tiles' NaN and infinite values leave the mma, and the rows that see them add them back. +template +__device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, + int p0, int W, int H, int HK, int G, float scale, int win, int ring) { constexpr int C = D / 8, K16 = D / 16, NB = D / 8, RB = WARPS / HPC, THREADS = 32 * WARPS; constexpr bool QREG = D <= 128; // the queries' fragments stay in registers, else in shared memory - extern __shared__ uint4 smem[]; - uint4* slots = smem; - uint4* qs = smem + NS * HALF * C; + extern __shared__ uint4 wsmem[]; + uint4* slots = wsmem; + uint4* qs = wsmem + NS * HALF * C; const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g = lane / 4, qd = lane % 4, mt = lane / 8; const int groups = G / HPC, hk = blockIdx.y / groups, hg = blockIdx.y % groups; const int row0 = blockIdx.x * 16 * RB, last = min(row0 + 16 * RB, W) - 1; @@ -224,8 +225,9 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric uint32_t pa[4][4]; float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; const int pos0 = p0 + wrow + g, pos1 = pos0 + 8; - auto take = [&](int it) -> const uint4* { // item it's slot, once every thread's copy has landed + auto take = [&](int it, bool odd) -> const uint4* { // item it's slot, once every thread's copy has landed wait_group(); + if (odd) attn_odd::clean(slots + (it % NS) * HALF * C, THREADS); // NaN and infinities out of the mma __syncthreads(); return slots + (it % NS) * HALF * C; }; @@ -273,11 +275,14 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric mma(o[i + 1], pa1, b[2], b[3]); } }; + const auto edge = [&](int t) { // tile t, where some row of the block cannot see some key + return (t0 + t) * BN < (win > 0 ? p0 + last - win + 1 : 0) || (t0 + t) * BN + BN - 1 > p0 + row0; + }; for (int t = 0; t < tiles; ++t) { const int it = 4 * t; - scores(*reinterpret_cast(&sc[0]), take(it)); + scores(*reinterpret_cast(&sc[0]), take(it, false)); done(it); - scores(*reinterpret_cast(&sc[4]), take(it + 1)); + scores(*reinterpret_cast(&sc[4]), take(it + 1, false)); done(it + 1); { // _tile's fold of the tile's 64 scores const int kb = (t0 + t) * BN + 2 * qd; @@ -331,13 +336,21 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); } } - values(pa[0], pa[1], take(it + 2)); + values(pa[0], pa[1], take(it + 2, ODD && edge(t))); done(it + 2); - values(pa[2], pa[3], take(it + 3)); + values(pa[2], pa[3], take(it + 3, ODD && edge(t))); done(it + 3); } wait_group<0>(); - const int r0 = wrow + g, r1 = r0 + 8; + const int r0 = wrow + g, r1 = r0 + 8, pw = p0 + wrow; + if constexpr (!ODD) { // a NaN or infinite output: the block runs again as ODD (finite values: the same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(o, r0 < W, r1 < W))) return true; + } else { + const auto row = [&](int k) { return vc + ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D; }; + const auto vis = [&](int r, int k) { return k >= (win > 0 ? pw + r - win + 1 : 0) && k <= pw + r; }; + for (int t = 0; t < tiles; ++t) + if (edge(t)) attn_odd::add_back(o, row, (t0 + t) * BN, min((t0 + t + 1) * BN, limit), lane, vis, l0, l1); + } #pragma unroll for (int i = 0; i < NB; ++i) { const int d = 8 * i + 2 * qd; @@ -348,8 +361,25 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric *reinterpret_cast(out + ((int64_t)r1 * H + head) * D + d) = pack(tdiv(o[i][2], l1), tdiv(o[i][3], l1)); } + return false; } +// A block with a NaN or infinite output, run again with ODD on (out of the common path). +template +__device__ __noinline__ void wattn_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, + __nv_bfloat16* out, int p0, int W, int H, int HK, int G, float scale, int win, + int ring) { + wattn_rows(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring); +} + +template +__global__ void __launch_bounds__(32 * WARPS, 1) +wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, int W, int H, int HK, + int G, float scale, int win, int ring) { + if (wattn_rows(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring)) + wattn_odd(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring); +} // Decode on tensor cores (d 128): a row's heads one mma tile; warps take every 4th 32-key tile, fold in warp order. extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( From 8f93377333d69ae2c424e6293fce98057090b382 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:52:39 +0200 Subject: [PATCH 23/41] cuda/window_attention, moe_glue, dense_rows: registered with the kernel contract --- zig/src/core/contract/registry.zig | 13 + zig/tests/cuda/contract_moe.zig | 378 +++++++++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 26 +- zig/tests/cuda/contract_window.zig | 225 +++++++++++++++++ 4 files changed, 641 insertions(+), 1 deletion(-) create mode 100644 zig/tests/cuda/contract_moe.zig create mode 100644 zig/tests/cuda/contract_window.zig diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b2cfbc0e18..ce90a057ea 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -15,6 +15,7 @@ const gemv_symbol = "_ZN12tf_lane_gemv11gemv_kernelILi64ELi64ELi8EEEvPK13__nv_bf const prompt_symbol = "_ZN11tf_prompt1613prompt_kernelILi{d}ELi128ELi128ELi2ELi2ELi2ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPviiiiii"; const qmmf_f32_symbol = "_ZN7tf_qmmf11qmmf_kernelILi0ELi{d}ELi64ELi1ELi4ELi4ELb1ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; const vision80_symbol = "_ZN20tf_prefill_attention13vision_kernelILi80ELi8EEEvPK13__nv_bfloat16PS1_iif"; +const wattn_symbol = "_ZN16tf_window_prompt12wattn_kernelILi128ELi8ELi{d}ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifii"; const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK13__nv_bfloat16PS1_iif"; /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. @@ -45,6 +46,18 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "window attention tf_window_part", .symbol = "tf_window_part", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 51, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 } } }, + .{ .launch = .{ .name = "window attention tf_window_merge", .symbol = "tf_window_merge", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "window attention tf_window_decode_tc", .symbol = "tf_window_decode_tc", .image = "window_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 166, .static = 0, .bound = 128 }, .{ .regs = 166, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 4, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{4}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 8, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{8}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "moe tf_moe_topk_sigmoid", .symbol = "tf_moe_topk_sigmoid", .image = "moe_glue", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 } } }, + .{ .launch = .{ .name = "moe tf_moe_topk_sigmoid_warp", .symbol = "tf_moe_topk_sigmoid_warp", .image = "moe_glue", .threads = 256 }, .usage = .{ .{ .regs = 63, .static = 0, .bound = 256 }, .{ .regs = 63, .static = 0, .bound = 256 }, .{ .regs = 63, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "moe tf_moe_combine", .symbol = "tf_moe_combine", .image = "moe_glue", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_1", .symbol = "tf_dense_rows_1", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_4", .symbol = "tf_dense_rows_4", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 45, .static = 0, .bound = 256 }, .{ .regs = 45, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_16", .symbol = "tf_dense_rows_16", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "dense rows tf_dense_rows_32", .symbol = "tf_dense_rows_32", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "tree attention tf_tree_paths", .symbol = "tf_tree_paths", .image = "tree_attention", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, diff --git a/zig/tests/cuda/contract_moe.zig b/zig/tests/cuda/contract_moe.zig new file mode 100644 index 0000000000..a49b79c14d --- /dev/null +++ b/zig/tests/cuda/contract_moe.zig @@ -0,0 +1,378 @@ +//! moe_glue.cu and dense_rows.cu (#798) under the contract on CUDA: sigmoid top-k, the slots' combine, dense rows. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const Dense = contract.dense.Dense; +const glue = cuda.moe_glue; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const pick_rule = "picks by fp32 logit + bias, larger first, ties to the lower id, a NaN never picked; weights sigmoid(logit)"; + +/// A routing's fixed form; shape: rows, experts (n) and top-k (k). +pub const Route = struct { name: []const u8, shared: bool, normalize: bool = false, scale: f32 = 1, warp: bool }; + +pub const routes = [_]Route{ + .{ .name = "cuda moe top-k tf_moe_topk_sigmoid_warp (Kolibri: 384 experts, top 6 + shared)", .shared = true, .warp = true }, + .{ .name = "cuda moe top-k tf_moe_topk_sigmoid_warp (256 experts, top 8 normalized x2.5)", .shared = false, .normalize = true, .scale = 2.5, .warp = true }, + .{ .name = "cuda moe top-k tf_moe_topk_sigmoid (block form, 513 to 4,096 experts)", .shared = true, .warp = false }, +}; + +pub fn topkSpec(r: Route) c.Spec { + const served: []const c.Shape = if (!r.warp) &.{ .{ .rows = 16, .n = 4096, .k = 64 }, .{ .rows = 33, .n = 513, .k = 6 } } else if (r.normalize) &.{ .{ .rows = 16, .n = 256, .k = 8 }, .{ .rows = 300, .n = 256, .k = 8 } } else &.{ .{ .rows = 16, .n = 384, .k = 6 }, .{ .rows = 300, .n = 384, .k = 6 } }; + const tiles: [3]c.Tile = if (r.warp) .{ .{ .size = 8 }, .{ .size = 32, .min = 8 }, .{} } else .{ .{}, .{ .size = 256, .min = 513 }, .{} }; + return .{ .name = r.name, .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pick_rule }, .tiles = tiles, .served = served }; +} + +pub const TopK = struct { + rig: rig_.Rig, + route: Route, + g: *const glue.MoeGlue, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + logits: std.ArrayList(f32) = .empty, + bias: std.ArrayList(f32) = .empty, + dev: [4]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *TopK) void { + self.logits.deinit(self.rig.gpu.gpa); + self.bias.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *TopK) g.Kernel { + return .{ .spec = topkSpec(self.route), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(TopK); + + fn slots(self: *const TopK) usize { + return self.shape.k + @intFromBool(self.route.shared); + } + + fn form(self: *const TopK) glue.Route { + return .{ .experts = self.shape.n, .topk = self.shape.k, .shared = self.route.shared, .normalize = self.route.normalize, .scale = self.route.scale }; + } + + pub fn prepare(self: *TopK, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const warp = shape.n <= 512 and shape.k <= 32; + if (warp != self.route.warp or shape.rows == 0 or shape.k == 0 or shape.k > shape.n or shape.n > glue.most_experts or shape.k > glue.most_topk) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + const ne = shape.n; + try self.logits.resize(a, shape.rows * ne); + try self.bias.resize(a, ne); + for (self.logits.items, 0..) |*l, i| { + l.* = rng.floatNorm(f32) * 2; + if (i % 7 == 3) l.* = self.logits.items[i - 1]; // equal neighbours: the tie rule decides + } + for (self.bias.items, 0..) |*b, i| b.* = if (i % 5 == 0) 0.25 else 0; + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, 3 }, .{ 1, 5 }, .{ 2, 7 }, .{ 3, 2 }, .{ 3, ne - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |p, v| if (p[0] < shape.rows) { + self.logits.items[p[0] * ne + p[1]] = v; + }; + } + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.logits.items)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.bias.items)); + self.dev[2] = try self.rig.output(shape.rows * self.slots() * 4); + self.dev[3] = try self.rig.output(shape.rows * self.slots() * 4); + return true; + } + + pub fn run(self: *TopK) !void { + try self.g.topk(self.rig.s, self.dev[0], self.dev[1], self.form(), self.dev[2], self.dev[3], self.shape.rows); + try self.rig.s.synchronize(); + } + + pub fn count(self: *TopK) usize { + return self.shape.rows * self.slots() * 8; + } + + pub fn reference(_: *TopK, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// Each row's picks by the stated rule exactly, weights within 2^-20 of float64, the shared slot last. + pub fn check(self: *TopK, _: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const out = try self.rig.outputBytes(); + const n = self.shape.rows * self.slots(); + const ne = self.shape.n; + const taken = try a.alloc(bool, ne); + defer a.free(taken); + for (0..self.shape.rows) |r| { + const lg = self.logits.items[r * ne ..][0..ne]; + @memset(taken, false); + var w64: [65]f64 = undefined; + var sum: f64 = 0; + for (0..self.shape.k) |k| { + var best: ?usize = null; + var bv: f32 = 0; + for (lg, 0..) |l, e| { + const v = l + self.bias.items[e]; + if (taken[e] or std.math.isNan(v)) continue; + if (best == null or v > bv) { + best = e; + bv = v; + } + } + const id = best orelse return try std.fmt.bufPrint(&self.msg, "row {d}: fewer than {d} numbers to pick", .{ r, self.shape.k }); + taken[id] = true; + const got = std.mem.readInt(u32, out[(r * self.slots() + k) * 4 ..][0..4], .little); + if (got != id) return try std.fmt.bufPrint(&self.msg, "row {d} pick {d} is {d} where the rule gives {d}", .{ r, k, got, id }); + w64[k] = 1 / (1 + @exp(-@as(f64, lg[id]))); + sum += w64[k]; + } + if (self.route.shared) w64[self.shape.k] = 1; + for (0..self.slots()) |k| { + const want = if (k == self.shape.k) 1 else (if (self.route.normalize) w64[k] / sum else w64[k]) * self.route.scale; + const got: f64 = @as(f32, @bitCast(std.mem.readInt(u32, out[(n + r * self.slots() + k) * 4 ..][0..4], .little))); + if (k == self.shape.k and std.mem.readInt(u32, out[(r * self.slots() + k) * 4 ..][0..4], .little) != ne) return try std.fmt.bufPrint(&self.msg, "row {d}: the shared slot is not expert {d}", .{ r, ne }); + if (!(@abs(got - want) <= @abs(want) * 0x1p-20 + 1e-30)) return try std.fmt.bufPrint(&self.msg, "row {d} slot {d}: weight {d:.8} where float64 gives {d:.8}", .{ r, k, got, want }); + } + } + return null; + } +}; + +/// The combine's kinds: y in, out; Kolibri combines fp32 expert rows into bf16. +pub const Combo = struct { name: []const u8, y: cuda.norms.Kind, out: cuda.norms.Kind, elem: c.Elem }; + +pub const combos = [_]Combo{ + .{ .name = "cuda moe combine tf_moe_combine (fp32 rows into bf16)", .y = .f32, .out = .bf16, .elem = .bf16 }, + .{ .name = "cuda moe combine tf_moe_combine (bf16 rows into bf16)", .y = .bf16, .out = .bf16, .elem = .bf16 }, + .{ .name = "cuda moe combine tf_moe_combine (fp32 rows into fp32)", .y = .f32, .out = .f32, .elem = .f32 }, + .{ .name = "cuda moe combine tf_moe_combine (fp16 rows into fp16, the generic loop)", .y = .f16, .out = .f16, .elem = .f16 }, +}; + +/// One fma a slot, in slot order. +fn slotChain(s: c.Shape) usize { + return s.k; +} + +/// Shape: rows, dims (n) and slots (k): Kolibri's 6 + shared over 2,560, eight + shared over 7,168. +pub fn combineSpec(x: Combo) c.Spec { + return .{ .name = x.name, .accum = .f32, .depth = slotChain, .out = x.elem, .tiles = .{ .{}, .{ .size = 256 }, .{ .size = 8 } }, .served = &.{ .{ .rows = 16, .n = 2560, .k = 7 }, .{ .rows = 300, .n = 2560, .k = 7 }, .{ .rows = 9, .n = 7168, .k = 9 } } }; +} + +pub const Combine = struct { + rig: rig_.Rig, + x: Combo, + g: *const glue.MoeGlue, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + y: std.ArrayList(f64) = .empty, + w: std.ArrayList(f32) = .empty, + dev: [3]u64 = @splat(0), + + pub fn deinit(self: *Combine) void { + self.y.deinit(self.rig.gpu.gpa); + self.w.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Combine) g.Kernel { + return .{ .spec = combineSpec(self.x), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Combine); + + fn put(kind: cuda.norms.Kind, bytes: []u8, i: usize, v: f32) f64 { + switch (kind) { + .f32 => std.mem.writeInt(u32, bytes[4 * i ..][0..4], @bitCast(v), .little), + .f16 => std.mem.writeInt(u16, bytes[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(v))), .little), + .bf16 => std.mem.writeInt(u16, bytes[2 * i ..][0..2], c.toBf16(v), .little), + } + return switch (kind) { + .f32 => v, + .f16 => @as(f16, @floatCast(v)), + .bf16 => c.fromBf16(c.toBf16(v)), + }; + } + + pub fn prepare(self: *Combine, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + const n = shape.rows * shape.k * shape.n; + try self.y.resize(a, n); + try self.w.resize(a, shape.rows * shape.k); + const esize: usize = if (self.x.y == .f32) 4 else 2; + const bytes = try a.alloc(u8, n * esize); + defer a.free(bytes); + for (self.w.items) |*x| x.* = if (fill == .stress) @floatFromInt(1 + 2 * rng.uintLessThan(u32, 512)) else rng.float(f32); + const h = shape.k / 2; + for (0..shape.rows) |r| for (0..shape.k) |k| for (0..shape.n) |d| { + const i = (r * shape.k + k) * shape.n + d; + var v = rng.floatNorm(f32) * (if (i % 97 == 0) @as(f32, 30) else 1); + if (fill == .stress) { + // integer products w y to 2^18, the second half of the slots cancelling the first (one nudged) + if (shape.k % 2 == 1 and k == shape.k - 1) v = 0 else if (k >= h) { + self.w.items[r * shape.k + k] = self.w.items[r * shape.k + k - h]; + v = (if (k == h) @as(f32, 1) else 0) - @as(f32, @floatCast(self.y.items[i - h * shape.n])); + } else v = @floatFromInt(1 + 2 * rng.uintLessThan(u32, 128)); + } + self.y.items[i] = put(self.x.y, bytes, i, v); + }; + if (fill == .nonfinite) { + const at = [_][3]usize{ .{ 0, 0, shape.n / 3 }, .{ 1, 1, shape.n / 2 }, .{ 2, 0, shape.n - 1 }, .{ 3, 0, 0 }, .{ 3, shape.k - 1, 0 } }; + const vals = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, vals) |p, v| if (p[0] < shape.rows and p[1] < shape.k) { + const i = (p[0] * shape.k + p[1]) * shape.n + p[2]; + self.y.items[i] = put(self.x.y, bytes, i, v); + }; + } + self.dev[0] = try self.rig.input(bytes); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.w.items)); + self.dev[2] = try self.rig.output(shape.rows * shape.n * c.size(self.x.elem)); + return true; + } + + pub fn run(self: *Combine) !void { + try self.g.combine(self.rig.s, self.dev[0], self.x.y, self.dev[1], self.dev[2], self.x.out, self.shape.rows, self.shape.k, self.shape.n); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Combine) usize { + return self.shape.rows * self.shape.n; + } + + fn terms(self: *Combine, r: usize, d: usize, t: []f64) void { + for (t, 0..) |*x, k| x.* = @as(f64, self.w.items[r * self.shape.k + k]) * self.y.items[(r * self.shape.k + k) * self.shape.n + d]; + } + + /// core/moe_ref.combine's sum, and the terms' magnitude. + pub fn reference(self: *Combine, want: []f64, mag: []f64) !void { + var t: [64]f64 = undefined; + for (0..self.shape.rows) |r| for (0..self.shape.n) |d| { + self.terms(r, d, t[0..self.shape.k]); + var s: f64 = 0; + var abs: f64 = 0; + for (t[0..self.shape.k]) |x| { + s += x; + abs += @abs(x); + } + want[r * self.shape.n + d] = s; + mag[r * self.shape.n + d] = abs; + }; + } + + pub fn simulate(self: *Combine, acc: c.Accum, order: g.Order, out: []f64) !void { + var t: [64]f64 = undefined; + for (0..self.shape.rows) |r| for (0..self.shape.n) |d| { + self.terms(r, d, t[0..self.shape.k]); + out[r * self.shape.n + d] = rows.sumIn(acc, order, t[0..self.shape.k]); + }; + } +}; + +/// A lane's every 256th octet in order (k / 32 fma), then the warp's 5-step tree. +fn laneChain(s: c.Shape) usize { + return s.k / 32 + 5; +} + +pub const dense_tiles = [_]usize{ 1, 4, 16, 32 }; + +/// Shape: rows, outputs (n) and k; each entry launches its own row tile (the wrapper picks it by row count). +pub fn denseSpec(i: usize) c.Spec { + const names = [_][]const u8{ "cuda dense rows tf_dense_rows_1", "cuda dense rows tf_dense_rows_4", "cuda dense rows tf_dense_rows_16", "cuda dense rows tf_dense_rows_32" }; + // Kolibri's router (384 x 2,560), an output head's odd 4,103 columns, 257 x 7,168; rows as the wrapper sends them + const m = [_][2]usize{ .{ 1, 1 }, .{ 4, 3 }, .{ 16, 9 }, .{ 32, 40 } }; + const served = [_][3]c.Shape{ + .{ .{ .rows = m[0][0], .n = 384, .k = 2560 }, .{ .rows = m[0][1], .n = 4103, .k = 2560 }, .{ .rows = m[0][0], .n = 257, .k = 7168 } }, + .{ .{ .rows = m[1][0], .n = 384, .k = 2560 }, .{ .rows = m[1][1], .n = 4103, .k = 2560 }, .{ .rows = m[1][0], .n = 257, .k = 7168 } }, + .{ .{ .rows = m[2][0], .n = 384, .k = 2560 }, .{ .rows = m[2][1], .n = 4103, .k = 2560 }, .{ .rows = m[2][0], .n = 257, .k = 7168 } }, + .{ .{ .rows = m[3][0], .n = 384, .k = 2560 }, .{ .rows = m[3][1], .n = 4103, .k = 2560 }, .{ .rows = m[3][0], .n = 257, .k = 7168 } }, + }; + return .{ .name = names[i], .accum = .f32, .depth = laneChain, .out = .f32, .tiles = .{ .{ .size = dense_tiles[i] }, .{ .size = 8 }, .{ .size = 256, .step = 8, .min = 8 } }, .served = &served[i] }; +} + +pub const DenseRows = struct { + rig: rig_.Rig, + tile: usize, + g: *const cuda.dense_rows.DenseRows, + d: ?Dense = null, + dev: [3]u64 = @splat(0), + + pub fn deinit(self: *DenseRows) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *DenseRows) g.Kernel { + return .{ .spec = denseSpec(self.tile), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(DenseRows); + + pub fn prepare(self: *DenseRows, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.k % 8 != 0 or shape.k == 0 or shape.n == 0 or shape.rows == 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 1); + const bits = try a.alloc(u16, d.w.len); + defer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed ^ 0x77); + const rng = prng.random(); + const k = shape.k; + for (bits) |*b| b.* = c.toBf16(if (fill == .stress) 0 else rng.float(f32) * 2 - 1); + if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + bits[j * k + i] = c.toBf16(v); + bits[j * k + k / 2 + i] = c.toBf16(-v); + }; + for (d.w, bits) |*w, b| w.* = c.fromBf16(b); + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(bits)); + self.dev[2] = try self.rig.output(d.outputs() * 4); + return true; + } + + /// The wrapper's launch, at this entry's row tile whatever the row count. + pub fn run(self: *DenseRows) !void { + const sh = self.d.?.shape; + var a: cuda.Args = .{}; + a.add(self.dev[0]); + a.add(@as(i64, @intCast(sh.k))); + a.add(self.dev[1]); + a.add(self.dev[2]); + a.add(@as(i64, @intCast(sh.n))); + for ([_]usize{ sh.rows, sh.n, sh.k }) |v| a.add(@as(i32, @intCast(v))); + const rt = dense_tiles[self.tile]; + const grid: cuda.Dim3 = .{ .x = @intCast((sh.n + 7) / 8), .y = @intCast((sh.rows + rt - 1) / rt), .z = 1 }; + try cuda.launch.launch(self.g.fns[self.tile], .{ .grid = grid, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *DenseRows) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *DenseRows, want: []f64, mag: []f64) !void { + return self.d.?.reference(want, mag); + } + + pub fn simulate(self: *DenseRows, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 8114049a57..032b24f962 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -10,6 +10,8 @@ const mm_k = @import("contract_mm.zig"); const rank_k = @import("contract_rank.zig"); const tree_k = @import("contract_tree_attn.zig"); const block_k = @import("contract_block_attn.zig"); +const window_k = @import("contract_window.zig"); +const moe_k = @import("contract_moe.zig"); const contract = core.contract; const g = contract.gates; @@ -23,7 +25,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "tree_attention", "block_attention" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "tree_attention", "block_attention", "window_attention", "moe_glue", "dense_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -112,6 +114,28 @@ pub fn run(gpu: check.Gpu) !void { var blocks: [2]block_k.Block = .{ .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = false }, .{ .rig = .{ .gpu = gpu, .s = s }, .ba = &block_attn, .causal = true } }; defer for (&blocks) |*x| x.deinit(); for (&blocks) |*x| try kernels.append(a, x.kernel()); + var window_fns = try cuda.window_attention.WindowAttention.load(gpu.d); + defer window_fns.unload(); + var windows: [window_k.cfgs.len]window_k.Window = undefined; + for (window_k.cfgs, &windows) |f, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .g = &window_fns }; + defer for (&windows) |*x| x.deinit(); + for (&windows) |*x| try kernels.append(a, x.kernel()); + var moe_fns = try cuda.moe_glue.MoeGlue.load(gpu.d); + defer moe_fns.unload(); + var topks: [moe_k.routes.len]moe_k.TopK = undefined; + for (moe_k.routes, &topks) |r, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .route = r, .g = &moe_fns }; + defer for (&topks) |*x| x.deinit(); + for (&topks) |*x| try kernels.append(a, x.kernel()); + var combines: [moe_k.combos.len]moe_k.Combine = undefined; + for (moe_k.combos, &combines) |m, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .x = m, .g = &moe_fns }; + defer for (&combines) |*x| x.deinit(); + for (&combines) |*x| try kernels.append(a, x.kernel()); + var dense_fns = try cuda.dense_rows.DenseRows.load(gpu.d); + defer dense_fns.unload(); + var denses: [moe_k.dense_tiles.len]moe_k.DenseRows = undefined; + for (&denses, 0..) |*x, i| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .tile = i, .g = &dense_fns }; + defer for (&denses) |*x| x.deinit(); + for (&denses) |*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.zig b/zig/tests/cuda/contract_window.zig new file mode 100644 index 0000000000..9108b3feb0 --- /dev/null +++ b/zig/tests/cuda/contract_window.zig @@ -0,0 +1,225 @@ +//! window_attention.cu (#798) under the contract on CUDA: decode and prompt rows over a ring of keys, head by head. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const wa = cuda.window_attention; + +/// Why the split decode (fp32 weights) runs attention_ref's bound: a score's rounding reaches its weight through exp. +const split_why = "a score's absolute rounding error is its weight's relative error: gates 1 and 3 hold outputs to attention_ref's bound"; + +/// A decode or prompt call's fixed heads, window and ring (the variant), and which kernels it runs. +pub const Cfg = struct { name: []const u8, h: usize, hk: usize, d: usize, win: usize, ring: usize, streams: usize = 1, prompt: bool = false, served: []const c.Shape }; + +pub const cfgs = [_]Cfg{ + // Kolibri's sliding layers (48 heads over 4, window 513): decode on tensor cores, prompt rows 4 heads a block + .{ .name = "cuda window decode tf_window_decode_tc, 48/4 heads, window 513, ring 640", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 640, .streams = 2, .served = &.{ .{ .rows = 16, .n = 48, .k = 3000 }, .{ .rows = 64, .n = 48, .k = 1000 } } }, + .{ .name = "cuda window decode tf_window_decode_tc, 16/1 heads, no window, ring 300", .h = 16, .hk = 1, .d = 128, .win = 0, .ring = 300, .served = &.{.{ .rows = 16, .n = 16, .k = 300 }} }, + .{ .name = "cuda window decode tf_window_part + tf_window_merge, 8/2 heads of 64, window 100, ring 128", .h = 8, .hk = 2, .d = 64, .win = 100, .ring = 128, .streams = 3, .served = &.{ .{ .rows = 16, .n = 8, .k = 3000 }, .{ .rows = 64, .n = 8, .k = 1000 } } }, + .{ .name = "cuda window prompt wattn_kernel<128, 8, 4, 8>, 48/4 heads, window 513, ring 2624", .h = 48, .hk = 4, .d = 128, .win = 513, .ring = 2624, .prompt = true, .served = &.{ .{ .rows = 128, .n = 48, .k = 0 }, .{ .rows = 96, .n = 48, .k = 2600 } } }, + .{ .name = "cuda window prompt wattn_kernel<128, 8, 8, 8>, 16/2 heads, no window, a row a position", .h = 16, .hk = 2, .d = 128, .win = 0, .ring = 0, .prompt = true, .served = &.{ .{ .rows = 64, .n = 16, .k = 100 }, .{ .rows = 48, .n = 16, .k = 0 } } }, +}; + +/// The P x V chain of a row: its window's keys, a tile at a time, then the tiles' fold. +fn keyChain(s: c.Shape) usize { + return s.k + 64 + 64; +} + +/// Decode: rows, query heads (n) and the positions' bound (k); prompt: chunk rows, query heads and the chunk's start. +pub fn spec(f: Cfg) c.Spec { + const key_tile: usize = if (f.prompt or f.d != 128) 64 else 32; + const row_tile: usize = if (f.prompt) 16 * (8 / (if ((f.h / f.hk) % 8 == 0) @as(usize, 8) else 4)) else 0; + const k_tile: c.Tile = if (f.prompt) .{ .size = 64, .min = 0 } else .{ .size = key_tile }; + return .{ .name = f.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = row_tile }, .{}, k_tile }, .served = f.served, .sums = .{ .why = if (f.d == 128) attn.why else split_why } }; +} + +pub const Window = struct { + rig: rig_.Rig, + f: Cfg, + g: *const wa.WindowAttention, + r: ?attn.Ref = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + pos: std.ArrayList(i32) = .empty, + slot: std.ArrayList(i32) = .empty, + q: std.ArrayList(u16) = .empty, + kc: std.ArrayList(u16) = .empty, + vc: std.ArrayList(u16) = .empty, + dev: [5]u64 = @splat(0), + scratch: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Window) void { + const a = self.rig.gpu.gpa; + if (self.r) |*r| r.deinit(); + inline for (.{ "pos", "slot", "q", "kc", "vc" }) |n| @field(self, n).deinit(a); + self.rig.deinit(); + } + + pub fn kernel(self: *Window) g.Kernel { + return .{ .spec = spec(self.f), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Window); + + fn win(self: *const Window) wa.Window { + const f = self.f; + return .{ .h = f.h, .hk = f.hk, .d = f.d, .win = f.win, .ring = f.ring, .scale = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(f.d)))) }; + } + + /// A row's position and its cache's first row (stream slot * ring; a prompt's cache holds one stream). + fn at(self: *const Window, r: usize) struct { p: usize, base: usize } { + if (self.f.prompt) return .{ .p = self.shape.k + r, .base = 0 }; + return .{ .p = @intCast(self.pos.items[r]), .base = @as(usize, @intCast(self.slot.items[r])) * self.f.ring }; + } + + fn cacheRow(self: *const Window, base: usize, key: usize) usize { + return base + if (self.f.ring > 0) key % self.f.ring else key; + } + + fn first(self: *const Window, p: usize) usize { + return if (self.f.win > 0 and p + 1 > self.f.win) p + 1 - self.f.win else 0; + } + + pub fn prepare(self: *Window, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const f = self.f; + if (shape.n != f.h or shape.rows == 0 or (!f.prompt and shape.k == 0)) return false; + if (!f.prompt and f.win == 0 and shape.k > f.ring) return false; // an unwindowed row's keys all fit the ring + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + const n = shape.rows; + try self.pos.resize(a, n); + try self.slot.resize(a, n); + for (self.pos.items, self.slot.items, 0..) |*p, *sl, 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); + } + const crow = if (f.prompt) (if (f.ring > 0) f.ring else shape.k + n) else f.streams * f.ring; + const kv = f.hk * f.d; + try self.q.resize(a, n * f.h * f.d); + try self.kc.resize(a, crow * kv); + try self.vc.resize(a, crow * kv); + for ([_][]u16{ self.q.items, self.kc.items, self.vc.items }) |x| for (x) |*e| { + e.* = attn.normal(rng); + }; + var most: usize = 1; + for (0..n) |r| { + const w = self.at(r); + most = @max(most, w.p + 1 - self.first(w.p)); + } + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, f.d, most); + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng), + } + for ([_][]const u8{ std.mem.sliceAsBytes(self.q.items), std.mem.sliceAsBytes(self.kc.items), std.mem.sliceAsBytes(self.vc.items), std.mem.sliceAsBytes(self.pos.items), std.mem.sliceAsBytes(self.slot.items) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + self.scratch = if (f.prompt) 0 else try self.rig.scratch(wa.WindowAttention.scratchBytes(self.win(), n)); + self.out = try self.rig.output(n * f.h * f.d * 2); + return true; + } + + fn cell(self: *Window, x: []u16, r: usize, key: usize, kh: usize, e: usize) *u16 { + const w = self.at(r); + return &x[(self.cacheRow(w.base, key) * self.f.hk + kh) * self.f.d + e]; + } + + /// NaN in a query; +inf and -inf in values; both in one column; NaN in a key; +inf in a window's first value. + fn plant(self: *Window) void { + const n = self.shape.rows; + const nan = c.toBf16(std.math.nan(f32)); + const inf = c.toBf16(std.math.inf(f32)); + const ninf = c.toBf16(-std.math.inf(f32)); + const hk = self.f.hk; + self.q.items[3] = nan; // row 0, head 0 + self.cell(self.vc.items, 0, self.at(0).p, 0, 7).* = inf; + if (n > 1) self.cell(self.vc.items, 1, self.at(1).p, (1 % hk), 9).* = ninf; + if (n > 2 and self.at(2).p >= self.first(self.at(2).p) + 1) { + self.cell(self.vc.items, 2, self.at(2).p, (2 % hk), 11).* = inf; + self.cell(self.vc.items, 2, self.at(2).p - 1, (2 % hk), 11).* = ninf; + } + if (n > 3) self.cell(self.kc.items, 3, self.at(3).p, (3 % hk), 0).* = nan; + if (n > 4) self.cell(self.vc.items, 4, self.first(self.at(4).p), (4 % hk), 13).* = inf; + } + + /// Gate 4: every key alike (each p is 1); values mirrored over the keys every row sees, 0 elsewhere. + fn stress(self: *Window, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const f = self.f; + const kv = f.hk * f.d; + const key = try a.alloc(u16, kv); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + for (0..self.kc.items.len / kv) |i| @memcpy(self.kc.items[i * kv ..][0..kv], key); + @memset(self.vc.items, 0); + const last = self.at(self.shape.rows - 1).p; + const lo = self.first(last); + const hi = self.at(0).p; + const col = try a.alloc(f32, hi + 1 - lo); + defer a.free(col); + for (0..f.hk) |kh| for (0..f.d) |e| { + attn.mirrored(rng, col); + for (col, lo..) |y, k| self.cell(self.vc.items, 0, k, kh, e).* = c.toBf16(y); + }; + } + + pub fn run(self: *Window) !void { + const n = self.shape.rows; + if (self.f.prompt) { + try self.g.prompt(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.out, self.shape.k, n, self.win()); + } else try self.g.decode(self.rig.s, self.dev[0], self.dev[1], self.dev[2], self.out, self.dev[3], self.dev[4], self.win(), self.scratch, n); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Window) usize { + return self.shape.rows * self.f.h * self.f.d; + } + + /// Output head i (row i / h, head i % h): keys (p - win, p] of its stream's ring. + pub fn load(self: *Window, i: usize, hd: *attn.Head) !void { + const f = self.f; + const d = f.d; + const r = i / f.h; + const head = i % f.h; + const kh = head / (f.h / f.hk); + const w = self.at(r); + const lo = self.first(w.p); + for (hd.q, 0..) |*x, e| x.* = c.fromBf16(self.q.items[(r * f.h + head) * d + e]); + for (lo..w.p + 1, 0..) |key, j| for (0..d) |e| { + const cr = (self.cacheRow(w.base, key) * f.hk + kh) * d + e; + hd.k[j * d + e] = c.fromBf16(self.kc.items[cr]); + hd.v[j * d + e] = c.fromBf16(self.vc.items[cr]); + }; + hd.n = w.p + 1 - lo; + } + + fn heads(self: *Window) usize { + return self.shape.rows * self.f.h; + } + + fn scale(self: *const Window) f64 { + return self.win().scale; + } + + pub fn check(self: *Window, gate: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), self.scale(), try self.rig.outputBytes(), gate); + } + + pub fn reference(self: *Window, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), self.scale(), want, mag); + } + + pub fn simulate(self: *Window, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), self.scale(), acc, order, out); + } +}; From 3c7a1860f19094bb489f2a7c7fd7ebdc95be2c4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:53:25 +0200 Subject: [PATCH 24/41] kolibri1: config checks, the tensor map and the loader on the shared block-FP8 layouts (cuda/fp8.zig, qlinear, fp8g experts) --- zig/src/families/kolibri1/config.zig | 305 ++++++++++++++++++++++++++ zig/src/families/kolibri1/cuda.zig | 13 ++ zig/src/families/kolibri1/names.zig | 212 ++++++++++++++++++ zig/src/families/kolibri1/pack.zig | 165 ++++++++++++++ zig/src/families/kolibri1/weights.zig | 100 +++++++++ 5 files changed, 795 insertions(+) create mode 100644 zig/src/families/kolibri1/config.zig create mode 100644 zig/src/families/kolibri1/cuda.zig create mode 100644 zig/src/families/kolibri1/names.zig create mode 100644 zig/src/families/kolibri1/pack.zig create mode 100644 zig/src/families/kolibri1/weights.zig diff --git a/zig/src/families/kolibri1/config.zig b/zig/src/families/kolibri1/config.zig new file mode 100644 index 0000000000..bf3a29a091 --- /dev/null +++ b/zig/src/families/kolibri1/config.zig @@ -0,0 +1,305 @@ +//! Kolibri-1's dimensions from config.json, and the layouts this engine refuses (FP8 128 x 128, sigmoid routing). +const std = @import("std"); + +pub const max_layers = 64; +pub const max_eos = 4; + +/// Why a check refused the checkpoint, kept for the caller's log line (tests read it instead). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + /// Keep the reason, cut short if it is long. + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +pub const Config = struct { + layers: usize, + hidden: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + vocab: usize, + experts: usize, // routed; the shared expert comes after them + top_k: usize, + moe_width: usize, + shared_width: usize, + window: usize, // sliding layers attend to this many keys back, the row's own included + full: [max_layers]bool = @splat(false), // full attention takes no positions; sliding layers take RoPE + rope_theta: f32, + eps: f32, + eos: [max_eos]u32 = @splat(0), + eos_count: usize = 0, + + pub fn qkvDim(self: Config) usize { + return (self.heads + 2 * self.kv_heads) * self.head_dim; + } + + pub fn fullCount(self: Config) usize { + var n: usize = 0; + for (self.full[0..self.layers]) |f| n += @intFromBool(f); + return n; + } + + pub fn isEos(self: Config, token: u32) bool { + for (self.eos[0..self.eos_count]) |e| if (e == token) return true; + return false; + } + + /// End tokens from generation_config.json's text, after config.json's own: replies end on either. + pub fn addEos(self: *Config, allocator: std.mem.Allocator, generation: []const u8) !void { + var parsed = try std.json.parseFromSlice(std.json.Value, allocator, generation, .{}); + defer parsed.deinit(); + if (parsed.value.object.get("eos_token_id")) |v| try eosFrom(self, v); + } + + /// `dir`/config.json, checked for the kernels' shapes, and generation_config.json's end tokens when it is there. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(text); + var why: Why = .{}; + var c = parse(gpa, text, &why) catch |e| { + if (why.len > 0) std.log.err("{s}: {s}", .{ path, why.text() }); + return e; + }; + checkShapes(c, &why) catch |e| { + std.log.err("{s}: {s}", .{ path, why.text() }); + return e; + }; + const gen_path = try std.fs.path.join(gpa, &.{ dir, "generation_config.json" }); + defer gpa.free(gen_path); + const gen = std.Io.Dir.cwd().readFileAlloc(io, gen_path, gpa, .limited(1 << 20)) catch |e| switch (e) { + error.FileNotFound => return c, + else => return e, + }; + defer gpa.free(gen); + try c.addEos(gpa, gen); + return c; + } +}; + +fn int(obj: std.json.ObjectMap, key: []const u8, why: *Why) !usize { + const v = obj.get(key) orelse { + why.set("config.json has no {s}", .{key}); + return error.BadConfig; + }; + if (v != .integer or v.integer <= 0) { + why.set("config.json's {s} is not a positive integer", .{key}); + return error.BadConfig; + } + return @intCast(v.integer); +} + +fn number(v: std.json.Value) !f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => error.BadConfig, + }; +} + +/// One end token, or a list of them; null entries are skipped and a token already listed is kept once. +fn eosFrom(c: *Config, v: std.json.Value) !void { + const items: []const std.json.Value = switch (v) { + .integer => &.{v}, + .array => |a| a.items, + else => return, + }; + for (items) |x| { + if (x != .integer) continue; + const t: u32 = @intCast(x.integer); + if (c.isEos(t)) continue; + if (c.eos_count == max_eos) return error.TooManyEos; + c.eos[c.eos_count] = t; + c.eos_count += 1; + } +} + +/// config.json's text; a refusal's reason goes to `why`. +pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); + defer parsed.deinit(); + const o = parsed.value.object; + const mt = o.get("model_type") orelse return error.BadConfig; + if (mt != .string or !std.mem.eql(u8, mt.string, "kolibri1")) { + why.set("model_type {f} is not kolibri1", .{std.json.fmt(mt, .{})}); + return error.NotKolibri1; + } + const q = if (o.get("quantization_config")) |v| (if (v == .object) v.object else null) else null; + const fp8 = if (q) |m| (if (m.get("quant_method")) |x| x == .string and std.mem.eql(u8, x.string, "fp8") else false) else false; + const blocks = if (q) |m| (if (m.get("weight_block_size")) |x| x == .array and x.array.items.len == 2 and + x.array.items[0] == .integer and x.array.items[0].integer == 128 and + x.array.items[1] == .integer and x.array.items[1].integer == 128 else false) else false; + if (!fp8 or !blocks) { + why.set("Kolibri-1 on CUDA reads its block-FP8 checkpoint (128 x 128 blocks, fp32 scales)", .{}); + return error.UnsupportedQuantization; + } + if (o.get("vision_config") != null) { + why.set("config.json declares a vision tower, and Kolibri-1 on CUDA has no image input yet", .{}); + return error.VisionTower; + } + if (o.get("norm_topk_prob")) |v| if (v == .bool and v.bool) { + why.set("Kolibri-1 routes by unnormalised sigmoid weights; norm_topk_prob is not supported", .{}); + return error.UnsupportedRouting; + }; + var c = Config{ + .layers = try int(o, "num_hidden_layers", why), + .hidden = try int(o, "hidden_size", why), + .heads = try int(o, "num_attention_heads", why), + .kv_heads = try int(o, "num_key_value_heads", why), + .head_dim = try int(o, "head_dim", why), + .vocab = try int(o, "vocab_size", why), + .experts = try int(o, "num_experts", why), + .top_k = try int(o, "num_experts_per_tok", why), + .moe_width = try int(o, "moe_intermediate_size", why), + .shared_width = try int(o, "shared_expert_intermediate_size", why), + .window = try int(o, "sliding_window", why), + .rope_theta = 10000.0, + .eps = @floatCast(try number(o.get("rms_norm_eps") orelse return error.BadConfig)), + }; + if (c.layers > max_layers) { + why.set("{d} layers; this build holds at most {d}", .{ c.layers, max_layers }); + return error.BadConfig; + } + // rope_parameters.rope_theta wins over the top-level rope_theta, as transformers reads them + if (o.get("rope_theta")) |v| c.rope_theta = @floatCast(try number(v)); + if (o.get("rope_parameters")) |v| if (v == .object) if (v.object.get("rope_theta")) |t| { + c.rope_theta = @floatCast(try number(t)); + }; + const types = o.get("layer_types") orelse { + why.set("config.json has no layer_types", .{}); + return error.BadConfig; + }; + if (types != .array or types.array.items.len != c.layers) { + why.set("layer_types does not name each of the {d} layers", .{c.layers}); + return error.BadConfig; + } + for (types.array.items, 0..) |t, i| { + if (t != .string) return error.BadConfig; + if (std.mem.eql(u8, t.string, "full_attention")) { + c.full[i] = true; + } else if (!std.mem.eql(u8, t.string, "sliding_attention")) { + why.set("layer {d} is {s}; Kolibri-1 has full_attention and sliding_attention layers", .{ i, t.string }); + return error.UnsupportedLayer; + } + } + if (o.get("eos_token_id")) |v| try eosFrom(&c, v); + return c; +} + +/// What the shared kernels take: head dim 128, query heads a key head in fours up to 16, whole 128 blocks, one width. +pub fn checkShapes(c: Config, why: *Why) !void { + const group = if (c.kv_heads > 0 and c.heads % c.kv_heads == 0) c.heads / c.kv_heads else 0; + if (c.head_dim != 128 or group == 0 or group % 4 != 0 or group > 16) { + why.set("head dim {d}, {d} query heads over {d} key heads: this engine takes head dim 128 and 4, 8, 12 or 16 query heads a key head", .{ c.head_dim, c.heads, c.kv_heads }); + return error.UnsupportedShapes; + } + if (c.hidden % 128 != 0 or c.moe_width % 128 != 0 or c.shared_width != c.moe_width) { + why.set("hidden {d}, experts {d} wide, shared expert {d}: this engine takes whole 128 blocks and one expert width", .{ c.hidden, c.moe_width, c.shared_width }); + return error.UnsupportedShapes; + } + if (c.experts > 4096 or c.top_k == 0 or c.top_k + 1 > 64 or c.top_k > c.experts or c.window < 2) { + why.set("{d} experts, top {d}, window {d}: outside what the routing and window kernels take", .{ c.experts, c.top_k, c.window }); + return error.UnsupportedShapes; + } +} + +/// Kolibri-1's config.json at the pinned revision (every fifth layer full), its long module list cut to two entries. +pub const pinned = + \\{"model_type": "kolibri1", "hidden_size": 2560, "num_hidden_layers": 50, "num_attention_heads": 48, + \\ "num_key_value_heads": 4, "head_dim": 128, "rms_norm_eps": 1e-06, "vocab_size": 128000, "rope_theta": 10000.0, + \\ "num_experts": 384, "num_experts_per_tok": 6, "moe_intermediate_size": 512, "shared_expert_intermediate_size": 512, + \\ "norm_topk_prob": false, "sliding_window": 513, "eos_token_id": 127906, + \\ "layer_types": ["sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", + \\ "full_attention"], + \\ "quantization_config": {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], + \\ "modules_to_not_convert": ["model.layers.0.mlp.gate", "model.layers.1.mlp.gate"]}} +; + +fn refused(text: []const u8, err: anyerror, word: []const u8) !void { + var why: Why = .{}; + try std.testing.expectError(err, parse(std.testing.allocator, text, &why)); + if (std.mem.indexOf(u8, why.text(), word) == null) { + std.debug.print("refusal \"{s}\" does not name \"{s}\"\n", .{ why.text(), word }); + return error.TestUnexpectedResult; + } +} + +fn replaced(from: []const u8, to: []const u8) ![]u8 { + return std.mem.replaceOwned(u8, std.testing.allocator, pinned, from, to); +} + +test "parse the pinned config, its shapes and end tokens" { + var why: Why = .{}; + var c = try parse(std.testing.allocator, pinned, &why); + try checkShapes(c, &why); + try std.testing.expectEqual(@as(usize, 50), c.layers); + try std.testing.expect(c.full[4] and c.full[49] and !c.full[0] and !c.full[48]); + try std.testing.expectEqual(@as(usize, 7168), c.qkvDim()); + try std.testing.expectEqual(@as(f32, 1e-6), c.eps); + try std.testing.expect(c.isEos(127906) and !c.isEos(127901)); + try c.addEos(std.testing.allocator, "{\"eos_token_id\": [127906, 127901], \"top_k\": 128}"); + try std.testing.expectEqual(@as(usize, 2), c.eos_count); + try std.testing.expect(c.isEos(127901)); +} + +test "refusals: not kolibri1, other quantization, normalised routing, other layer kinds, too many layers" { + const not = try replaced("\"model_type\": \"kolibri1\"", "\"model_type\": \"qwen3_moe\""); + defer std.testing.allocator.free(not); + try refused(not, error.NotKolibri1, "kolibri1"); + const blocks = try replaced("[128, 128]", "[64, 64]"); + defer std.testing.allocator.free(blocks); + try refused(blocks, error.UnsupportedQuantization, "128 x 128"); + const method = try replaced("\"quant_method\": \"fp8\"", "\"quant_method\": \"awq\""); + defer std.testing.allocator.free(method); + try refused(method, error.UnsupportedQuantization, "FP8"); + const norm = try replaced("\"norm_topk_prob\": false", "\"norm_topk_prob\": true"); + defer std.testing.allocator.free(norm); + try refused(norm, error.UnsupportedRouting, "norm_topk_prob"); + const layer = try replaced("\"full_attention\"]", "\"linear_attention\"]"); + defer std.testing.allocator.free(layer); + try refused(layer, error.UnsupportedLayer, "linear_attention"); + const tower = try replaced("\"model_type\": \"kolibri1\",", "\"model_type\": \"kolibri1\", \"vision_config\": {},"); + defer std.testing.allocator.free(tower); + try refused(tower, error.VisionTower, "vision tower"); + const short = try replaced("\"num_hidden_layers\": 50", "\"num_hidden_layers\": 49"); + defer std.testing.allocator.free(short); + try refused(short, error.BadConfig, "49 layers"); +} + +test "rope_parameters' theta wins, and other shapes are refused before any kernel runs" { + const rope = try replaced("\"rope_theta\": 10000.0", "\"rope_theta\": 10000.0, \"rope_parameters\": {\"rope_theta\": 5e5}"); + defer std.testing.allocator.free(rope); + var why: Why = .{}; + const c = try parse(std.testing.allocator, rope, &why); + try std.testing.expectEqual(@as(f32, 5e5), c.rope_theta); + try checkShapes(c, &why); + var narrow = c; + narrow.head_dim = 64; + try std.testing.expectError(error.UnsupportedShapes, checkShapes(narrow, &why)); + var odd = c; + odd.kv_heads = 8; // six query heads a key head + try std.testing.expectError(error.UnsupportedShapes, checkShapes(odd, &why)); + var wide = c; + wide.shared_width = 1024; + try std.testing.expectError(error.UnsupportedShapes, checkShapes(wide, &why)); +} diff --git a/zig/src/families/kolibri1/cuda.zig b/zig/src/families/kolibri1/cuda.zig new file mode 100644 index 0000000000..0ceab59ddc --- /dev/null +++ b/zig/src/families/kolibri1/cuda.zig @@ -0,0 +1,13 @@ +//! Kolibri-1 (block FP8) on CUDA: config, weights, the layer graph on shared operators, the lane backend. + +pub const config = @import("config.zig"); +pub const names = @import("names.zig"); +pub const pack = @import("pack.zig"); +pub const weights = @import("weights.zig"); + +test { + _ = @import("config.zig"); + _ = @import("names.zig"); + _ = @import("pack.zig"); + _ = @import("weights.zig"); +} diff --git a/zig/src/families/kolibri1/names.zig b/zig/src/families/kolibri1/names.zig new file mode 100644 index 0000000000..bec7f25ab6 --- /dev/null +++ b/zig/src/families/kolibri1/names.zig @@ -0,0 +1,212 @@ +//! Kolibri-1's checkpoint tensors by name, each checked for dtype and shape before any upload. +const std = @import("std"); +const core = @import("core"); +const Config = @import("config.zig").Config; + +const Tensor = core.checkpoint.Tensor; +const Checkpoint = core.checkpoint.Checkpoint; + +pub const block = 128; // the checkpoint's scale block, both ways + +/// A block-FP8 projection as stored: e4m3 codes [n, k] and fp32 scale_inv [ceil(n/128), ceil(k/128)]. +pub const Fp8 = struct { w: Tensor, s: Tensor }; +pub const Expert = struct { gate: Fp8, up: Fp8, down: Fp8 }; + +pub const Layer = struct { + input_norm: Tensor, + q: Fp8, + k: Fp8, + v: Fp8, + q_norm: Tensor, + k_norm: Tensor, + o: Fp8, + post_attn_norm: Tensor, + pre_moe_norm: Tensor, // post_attention_layernorm: the norm before the MoE, despite its name + router: Tensor, // [E, D] bf16 + bias: Tensor, // [E] expert_bias + experts: []Expert, // E routed, then the shared expert + post_moe_norm: Tensor, +}; + +pub const Map = struct { + gpa: std.mem.Allocator, + embed: Tensor, + layers: []Layer, + norm: Tensor, + head: Tensor, + + pub fn deinit(m: *Map) void { + for (m.layers) |l| m.gpa.free(l.experts); + m.gpa.free(m.layers); + m.* = undefined; + } +}; + +fn blocks(n: usize) usize { + return (n + block - 1) / block; +} + +const Names = struct { + ck: *Checkpoint, + buf: [160]u8 = undefined, + + fn name(self: *Names, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(&self.buf, fmt, args); + } + + fn bf16(self: *Names, shape: []const usize, comptime fmt: []const u8, args: anytype) !Tensor { + return self.ck.expect(try self.name(fmt, args), .bf16, shape); + } + + fn fp8(self: *Names, n: usize, k: usize, comptime fmt: []const u8, args: anytype) !Fp8 { + const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .f8_e4m3, &.{ n, k }); + const s = try self.ck.expect(try self.name(fmt ++ ".weight_scale_inv", args), .f32, &.{ blocks(n), blocks(k) }); + return .{ .w = w, .s = s }; + } + + fn expert(self: *Names, c: Config, comptime pre: []const u8, args: anytype) !Expert { + return .{ + .gate = try self.fp8(c.moe_width, c.hidden, pre ++ ".gate_proj", args), + .up = try self.fp8(c.moe_width, c.hidden, pre ++ ".up_proj", args), + .down = try self.fp8(c.hidden, c.moe_width, pre ++ ".down_proj", args), + }; + } +}; + +/// Every tensor the loader reads, in its order; a missing, mistyped or leftover tensor refuses the checkpoint. +pub fn map(gpa: std.mem.Allocator, ck: *Checkpoint, c: Config) !Map { + if (c.shared_width != c.moe_width) { + std.log.err("shared expert {d} wide, routed {d}: the experts stack only at one width", .{ c.shared_width, c.moe_width }); + return error.UnsupportedShapes; + } + var nm: Names = .{ .ck = ck }; + const d = c.hidden; + var m: Map = .{ .gpa = gpa, .embed = undefined, .layers = try gpa.alloc(Layer, c.layers), .norm = undefined, .head = undefined }; + var done: usize = 0; + errdefer { + for (m.layers[0..done]) |l| gpa.free(l.experts); + gpa.free(m.layers); + } + m.embed = try nm.bf16(&.{ c.vocab, d }, "model.embed_tokens.weight", .{}); + for (m.layers, 0..) |*l, i| { + const ex = try gpa.alloc(Expert, c.experts + 1); + errdefer gpa.free(ex); + for (ex[0..c.experts], 0..) |*e, j| e.* = try nm.expert(c, "model.layers.{d}.mlp.experts.{d}", .{ i, j }); + ex[c.experts] = try nm.expert(c, "model.layers.{d}.mlp.shared_experts", .{i}); + const qd = c.heads * c.head_dim; + const kvd = c.kv_heads * c.head_dim; + l.* = .{ + .input_norm = try nm.bf16(&.{d}, "model.layers.{d}.input_layernorm.weight", .{i}), + .q = try nm.fp8(qd, d, "model.layers.{d}.self_attn.q_proj", .{i}), + .k = try nm.fp8(kvd, d, "model.layers.{d}.self_attn.k_proj", .{i}), + .v = try nm.fp8(kvd, d, "model.layers.{d}.self_attn.v_proj", .{i}), + .q_norm = try nm.bf16(&.{c.head_dim}, "model.layers.{d}.self_attn.q_norm.weight", .{i}), + .k_norm = try nm.bf16(&.{c.head_dim}, "model.layers.{d}.self_attn.k_norm.weight", .{i}), + .o = try nm.fp8(d, qd, "model.layers.{d}.self_attn.o_proj", .{i}), + .post_attn_norm = try nm.bf16(&.{d}, "model.layers.{d}.post_attn_norm.weight", .{i}), + .pre_moe_norm = try nm.bf16(&.{d}, "model.layers.{d}.post_attention_layernorm.weight", .{i}), + .router = try nm.bf16(&.{ c.experts, d }, "model.layers.{d}.mlp.gate.weight", .{i}), + .bias = try nm.bf16(&.{c.experts}, "model.layers.{d}.moe.router.expert_bias", .{i}), + .experts = ex, + .post_moe_norm = try nm.bf16(&.{d}, "model.layers.{d}.post_ffn_norm.weight", .{i}), + }; + done += 1; + } + m.norm = try nm.bf16(&.{d}, "model.norm.weight", .{}); + m.head = try nm.bf16(&.{ c.vocab, d }, "lm_head.weight", .{}); + if (ck.unused() != 0) return error.UnusedTensors; + return m; +} + +const testing = std.testing; +const io = testing.io; + +const tiny: Config = .{ .layers = 2, .hidden = 256, .heads = 2, .kv_heads = 1, .head_dim = 128, .vocab = 8, .experts = 2, .top_k = 1, .moe_width = 128, .shared_width = 128, .window = 4, .rope_theta = 1e4, .eps = 1e-6 }; + +/// A zero-filled single-file checkpoint with each entry's dtype and shape. +fn writeTiny(a: std.mem.Allocator, tmp: testing.TmpDir) !void { + var entries: std.ArrayList(struct { name: []const u8, dtype: []const u8, shape: []const usize }) = .empty; + const c = tiny; + const Add = struct { + fn fp8(list: anytype, al: std.mem.Allocator, pre: []const u8, n: usize, k: usize) !void { + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k }) }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_inv", .{pre}), .dtype = "F32", .shape = try al.dupe(usize, &.{ blocks(n), blocks(k) }) }); + } + fn bf(list: anytype, al: std.mem.Allocator, name: []const u8, shape: []const usize) !void { + try list.append(al, .{ .name = name, .dtype = "BF16", .shape = try al.dupe(usize, shape) }); + } + }; + const d = c.hidden; + try Add.bf(&entries, a, "model.embed_tokens.weight", &.{ c.vocab, d }); + try Add.bf(&entries, a, "model.norm.weight", &.{d}); + try Add.bf(&entries, a, "lm_head.weight", &.{ c.vocab, d }); + for (0..c.layers) |i| { + const p = try std.fmt.allocPrint(a, "model.layers.{d}.", .{i}); + for ([_][]const u8{ "input_layernorm", "post_attn_norm", "post_attention_layernorm", "post_ffn_norm" }) |n| + try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}{s}.weight", .{ p, n }), &.{d}); + for ([_][]const u8{ "q_norm", "k_norm" }) |n| try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}self_attn.{s}.weight", .{ p, n }), &.{c.head_dim}); + try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}mlp.gate.weight", .{p}), &.{ c.experts, d }); + try Add.bf(&entries, a, try std.fmt.allocPrint(a, "{s}moe.router.expert_bias", .{p}), &.{c.experts}); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}self_attn.q_proj", .{p}), c.heads * c.head_dim, d); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}self_attn.k_proj", .{p}), c.kv_heads * c.head_dim, d); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}self_attn.v_proj", .{p}), c.kv_heads * c.head_dim, d); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}self_attn.o_proj", .{p}), d, c.heads * c.head_dim); + for (0..c.experts + 1) |j| { + const e = if (j < c.experts) try std.fmt.allocPrint(a, "{s}mlp.experts.{d}", .{ p, j }) else try std.fmt.allocPrint(a, "{s}mlp.shared_experts", .{p}); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}.gate_proj", .{e}), c.moe_width, d); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}.up_proj", .{e}), c.moe_width, d); + try Add.fp8(&entries, a, try std.fmt.allocPrint(a, "{s}.down_proj", .{e}), d, c.moe_width); + } + } + var header: std.Io.Writer.Allocating = .init(a); + try header.writer.writeAll("{"); + var at: usize = 0; + var first = true; + for (entries.items) |e| { + var n: usize = if (e.dtype[0] == 'F' and e.dtype[1] == '8') 1 else if (e.dtype[0] == 'B') 2 else 4; + for (e.shape) |x| n *= x; + try header.writer.print("{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ if (first) "" else ",", e.name, e.dtype }); + for (e.shape, 0..) |x, i| try header.writer.print("{s}{d}", .{ if (i == 0) "" else ",", x }); + try header.writer.print("],\"data_offsets\":[{d},{d}]}}", .{ at, at + n }); + first = false; + at += n; + } + try header.writer.writeAll("}"); + const h = header.written(); + const file = try a.alloc(u8, 8 + h.len + at); + @memset(file, 0); + std.mem.writeInt(u64, file[0..8], h.len, .little); + @memcpy(file[8..][0..h.len], h); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file }); +} + +fn tmpDir(a: std.mem.Allocator, tmp: testing.TmpDir) ![]u8 { + const cwd = try std.process.currentPathAlloc(io, a); + return std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); +} + +test "every tensor of a checkpoint maps, the shared expert last" { + var arena: std.heap.ArenaAllocator = .init(testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeTiny(a, tmp); + var ck = try Checkpoint.openModel(testing.allocator, io, try tmpDir(a, tmp)); + defer ck.close(); + var m = try map(testing.allocator, &ck, tiny); + defer m.deinit(); + try testing.expectEqual(@as(usize, 3), m.layers[1].experts.len); + try testing.expectEqual(@as(usize, 2), m.layers[1].experts[2].down.s.dim(0)); // the shared expert, last: 256 / 128 + try testing.expectEqual(core.safetensors.DType.f8_e4m3, m.layers[0].q.w.dtype); +} + +test "the pinned checkpoint (TF_KOLIBRI_DIR): 116,303 tensors, none left over" { + const dir = testing.environ.getPosix("TF_KOLIBRI_DIR") orelse return error.SkipZigTest; + const cfg = try Config.read(testing.allocator, io, dir); + var ck = try Checkpoint.openModel(testing.allocator, io, dir); + defer ck.close(); + var m = try map(testing.allocator, &ck, cfg); + defer m.deinit(); + try testing.expectEqual(@as(usize, 385), m.layers[49].experts.len); +} diff --git a/zig/src/families/kolibri1/pack.zig b/zig/src/families/kolibri1/pack.zig new file mode 100644 index 0000000000..8c93570cf6 --- /dev/null +++ b/zig/src/families/kolibri1/pack.zig @@ -0,0 +1,165 @@ +//! One layer in the device layouts, packed on the host: lane FP8 tiles, stacked experts (shared last), fp32 norms. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const Config = @import("config.zig").Config; +const names = @import("names.zig"); + +const Tensor = core.checkpoint.Tensor; +const fp8 = cuda.fp8; +const fx8 = cuda.experts.Fp8Experts; + +/// A block-FP8 projection in the lane's layout (cuda/fp8.zig): codes [npad/64][k/64][8][32][2][8], fp32 scale tiles. +pub const Linear = struct { w8: []u8, bs: []f32, n: usize, k: usize }; + +pub const Layer = struct { + input_norm: []f32, + qkv: Linear, // q, k and v rows stacked, as attention reads them in one call + q_norm: []f32, + k_norm: []f32, + o: Linear, + post_attn_norm: []f32, + pre_moe_norm: []f32, + router: []u8, // [E, D] bf16 as stored + bias: []f32, + up: []u32, // experts.Layer.up (fp8g): [E][NI/32][D/32][2][32][2][4] + down: []u32, // [E][D/32][NI/32][1][32][2][4] + up_scale: []f32, // [E][2][NI/128][D/128] (gate, up) + down_scale: []f32, // [E][1][D/128][NI/128] + + pub fn deinit(l: *Layer, gpa: std.mem.Allocator) void { + inline for (.{ "input_norm", "q_norm", "k_norm", "post_attn_norm", "pre_moe_norm", "bias", "up_scale", "down_scale" }) |f| + gpa.free(@field(l, f)); + for ([_]Linear{ l.qkv, l.o }) |x| { + gpa.free(x.w8); + gpa.free(x.bs); + } + gpa.free(l.router); + gpa.free(l.up); + gpa.free(l.down); + l.* = undefined; + } +}; + +/// `.float()` of a bf16 tensor: exact. +pub fn widen(gpa: std.mem.Allocator, t: Tensor) ![]f32 { + if (t.dtype != .bf16) return error.UnexpectedTensor; + const out = try gpa.alloc(f32, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16); + return out; +} + +/// fp32 values of a tensor whose bytes may sit at any offset of the mapped file. +fn floats(dst: []f32, t: Tensor) void { + for (dst, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); +} + +/// Row-stacked block-FP8 parts as one projection: codes and scales concatenated by rows, then packed. +fn linear(gpa: std.mem.Allocator, parts: []const names.Fp8) !Linear { + const k = parts[0].w.dim(1); + var n: usize = 0; + var blocks: usize = 0; + for (parts) |p| { + if (p.w.dim(1) != k or p.w.dim(0) % names.block != 0) return error.UnexpectedTensor; + n += p.w.dim(0); + blocks += p.s.numel(); + } + const codes = try gpa.alloc(u8, n * k); + defer gpa.free(codes); + const inv = try gpa.alloc(f32, blocks); + defer gpa.free(inv); + var at: usize = 0; + var sat: usize = 0; + for (parts) |p| { + @memcpy(codes[at..][0..p.w.bytes.len], p.w.bytes); + at += p.w.bytes.len; + floats(inv[sat..][0..p.s.numel()], p.s); + sat += p.s.numel(); + } + const npad = fp8.padded(n); + const w8 = try gpa.alloc(u8, npad * k); + errdefer gpa.free(w8); + fp8.packCodes(w8, codes, n, k); + const bs = try gpa.alloc(f32, npad * (k / fp8.group)); + fp8.packScales(bs, inv, n, k); + return .{ .w8 = w8, .bs = bs, .n = n, .k = k }; +} + +/// Experts as cuda/experts.zig's fp8g layer takes them: gate and up interleaved, down alone, scales gate then up. +fn experts(gpa: std.mem.Allocator, l: *Layer, ex: []const names.Expert, c: Config) !void { + const ni = c.moe_width; + const d = c.hidden; + const one = ni * d; + const sb = (ni / names.block) * (d / names.block); + l.up = try gpa.alloc(u32, ex.len * one / 2); + l.down = try gpa.alloc(u32, ex.len * one / 4); + l.up_scale = try gpa.alloc(f32, ex.len * 2 * sb); + l.down_scale = try gpa.alloc(f32, ex.len * sb); + for (ex, 0..) |e, i| { + try fx8.packGateUp(gpa, l.up[i * one / 2 ..][0 .. one / 2], e.gate.w.bytes, e.up.w.bytes, ni, d); + fx8.packOne(l.down[i * one / 4 ..][0 .. one / 4], e.down.w.bytes, d, ni); + floats(l.up_scale[2 * i * sb ..][0..sb], e.gate.s); + floats(l.up_scale[(2 * i + 1) * sb ..][0..sb], e.up.s); + floats(l.down_scale[i * sb ..][0..sb], e.down.s); + } +} + +/// The packed layer for one mapped layer. +pub fn layer(gpa: std.mem.Allocator, m: names.Layer, c: Config) !Layer { + var l: Layer = undefined; + l.qkv = try linear(gpa, &.{ m.q, m.k, m.v }); + l.o = try linear(gpa, &.{m.o}); + l.input_norm = try widen(gpa, m.input_norm); + l.q_norm = try widen(gpa, m.q_norm); + l.k_norm = try widen(gpa, m.k_norm); + l.post_attn_norm = try widen(gpa, m.post_attn_norm); + l.pre_moe_norm = try widen(gpa, m.pre_moe_norm); + l.bias = try widen(gpa, m.bias); + l.router = try gpa.dupe(u8, m.router.bytes); + try experts(gpa, &l, m.experts, c); + return l; +} + +const testing = std.testing; + +fn tensor(dtype: core.safetensors.DType, shape: []const usize, bytes: []const u8) Tensor { + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = undefined, .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +test "q, k and v stack by rows into one projection, codes and scales alike" { + const a = testing.allocator; + const k = 128; + var codes: [3][128 * k]u8 = undefined; + var inv: [3][1]f32 = .{ .{2}, .{3}, .{5} }; + for (&codes, 0..) |*c, p| for (c, 0..) |*b, i| { + b.* = @truncate(i * 13 + p * 7 + 1); + }; + var parts: [3]names.Fp8 = undefined; + for (&parts, 0..) |*p, i| p.* = .{ .w = tensor(.f8_e4m3, &.{ 128, k }, &codes[i]), .s = tensor(.f32, &.{ 1, 1 }, std.mem.asBytes(&inv[i])) }; + const l = try linear(a, &parts); + defer a.free(l.w8); + defer a.free(l.bs); + try testing.expectEqual(@as(usize, 384), l.n); + const all = try a.alloc(u8, 3 * 128 * k); + defer a.free(all); + for (codes, 0..) |c, i| @memcpy(all[i * 128 * k ..][0 .. 128 * k], &c); + const want = try a.alloc(u8, 384 * k); + defer a.free(want); + fp8.packCodes(want, all, 384, k); + try testing.expectEqualSlices(u8, want, l.w8); + var ws: [384 * 2]f32 = undefined; + fp8.packScales(&ws, &.{ 2, 3, 5 }, 384, k); + try testing.expectEqualSlices(f32, &ws, l.bs); +} + +test "bf16 widens to fp32 exactly, and other dtypes are refused" { + const bits = [_]u16{ 0x3f80, 0xc000, 0x0001 }; + const v = try widen(testing.allocator, tensor(.bf16, &.{3}, std.mem.sliceAsBytes(&bits))); + defer testing.allocator.free(v); + try testing.expectEqual(@as(f32, 1), v[0]); + try testing.expectEqual(@as(f32, -2), v[1]); + try testing.expectEqual(@as(u32, 0x00010000), @as(u32, @bitCast(v[2]))); + try testing.expectError(error.UnexpectedTensor, widen(testing.allocator, tensor(.f32, &.{1}, std.mem.sliceAsBytes(&bits)[0..4]))); +} diff --git a/zig/src/families/kolibri1/weights.zig b/zig/src/families/kolibri1/weights.zig new file mode 100644 index 0000000000..afea9d253e --- /dev/null +++ b/zig/src/families/kolibri1/weights.zig @@ -0,0 +1,100 @@ +//! Kolibri-1's weights on the GPU: each layer packed on the host, then uploaded; projections as qlinear fp8g weights. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const Config = @import("config.zig").Config; +const names = @import("names.zig"); +const pack = @import("pack.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const fp8 = cuda.fp8; + +pub const Layer = struct { + input_norm: u64, // fp32 [D] + qkv: cuda.qlinear.Weight, // fp8g + q_norm: u64, // fp32 [head_dim] + k_norm: u64, + o: cuda.qlinear.Weight, + post_attn_norm: u64, + pre_moe_norm: u64, + router: u64, // bf16 [E, D] + bias: u64, // fp32 [E] + experts: cuda.experts.Layer, // E routed, then the shared expert (format fp8g) + post_moe_norm: u64, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: Config, + embed: u64 = 0, // bf16 [V, D] + norm: u64 = 0, // fp32 [D] + head: u64 = 0, // bf16 [V, D] + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + try w.buffers.append(w.gpa, b); + return b.ptr; + } + + fn widened(w: *Weights, d: *const cuda.Driver, t: core.checkpoint.Tensor) !u64 { + const v = try pack.widen(w.gpa, t); + defer w.gpa.free(v); + return w.upload(d, std.mem.sliceAsBytes(v)); + } + + fn linear(w: *Weights, d: *const cuda.Driver, l: pack.Linear) !cuda.qlinear.Weight { + return .{ .fp8g = fp8.weight(try w.upload(d, l.w8), try w.upload(d, std.mem.sliceAsBytes(l.bs)), l.n, l.k) }; + } +}; + +/// config.json, every tensor mapped and checked, then layer by layer; `layers` caps them (tests). +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, layers: ?usize) !Weights { + const c = try Config.read(gpa, io, dir); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var m = try names.map(gpa, &ck, c); + defer m.deinit(); + var w: Weights = .{ .gpa = gpa, .config = c }; + errdefer w.deinit(); + const count = @min(layers orelse c.layers, c.layers); + w.layers = try gpa.alloc(Layer, count); + w.embed = try w.upload(d, m.embed.bytes); + for (w.layers, m.layers[0..count]) |*out, ml| { + var h = try pack.layer(gpa, ml, c); + defer h.deinit(gpa); + out.* = .{ + .input_norm = try w.upload(d, std.mem.sliceAsBytes(h.input_norm)), + .qkv = try w.linear(d, h.qkv), + .q_norm = try w.upload(d, std.mem.sliceAsBytes(h.q_norm)), + .k_norm = try w.upload(d, std.mem.sliceAsBytes(h.k_norm)), + .o = try w.linear(d, h.o), + .post_attn_norm = try w.upload(d, std.mem.sliceAsBytes(h.post_attn_norm)), + .pre_moe_norm = try w.upload(d, std.mem.sliceAsBytes(h.pre_moe_norm)), + .router = try w.upload(d, h.router), + .bias = try w.upload(d, std.mem.sliceAsBytes(h.bias)), + .experts = .{ + .format = .fp8g, + .up = try w.upload(d, std.mem.sliceAsBytes(h.up)), + .down = try w.upload(d, std.mem.sliceAsBytes(h.down)), + .up_scale = try w.upload(d, std.mem.sliceAsBytes(h.up_scale)), + .down_scale = try w.upload(d, std.mem.sliceAsBytes(h.down_scale)), + .width = c.moe_width, + .dims = c.hidden, + .experts = c.experts + 1, + }, + .post_moe_norm = try w.widened(d, ml.post_moe_norm), + }; + } + w.norm = try w.widened(d, m.norm); + w.head = try w.upload(d, m.head.bytes); + return w; +} From de56c8dc928d07ab4ed59843b0957dea20b3dea2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:53:25 +0200 Subject: [PATCH 25/41] kolibri1: the layer graph on shared CUDA operators (window, tree, qk rows, norms, sigmoid routing, grouped FP8 experts), one-row step graphs --- zig/src/families/kolibri1/cuda.zig | 6 + zig/src/families/kolibri1/forward.zig | 382 ++++++++++++++++++++++++++ zig/src/families/kolibri1/kernels.zig | 176 ++++++++++++ zig/src/families/kolibri1/tree.zig | 66 +++++ 4 files changed, 630 insertions(+) create mode 100644 zig/src/families/kolibri1/forward.zig create mode 100644 zig/src/families/kolibri1/kernels.zig create mode 100644 zig/src/families/kolibri1/tree.zig diff --git a/zig/src/families/kolibri1/cuda.zig b/zig/src/families/kolibri1/cuda.zig index 0ceab59ddc..85e9250575 100644 --- a/zig/src/families/kolibri1/cuda.zig +++ b/zig/src/families/kolibri1/cuda.zig @@ -4,10 +4,16 @@ pub const config = @import("config.zig"); pub const names = @import("names.zig"); pub const pack = @import("pack.zig"); pub const weights = @import("weights.zig"); +pub const tree = @import("tree.zig"); +pub const kernels = @import("kernels.zig"); +pub const forward = @import("forward.zig"); test { _ = @import("config.zig"); _ = @import("names.zig"); _ = @import("pack.zig"); _ = @import("weights.zig"); + _ = @import("tree.zig"); + _ = @import("kernels.zig"); + _ = @import("forward.zig"); } diff --git a/zig/src/families/kolibri1/forward.zig b/zig/src/families/kolibri1/forward.zig new file mode 100644 index 0000000000..ad842336ed --- /dev/null +++ b/zig/src/families/kolibri1/forward.zig @@ -0,0 +1,382 @@ +//! Kolibri-1's layer graph over streams' serial chains of rows: prompt chunks, decode rows, one-row step graphs. +const std = @import("std"); +const cuda = @import("cuda"); +const Config = @import("config.zig").Config; +const wts = @import("weights.zig"); +const kern = @import("kernels.zig"); +const tree = @import("tree.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const grouped = cuda.grouped; +const experts = cuda.experts; +const TreeAttention = cuda.tree_attention.TreeAttention; + +/// Prompt rows a forward takes: the chunk grid, from position 0. +pub const prompt_rows = 2048; +/// Decode rows a forward takes over every stream's window. +pub const round_rows = 64; +const key_block = 64; // the window kernels' key tile + +/// Keys a sliding layer keeps: a chunk's rows written before its first row reads back its window. +pub fn ringSize(rows: usize, window: usize) usize { + return (rows + window + key_block - 1) / key_block * key_block; +} + +/// A stream's rows in one forward: its cache slot, first position, tokens. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + c: Config, + context: usize, + slots: usize, + ring: usize, + max_rows: usize, + groups: usize, // tree attention's key groups at the context and a full chain + k: []u64, // per layer: the cache [slots, len, HK, D] bf16 (len: context for full layers, the ring for sliding) + v: []u64, + norms: []u64, // per layer: q_norm then k_norm, fp32 [2, D] + cos: u64, + sin: u64, + origin: u64, // the bf16 base tree attention's cache offsets count from + s: Scratch, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + const Scratch = struct { ids: u64, pos: u64, slot: u64, res: u64, x: u64, qkv: u64, q: u64, kn: u64, vn: u64, att: u64, o: u64, part: u64, logits_e: u64, pick: u64, wts: u64, plan: grouped.Plan, act: u64, y: u64, mo: u64, flat: u64, paths: u64, depths: u64, offs: u64, sel: u64, win: u64, tree: u64 }; + + fn alloc(m: *Model, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + fn zeros(m: *Model, s: cuda.Stream, bytes: usize) !u64 { + const p = try m.alloc(bytes); + try m.d.check(m.d.api.cuMemsetD8Async(p, 0, @max(bytes, 256), s.handle), "cuMemsetD8Async"); + return p; + } + + /// Bytes one stream's caches take: full layers hold `context` positions, sliding layers the ring. + pub fn streamBytes(c: Config, context: usize) usize { + var total: usize = 0; + for (c.full[0..c.layers]) |full| total += 2 * (if (full) context else ringSize(prompt_rows, c.window)) * c.kv_heads * c.head_dim * 2; + return total; + } + + /// Caches for `slots` streams of `context` positions, RoPE tables, and scratch for `max_rows` rows a forward. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, max_rows: usize) !Model { + var c = w.config; + c.layers = w.layers.len; // a model loaded with fewer layers (checks) runs those alone + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .c = c, .context = context, .slots = slots, .ring = ringSize(prompt_rows, c.window), .max_rows = max_rows, .groups = TreeAttention.groups(context + tree.max_rows), .k = &.{}, .v = &.{}, .norms = &.{}, .cos = 0, .sin = 0, .origin = 0, .s = undefined }; + errdefer m.deinit(); + const dh = c.head_dim; + m.k = try gpa.alloc(u64, c.layers); + m.v = try gpa.alloc(u64, c.layers); + m.norms = try gpa.alloc(u64, c.layers); + for (0..c.layers) |i| { + const len = if (c.full[i]) context else m.ring; + const bytes = slots * len * c.kv_heads * dh * 2; + m.k[i] = try m.zeros(s, bytes); + m.v[i] = try m.zeros(s, bytes); + m.norms[i] = try m.alloc(2 * dh * 4); + const L = w.layers[i]; + try d.check(d.api.cuMemcpyDtoDAsync_v2(m.norms[i], L.q_norm, dh * 4, s.handle), "cuMemcpyDtoDAsync"); + try d.check(d.api.cuMemcpyDtoDAsync_v2(m.norms[i] + dh * 4, L.k_norm, dh * 4, s.handle), "cuMemcpyDtoDAsync"); + } + try m.rope(s); + m.origin = try m.zeros(s, 256); + try m.scratch(k); + try s.synchronize(); + return m; + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.k); + m.gpa.free(m.v); + m.gpa.free(m.norms); + } + + /// RoPE tables: angle p / theta^(2j / D) in float64, cos and sin rounded to fp32, the halves repeated (neox). + fn rope(m: *Model, s: cuda.Stream) !void { + const dh = m.c.head_dim; + const cos = try m.gpa.alloc(f32, m.context * dh); + defer m.gpa.free(cos); + const sin = try m.gpa.alloc(f32, m.context * dh); + defer m.gpa.free(sin); + const theta: f64 = m.c.rope_theta; + for (0..m.context) |p| for (0..dh / 2) |j| { + const frac = @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(dh)); + const ang = @as(f64, @floatFromInt(p)) / std.math.pow(f64, theta, frac); + for ([_]usize{ j, j + dh / 2 }) |col| { + cos[p * dh + col] = @floatCast(@cos(ang)); + sin[p * dh + col] = @floatCast(@sin(ang)); + } + }; + m.cos = try m.alloc(cos.len * 4); + m.sin = try m.alloc(sin.len * 4); + try m.upload(s, m.cos, std.mem.sliceAsBytes(cos)); + try m.upload(s, m.sin, std.mem.sliceAsBytes(sin)); + try s.synchronize(); // the host tables are freed on return + } + + pub fn upload(m: *Model, s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); + } + + /// The lane's slice partials at most, over every row count a forward takes. + fn partBytes(m: *const Model, k: *const kern.Kernels) usize { + var most: usize = 0; + for (m.w.layers) |L| for ([_]cuda.qlinear.Weight{ L.qkv, L.o }) |q| { + var rows: usize = 1; + while (rows <= m.max_rows) : (rows += 1) most = @max(most, k.lane.partBytes(rows, q.fp8g)); + }; + return most; + } + + fn scratch(m: *Model, k: *const kern.Kernels) !void { + const c = m.c; + const r = m.max_rows; + const dh = c.head_dim; + const pairs = r * (c.top_k + 1); + const count = c.experts + 1; + const wide = pairs > grouped.small; + const items = @max(grouped.maxItems(pairs, count, 16), grouped.maxItems(pairs, count, experts.promptTile(.fp8g))); + const cap: usize = @min(r, round_rows); + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .slot = try m.alloc(r * 4), + .res = try m.alloc(r * c.hidden * 4), + .x = try m.alloc(r * c.hidden * 2), + .qkv = try m.alloc(r * c.qkvDim() * 2), + .q = try m.alloc(r * c.heads * dh * 2), + .kn = try m.alloc(r * c.kv_heads * dh * 2), + .vn = try m.alloc(r * c.kv_heads * dh * 2), + .att = try m.alloc(r * c.heads * dh * 2), + .o = try m.alloc(r * c.hidden * 2), + .part = try m.alloc(m.partBytes(k)), + .logits_e = try m.alloc(r * c.experts * 4), + .pick = try m.alloc(pairs * 4), + .wts = try m.alloc(pairs * 4), + .plan = .{ + .members = try m.alloc(pairs * 4), + .items = try m.alloc(items * 12), + .counts = try m.alloc(8), + .rank = try m.alloc((if (wide) pairs else 1) * 4), + .hist = try m.alloc((if (wide) (pairs + 1023) / 1024 * count else 1) * 4), + }, + .act = try m.alloc(pairs * c.moe_width * 2), + .y = try m.alloc(pairs * c.hidden * 4), + .mo = try m.alloc(r * c.hidden * 2), + .flat = try m.alloc((2 * cap + 4 * m.slots) * 4), + .paths = try m.alloc(cap * tree.max_rows * 4), + .depths = try m.alloc(cap * 4), + .offs = try m.alloc(c.layers * m.slots * 2 * 8), + .sel = try m.alloc(r * c.hidden * 2), + .tree = try m.alloc(TreeAttention.scratchBytes(cap, c.heads, c.kv_heads, m.groups, dh)), + .win = try m.alloc(cuda.window_attention.WindowAttention.scratchBytes(m.window(), cap)), + }; + } + + /// A sliding layer's keys: (pos - window, pos], key p at ring row p % ring. + fn window(m: *const Model) cuda.window_attention.Window { + const c = m.c; + return .{ .h = c.heads, .hk = c.kv_heads, .d = c.head_dim, .win = c.window, .ring = m.ring, .scale = scale(c) }; + } + + fn scale(c: Config) f32 { + return @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(c.head_dim)))); + } + + /// fp32 logits [rows.len, V] at `out` of the chains' rows `rows` (default: each chain's last). + pub fn forward(m: *Model, o: kern.Ops, chains: []const Chain, prompt: bool, rows: ?[]const usize, out: u64) !void { + const gpa = m.gpa; + const s = o.s; + if (prompt and chains.len != 1) return error.PromptTakesOneStream; + if (!prompt and chains.len > m.slots) return error.TooManyStreams; + var n: usize = 0; + for (chains) |ch| { + if (ch.tokens.len == 0) return error.EmptyChain; + if (ch.tokens.len > m.ring - m.c.window) return error.ChainOverrunsRing; + if (ch.p0 + ch.tokens.len > m.context) return error.PastContext; + if (ch.slot >= m.slots) return error.NoSuchSlot; + n += ch.tokens.len; + } + if (n > m.max_rows or (!prompt and n > round_rows)) return error.TooManyRows; + const host = try gpa.alloc(i32, 3 * n); + defer gpa.free(host); + var at: usize = 0; + for (chains) |ch| for (ch.tokens, 0..) |t, i| { + host[at] = @intCast(t); + host[n + at] = @intCast(ch.p0 + i); + host[2 * n + at] = @intCast(ch.slot); + at += 1; + }; + try m.upload(s, m.s.ids, std.mem.sliceAsBytes(host[0..n])); + try m.upload(s, m.s.pos, std.mem.sliceAsBytes(host[n .. 2 * n])); + try m.upload(s, m.s.slot, std.mem.sliceAsBytes(host[2 * n ..])); + var t: ?cuda.tree_attention.Rows = null; + var ng: usize = 0; + var plan: ?tree.Plan = null; + defer if (plan) |p| gpa.free(p.flat); + if (!prompt) { + const windows = try gpa.alloc(tree.Window, chains.len); + defer gpa.free(windows); + var most: usize = 0; + for (chains, windows) |ch, *win| { + win.* = .{ .rows = ch.tokens.len, .committed = ch.p0 }; + most = @max(most, ch.p0 + ch.tokens.len); + } + plan = try tree.build(gpa, windows); + try m.upload(s, m.s.flat, std.mem.sliceAsBytes(plan.?.flat)); + t = m.treeRows(m.s.flat, plan.?, m.s.paths, m.s.depths, m.s.offs); + try o.k.tree.paths(s, m.s.flat + plan.?.parentsAt(), m.s.paths, m.s.depths, n); + try m.offsets(s, chains, m.s.offs); + ng = TreeAttention.groups(most); + } + var ends: std.ArrayList(usize) = .empty; + defer ends.deinit(gpa); + const pick = rows orelse blk: { + var e: usize = 0; + for (chains) |ch| { + e += ch.tokens.len; + try ends.append(gpa, e - 1); + } + break :blk ends.items; + }; + for (pick) |r| if (r >= n) return error.NoSuchRow; + try m.launch(o, n, if (prompt) chains[0] else null, t, chains.len, ng, @min(m.max_rows, round_rows), pick, out); + } + + fn treeRows(m: *const Model, base: u64, p: tree.Plan, paths: u64, depths: u64, offs: u64) cuda.tree_attention.Rows { + return .{ .row_stream = base, .streams = base + p.streamsAt(), .offs = offs, .origin = m.origin, .paths = paths, .depths = depths }; + } + + /// Each full layer's cache offsets a chain, K then V, in bf16 elements from `origin`: [layer][chain][2]. + fn offsets(m: *Model, s: cuda.Stream, chains: []const Chain, dst: u64) !void { + const c = m.c; + const per = m.context * c.kv_heads * c.head_dim * 2; + const offs = try m.gpa.alloc(i64, c.layers * 2 * chains.len); + defer m.gpa.free(offs); + for (0..c.layers) |i| for (chains, 0..) |ch, j| { + offs[(i * chains.len + j) * 2] = @intCast((m.k[i] + ch.slot * per - m.origin) / 2); + offs[(i * chains.len + j) * 2 + 1] = @intCast((m.v[i] + ch.slot * per - m.origin) / 2); + }; + try m.upload(s, dst, std.mem.sliceAsBytes(offs)); + } + + /// The layers over staged inputs, no host copies (a graph can hold them); `prompt_chain` set: a prompt chunk. + fn launch(m: *Model, o: kern.Ops, n: usize, prompt_chain: ?Chain, t: ?cuda.tree_attention.Rows, streams: usize, ng: usize, cap: usize, pick: []const usize, out: u64) !void { + const c = m.c; + const s = o.s; + const prompt = prompt_chain != null; + const w = m.w; + const h = c.heads; + const hk = c.kv_heads; + const dh = c.head_dim; + const lin: cuda.qlinear.Linear = .{ .lane = &o.k.lane }; + try o.embed(w.embed, m.s.ids, n, c.hidden, m.s.x, m.s.res); + try o.rms(m.s.res, w.layers[0].input_norm, m.s.x, n, c.hidden, c.eps); + const slots_e = c.top_k + 1; + const pairs = n * slots_e; + const tile = if (prompt) experts.promptTile(.fp8g) else 16; + const items = grouped.maxItems(pairs, c.experts + 1, tile); + const gx: experts.Grouped = .{ .fp8g = &o.k.fp8x }; + for (w.layers, 0..) |L, i| { + try (if (prompt) lin.prompt(s, L.qkv, .{ .x = m.s.x, .part = m.s.part }, m.s.qkv, n) else lin.decode(s, L.qkv, .{ .x = m.s.x, .part = m.s.part }, m.s.qkv, n)); + const len = if (c.full[i]) m.context else m.ring; + try o.qkv(.{ .rows = m.s.qkv, .norms = m.norms[i], .cos = m.cos, .sin = m.sin, .pos = m.s.pos, .slots = m.s.slot, .q = m.s.q, .k = m.s.kn, .v = m.s.vn, .kc = m.k[i], .vc = m.v[i], .len = len }, n, h, hk, dh, c.eps, !c.full[i]); + if (prompt_chain) |ch| { + const per = len * hk * dh * 2; // a stream's rows of this layer's cache + var win = m.window(); + if (c.full[i]) { + win.win = 0; // every key before the row + win.ring = m.context; + } + try o.promptRows(m.s.q, m.k[i] + ch.slot * per, m.v[i] + ch.slot * per, m.s.att, ch.p0, n, win); + } else if (!c.full[i]) { + try o.sliding(m.s.q, m.k[i], m.v[i], m.s.att, m.s.pos, m.s.slot, n, m.window(), m.s.win); + } else { + var rows = t.?; + rows.offs += i * streams * 16; + try o.full(m.s.q, m.s.kn, m.s.vn, rows, m.s.tree, cap, ng, m.s.att, n, h, hk, dh, scale(c)); + } + try (if (prompt) lin.prompt(s, L.o, .{ .x = m.s.att, .part = m.s.part }, m.s.o, n) else lin.decode(s, L.o, .{ .x = m.s.att, .part = m.s.part }, m.s.o, n)); + try o.addRms(m.s.o, m.s.res, L.post_attn_norm, L.pre_moe_norm, m.s.x, n, c.hidden, c.eps); + // routed experts by sigmoid + bias, the shared expert as one more slot of weight 1, then the weighted sum + try o.dense(m.s.x, c.hidden, L.router, m.s.logits_e, n, c.hidden, c.experts); + try o.topk(m.s.logits_e, L.bias, m.s.pick, m.s.wts, n, c.experts, c.top_k); + try o.k.router.route(s, m.s.pick, pairs, c.experts + 1, tile, m.s.plan); + const up_in: experts.Rows = .{ .x = m.s.x, .stride = c.hidden, .slots = slots_e }; + const down_in: experts.Rows = .{ .x = m.s.act, .stride = c.moe_width }; + if (prompt) { + try gx.prompt(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1); + try gx.prompt(s, .down_bf16, down_in, L.experts, m.s.plan, m.s.y, items, -1); + } else { + try gx.decode(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1); + try gx.decode(s, .down_f32, down_in, L.experts, m.s.plan, m.s.y, items, -1); + } + try o.combine(m.s.y, prompt, m.s.wts, m.s.mo, n, slots_e, c.hidden); + const after = if (i + 1 < w.layers.len) w.layers[i + 1].input_norm else w.norm; + try o.addRms(m.s.mo, m.s.res, L.post_moe_norm, after, m.s.x, n, c.hidden, c.eps); + } + for (pick, 0..) |r, j| try m.d.check(m.d.api.cuMemcpyDtoDAsync_v2(m.s.sel + j * c.hidden * 2, m.s.x + r * c.hidden * 2, c.hidden * 2, s.handle), "cuMemcpyDtoDAsync"); + try o.dense(m.s.sel, c.hidden, w.head, out, pick.len, c.hidden, c.vocab); + } + + /// One stream's one-row step as a CUDA graph: its stream row, paths and offsets fixed; a step writes token and row. + pub const Step = struct { exec: cuda.graph.Exec, slot: usize, flat: u64, logits: u64 }; + + pub fn captureStep(m: *Model, o: kern.Ops, slot: usize) !Step { + const c = m.c; + const s = o.s; + const p = try tree.build(m.gpa, &.{.{ .rows = 1, .committed = 0 }}); + defer m.gpa.free(p.flat); + const flat = try m.alloc(p.flat.len * 4); + try m.upload(s, flat, std.mem.sliceAsBytes(p.flat)); + const paths = try m.alloc(tree.max_rows * 4); + const depths = try m.alloc(4); + try o.k.tree.paths(s, flat + p.parentsAt(), paths, depths, 1); // a lone root: its path never changes + const offs = try m.alloc(c.layers * 16); + try m.offsets(s, &.{.{ .slot = slot, .p0 = 0, .tokens = &.{} }}, offs); + const logits = try m.alloc(c.vocab * 4); + try s.synchronize(); + try cuda.graph.beginCapture(s, .thread_local); + m.launch(o, 1, null, m.treeRows(flat, p, paths, depths, offs), 1, m.groups, 1, &.{0}, logits) catch |e| { + if (cuda.graph.endCapture(s)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return e; + }; + var g = try cuda.graph.endCapture(s); + defer g.deinit(); + return .{ .exec = try g.instantiate(), .slot = slot, .flat = flat, .logits = logits }; + } + + /// The step: `token` at position `p0` (the cache length before it) through the captured graph. + pub fn step(m: *Model, o: kern.Ops, g: *const Step, token: u32, p0: usize) !void { + const s = o.s; + if (p0 + 1 > m.context) return error.PastContext; + const staged = [3]i32{ @intCast(token), @intCast(p0), @intCast(g.slot) }; + try m.upload(s, m.s.ids, std.mem.asBytes(&staged[0])); + try m.upload(s, m.s.pos, std.mem.asBytes(&staged[1])); + try m.upload(s, m.s.slot, std.mem.asBytes(&staged[2])); + const row = tree.stepRow(p0); + try m.upload(s, g.flat + 4, std.mem.asBytes(&row)); + try g.exec.launchOn(s); + } +}; + +test "the ring holds a prompt chunk and a window, in whole key tiles" { + try std.testing.expectEqual(@as(usize, 2624), ringSize(prompt_rows, 513)); + try std.testing.expect(ringSize(prompt_rows, 513) - 513 >= prompt_rows); +} diff --git a/zig/src/families/kolibri1/kernels.zig b/zig/src/families/kolibri1/kernels.zig new file mode 100644 index 0000000000..de488c55f7 --- /dev/null +++ b/zig/src/families/kolibri1/kernels.zig @@ -0,0 +1,176 @@ +//! Kolibri-1's launches, every one a shared CUDA operator: lane FP8, grouped experts, norms, routing, attention. +const std = @import("std"); +const cuda = @import("cuda"); + +const grouped = cuda.grouped; +const experts = cuda.experts; + +/// The torch-op images' functions this family calls (argmax by torch's first-maximum rule, bf16 casts). +const torch_argmax = "tf_argmax_rows_i32_kernel"; +const torch_to_f32 = "tf_bf16_to_f32_kernel"; +const torch_to_bf16 = "tf_f32_to_bf16_kernel"; + +pub const Kernels = struct { + lane: cuda.qmmf.Lane, + fp8x: experts.Fp8Experts, + router: grouped.Router, + mods: [4]cuda.Module, // experts (the plan kernels), sample, torch_argmax, torch_pointwise + argmax_fn: cuda.Function, + to_f32_fn: cuda.Function, + to_bf16_fn: cuda.Function, + draws: cuda.sample.Draws, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + moe: cuda.moe_glue.MoeGlue, + dense: cuda.dense_rows.DenseRows, + qk: cuda.qk_rows.QkRows, + win: cuda.window_attention.WindowAttention, + tree: cuda.tree_attention.TreeAttention, + + /// Every module is a shared image built into the runtime: nothing captured, nothing of the family's own. + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + const kk = cuda.kernels; + const images = [_][]const u8{ kk.experts, kk.sample, kk.torch_argmax, kk.torch_pointwise }; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + for (images, 0..) |img, i| { + k.mods[i] = try cuda.Module.load(d, img); + loaded += 1; + } + k.router = try grouped.Router.resolve(k.mods[0]); + k.draws = try cuda.sample.Draws.resolve(k.mods[1]); + k.argmax_fn = try k.mods[2].function(torch_argmax); + k.to_f32_fn = try k.mods[3].function(torch_to_f32); + k.to_bf16_fn = try k.mods[3].function(torch_to_bf16); + k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10); + errdefer k.lane.unload(); + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + k.fp8x = try experts.Fp8Experts.load(d, sms); + errdefer k.fp8x.unload(); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.moe = try cuda.moe_glue.MoeGlue.load(d); + errdefer k.moe.unload(); + k.dense = try cuda.dense_rows.DenseRows.load(d); + errdefer k.dense.unload(); + k.qk = try cuda.qk_rows.QkRows.load(d); + errdefer k.qk.unload(); + k.win = try cuda.window_attention.WindowAttention.load(d); + errdefer k.win.unload(); + k.tree = try cuda.tree_attention.TreeAttention.load(d); + k.tree.sms = sms; + return k; + } + + pub fn deinit(k: *Kernels) void { + k.embed.unload(); + k.norms.unload(); + k.moe.unload(); + k.dense.unload(); + k.qk.unload(); + k.win.unload(); + k.tree.unload(); + k.fp8x.unload(); + k.lane.unload(); + for (&k.mods) |*m| m.unload(); + } +}; + +/// launch_shape.h's block count: one block a 256 elements, at most 65535. +fn blocks(count: usize) u32 { + return @intCast(@max(1, @min((count + 255) / 256, 65535))); +} + +/// Launches on one stream. +pub const Ops = struct { + k: *const Kernels, + s: cuda.Stream, + + fn one(o: Ops, f: cuda.Function, grid: u32, block: u32, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = grid }, .block = .{ .x = block } }, o.s, a); + } + + /// The fp32 residual of `rows` tokens: their bf16 table rows into `tmp`, widened exactly into `out`. + pub fn embed(o: Ops, table: u64, ids: u64, rows: usize, d: usize, tmp: u64, out: u64) !void { + try o.k.embed.bf16(o.s, ids, table, tmp, rows, d); + var a: cuda.Args = .{}; + a.add(tmp); + a.add(out); + a.add(@as(u64, rows * d)); + try o.one(o.k.to_f32_fn, blocks(rows * d), 256, &a); + } + + /// Each row's argmax over `v` fp32 logits, the first maximum on ties (torch.argmax), into `out` (int32 a row). + pub fn argmax(o: Ops, logits: u64, rows: usize, v: usize, out: u64) !void { + var a: cuda.Args = .{}; + a.add(logits); + a.add(out); + for ([_]usize{ rows, v, v }) |x| a.add(@as(i64, @intCast(x))); + a.add(@as(i32, 1)); // dtype 1: fp32 + try o.one(o.k.argmax_fn, @intCast(rows), 1024, &a); + } + + /// bf16 copies of `n` fp32 logits (round to nearest even), as sample.cu's draws read them. + pub fn toBf16(o: Ops, x: u64, y: u64, n: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(y); + a.add(@as(u64, n)); + try o.one(o.k.to_bf16_fn, blocks(n), 256, &a); + } + + /// RMS norm of fp32 rows, bf16 out. + pub fn rms(o: Ops, x: u64, w: u64, out: u64, rows: usize, n: usize, eps: f32) !void { + try o.k.norms.rms(o.s, x, .f32, w, out, .bf16, rows, n, eps); + } + + /// res += rms(x) * w1 in place (fp32); out = bf16(rms(res) * w2). + pub fn addRms(o: Ops, x: u64, res: u64, w1: u64, w2: u64, out: u64, rows: usize, n: usize, eps: f32) !void { + try o.k.norms.addRms(o.s, x, .bf16, res, w1, w2, out, .bf16, rows, n, eps); + } + + pub const Qkv = struct { rows: u64, norms: u64, cos: u64, sin: u64, pos: u64, slots: u64, q: u64, k: u64, v: u64, kc: u64, vc: u64, len: usize }; + + /// q and k RMS-normed (RoPE'd on sliding layers), v as is; k and v also into the cache at pos % len. + pub fn qkv(o: Ops, a: Qkv, n: usize, h: usize, hk: usize, d: usize, eps: f32, rope: bool) !void { + const shape: cuda.qk_rows.Shape = .{ .norms = a.norms, .cos = if (rope) a.cos else null, .sin = if (rope) a.sin else null, .rot = if (rope) d else 0, .eps = eps }; + try o.k.qk.rows(o.s, a.rows, .{ .h = h, .hk = hk, .dim = d }, shape, a.pos, a.slots, .{ .q = a.q, .k = a.k, .v = a.v, .kc = a.kc, .vc = a.vc, .len = a.len }, n); + } + + /// Decode rows of a sliding layer: each row over its stream's ring, keys (pos - window, pos]. + pub fn sliding(o: Ops, q: u64, kc: u64, vc: u64, out: u64, pos: u64, slots: u64, n: usize, w: cuda.window_attention.Window, scratch: u64) !void { + try o.k.win.decode(o.s, q, kc, vc, out, pos, slots, w, scratch, n); + } + + /// Prompt rows of one stream at p0, p0 + 1, ...: a ring and window (sliding), or every key before (win 0, full). + pub fn promptRows(o: Ops, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: cuda.window_attention.Window) !void { + try o.k.win.prompt(o.s, q, kc, vc, out, p0, n, w); + } + + /// Decode rows of a full layer: committed keys from each stream's cache, then the row's own chain. + pub fn full(o: Ops, q: u64, kn: u64, vn: u64, rows: cuda.tree_attention.Rows, scratch: u64, cap: usize, ng: usize, out: u64, n: usize, h: usize, hk: usize, d: usize, scale: f32) !void { + if (n > cap) return error.TreeRowsPastScratch; + const sc = cuda.tree_attention.TreeAttention.scratch(scratch, cap, h, hk, ng, d); + try o.k.tree.run(o.s, q, kn, vn, rows, sc, out, n, h, hk, d, ng, scale); + } + + /// x [m, D] bf16 times rows [n, D] bf16 into fp32 [m, n] (the router and the head). + pub fn dense(o: Ops, x: u64, x_stride: usize, rows: u64, out: u64, m: usize, d: usize, n: usize) !void { + try o.k.dense.rows(o.s, x, x_stride, rows, out, n, m, n, d); + } + + /// Each row's top-k experts by sigmoid logit + bias, unnormalised sigmoid weights, then the shared expert's slot. + pub fn topk(o: Ops, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, ne: usize, k: usize) !void { + try o.k.moe.topk(o.s, logits, bias, .{ .experts = ne, .topk = k, .shared = true }, pick, wts, rows); + } + + /// y [R, S, D] (fp32 decode, bf16 prompt) weighted by wts [R, S] in slot order -> [R, D] bf16. + pub fn combine(o: Ops, y: u64, y_bf16: bool, wts: u64, out: u64, rows: usize, s: usize, d: usize) !void { + try o.k.moe.combine(o.s, y, if (y_bf16) .bf16 else .f32, wts, out, .bf16, rows, s, d); + } +}; diff --git a/zig/src/families/kolibri1/tree.zig b/zig/src/families/kolibri1/tree.zig new file mode 100644 index 0000000000..93b972d149 --- /dev/null +++ b/zig/src/families/kolibri1/tree.zig @@ -0,0 +1,66 @@ +//! The host side of cuda/tree_attention for serial chains: each row's stream, each stream's table row, the parents. +const std = @import("std"); + +/// The longest chain a stream's window may hold (cuda/tree_attention's path limit). +pub const max_rows = 128; + +/// A stream's window: its rows and the keys it has committed before them. +pub const Window = struct { rows: usize, committed: usize }; + +/// int32 lists in one buffer: row_stream [n], streams [s][4] (first row, rows, committed, 0), parents [n]. +pub const Plan = struct { + flat: []i32, + rows: usize, + streams: usize, + + /// Byte offsets of the three lists inside `flat`. + pub fn streamsAt(p: Plan) usize { + return p.rows * 4; + } + + pub fn parentsAt(p: Plan) usize { + return (p.rows + 4 * p.streams) * 4; + } +}; + +/// The plan of `windows`, each a chain (row 0's parent -1, row r's r - 1, global rows); the caller frees `flat`. +pub fn build(gpa: std.mem.Allocator, windows: []const Window) !Plan { + var n: usize = 0; + for (windows) |w| { + if (w.rows < 1 or w.rows > max_rows) return error.WindowWidth; + n += w.rows; + } + const flat = try gpa.alloc(i32, 2 * n + 4 * windows.len); + var start: usize = 0; + for (windows, 0..) |w, s| { + for (0..w.rows) |r| { + flat[start + r] = @intCast(s); + flat[n + 4 * windows.len + start + r] = if (r == 0) -1 else @intCast(start + r - 1); + } + for ([_]usize{ start, w.rows, w.committed, 0 }, 0..) |v, k| flat[n + 4 * s + k] = @intCast(v); + start += w.rows; + } + return .{ .flat = flat, .rows = n, .streams = windows.len }; +} + +/// The stream row a one-row step writes before its graph replays: first row 0, one row, `committed` keys. +pub fn stepRow(committed: usize) [4]i32 { + return .{ 0, 1, @intCast(committed), 0 }; +} + +test "two chains: rows' streams, the streams' table and global parents" { + const gpa = std.testing.allocator; + const p = try build(gpa, &.{ .{ .rows = 2, .committed = 600 }, .{ .rows = 1, .committed = 40 } }); + defer gpa.free(p.flat); + const want = [_]i32{ 0, 0, 1, 0, 2, 600, 0, 2, 1, 40, 0, -1, 0, -1 }; + try std.testing.expectEqualSlices(i32, &want, p.flat); + try std.testing.expectEqual(@as(usize, 12), p.streamsAt()); + try std.testing.expectEqual(@as(usize, 44), p.parentsAt()); +} + +test "a window past the path limit, or an empty one, is refused" { + const gpa = std.testing.allocator; + try std.testing.expectError(error.WindowWidth, build(gpa, &.{.{ .rows = max_rows + 1, .committed = 0 }})); + try std.testing.expectError(error.WindowWidth, build(gpa, &.{.{ .rows = 0, .committed = 0 }})); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 77, 0 }, &stepRow(77)); +} From 0a267b0b8fd4eb1c8dd199c72d711533fc066833 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:53:25 +0200 Subject: [PATCH 26/41] kolibri1: the lane backend (prompt steps, shared verify windows, keep) and prompt-cache states at any row --- zig/src/families/kolibri1/backend.zig | 363 ++++++++++++++++++++++++++ zig/src/families/kolibri1/cuda.zig | 4 + zig/src/families/kolibri1/snap.zig | 109 ++++++++ 3 files changed, 476 insertions(+) create mode 100644 zig/src/families/kolibri1/backend.zig create mode 100644 zig/src/families/kolibri1/snap.zig diff --git a/zig/src/families/kolibri1/backend.zig b/zig/src/families/kolibri1/backend.zig new file mode 100644 index 0000000000..1f5154330f --- /dev/null +++ b/zig/src/families/kolibri1/backend.zig @@ -0,0 +1,363 @@ +//! Kolibri-1 behind lanes.Backend: a cache slot a stream, prompts a chunk a step, one forward a round of windows. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const fwd = @import("forward.zig"); +const kern = @import("kernels.zig"); +const snap = @import("snap.zig"); + +const be = lanes.backend; +const Stream = lanes.Stream; + +/// Prompt rows a step takes: the forward's chunk grid, from position 0. +pub const prompt_rows = fwd.prompt_rows; +/// Window rows a round takes over every stream. +pub const round_rows = fwd.round_rows; +/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. +pub const max_window = 16; +/// Drawn tokens a handle names (a ring; read before the round loop could wrap it). +const ring = 1024; + +const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's fp32 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts +}; + +pub const Kolibri = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + m: *fwd.Model, + k: *const kern.Kernels, + s: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + logits: cuda.DeviceBuffer, // fp32 [round_rows, V] + logits16: cuda.DeviceBuffer, // their bf16 copy for keyed draws + tok: cuda.DeviceBuffer, // drawn tokens: the handle ring, then a round's rows + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot + next: u64 = 0, + steps: []?fwd.Model.Step, // a slot's one-row step, captured at its first use (null: not yet) + graphed: bool, // TENSORFOLD_KOLIBRI_GRAPHS=0 runs every step eager + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *fwd.Model, k: *const kern.Kernels, s: cuda.Stream) !Kolibri { + const v = m.c.vocab; + const free = try gpa.alloc(bool, m.slots); + errdefer gpa.free(free); + @memset(free, true); + var logits = try cuda.DeviceBuffer.alloc(d, round_rows * v * 4); + errdefer logits.free(); + var logits16 = try cuda.DeviceBuffer.alloc(d, round_rows * v * 2); + errdefer logits16.free(); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + errdefer meta.free(); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slots); + errdefer gpa.free(bufs); + var made: usize = 0; + errdefer for (bufs[0..made]) |*b| b.free(); + for (bufs, 0..) |*b, i| { + b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) v * 4 else @sizeOf(cuda.sample.Rule)); + made += 1; + } + const steps = try gpa.alloc(?fwd.Model.Step, m.slots); + @memset(steps, null); + return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .logits = logits, .logits16 = logits16, .tok = tok, .meta = meta, .slot_bufs = bufs, .steps = steps, .graphed = !envOff("TENSORFOLD_KOLIBRI_GRAPHS") }; + } + + pub fn deinit(x: *Kolibri) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| l.rows.deinit(x.gpa); + x.lanes.deinit(x.gpa); + for (x.slot_bufs) |*b| b.free(); + x.gpa.free(x.slot_bufs); + for (x.steps) |*g| if (g.*) |*e| e.exec.deinit(); + x.gpa.free(x.steps); + x.gpa.free(x.free); + x.logits.free(); + x.logits16.free(); + x.tok.free(); + x.meta.free(); + } + + pub fn backend(x: *Kolibri) be.Backend { + return .{ .ptr = x, .vtable = &vtable }; + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true }; + + /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a forward. + pub fn facts(x: *const Kolibri) lanes.Model { + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slots, round_rows)) }; + } + + fn envOff(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return false; + return std.mem.eql(u8, std.mem.span(v), "0"); + } + + fn self(ptr: *anyopaque) *Kolibri { + return @ptrCast(@alignCast(ptr)); + } + + fn ops(x: *const Kolibri) kern.Ops { + return .{ .k = x.k, .s = x.s }; + } + + fn upload(x: *Kolibri, dst: u64, bytes: []const u8) !void { + try x.d.check(x.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, x.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Rows [row0, row0 + rows) of `logits` drawn into tok[at..], row r at `position0 + r` (argmax when greedy). + fn drawRows(x: *Kolibri, l: *const Lane, logits: u64, row0: usize, rows: usize, position0: u64, at: usize, meta_at: usize) !void { + const v = x.m.c.vocab; + const o = x.ops(); + const src = logits + row0 * v * 4; + if (!l.sampled) return o.argmax(src, rows, v, x.tok.ptr + at * 4); + try o.toBf16(src, x.logits16.ptr, rows * v); + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try x.upload(mp, std.mem.asBytes(&meta)); + try x.k.draws.draw(x.s, x.logits16.ptr, v, l.rule, mp, 0, x.tok.ptr + at * 4, rows, null, null); + } + + /// The stream's cache state at `at` for the prompt cache, copied behind the chunk that ends there. + pub fn saveState(x: *Kolibri, owner: ?*anyopaque, at: u32) !*snap.Saved { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if ((l.fill orelse l.len) != at) return error.NotAtMark; // a state is copied only where the pass stands + return snap.save(x.gpa, x.m, x.s, l.slot, at); + } + + /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). + pub fn restoreState(x: *Kolibri, owner: ?*anyopaque, k: *const snap.Saved) !void { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + try snap.restore(x.m, x.s, l.slot, k); + l.len = k.at; + } + + pub fn dropState(x: *Kolibri, k: *snap.Saved) void { + snap.drop(x.gpa, x.s, k); + } + + /// The lowest free slot. + fn freeSlot(x: *const Kolibri) ?usize { + return std.mem.indexOfScalar(bool, x.free, true); + } + + /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. + fn opening(x: *Kolibri, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (p.len + @as(usize, s.max_new) + max_window > x.m.context) return error.PastContext; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + const slot = x.freeSlot() orelse { + x.lanes.removeByPtr(got.key_ptr); + return error.NoFreeSlot; + }; + x.free[slot] = false; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr }; + } + const l = got.value_ptr; + l.sampled = false; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try x.upload(l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + l.fill = x.resumeAt(s, l, p.len); + s.fill_at.store(@intCast(l.fill.?), .monotonic); + return l; + } + + /// The kept state's position, restored into the slot (0: none, or not this prompt's); any row resumes exactly. + fn resumeAt(x: *Kolibri, s: *Stream, l: *const Lane, len: usize) usize { + const saved = s.reuse.saved orelse return 0; + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len) { + s.reuse_failed = true; + return 0; + } + snap.restore(x.m, x.s, l.slot, k) catch { + s.reuse_failed = true; + return 0; + }; + s.cached = k.at; + return k.at; + } + + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written. + fn chunk(x: *Kolibri, s: *Stream, l: *Lane) !bool { + const p = s.prompt(); + const at = l.fill.?; + if (s.isCancelled()) return error.Cancelled; + var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows); + for (s.reuse.marks) |mk| if (mk > at and mk < end) { + end = mk; + }; + try x.m.forward(x.ops(), &.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, null, l.last); + s.fill_at.store(@intCast(end), .monotonic); + const done = end == p.len; + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l); + } + + fn handle(x: *Kolibri) usize { + const h = x.next; + x.next += 1; + return @intCast(h % ring); + } + + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.next; + try x.drawRows(l, l.last, 0, 1, position, x.handle(), 0); + return h; + } + + fn value(x: *Kolibri, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| read(x, h), + }; + } + + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const t = try x.value(feed); + if (position != l.len + 1) return error.PositionMismatch; + const logits = try x.step(l, t); + l.len += 1; + const h = x.next; + try x.drawRows(l, logits, 0, 1, position, x.handle(), 0); + return h; + } + + /// One row at the lane's cache length: the slot's captured graph (captured at its first step), else eager. + fn step(x: *Kolibri, l: *const Lane, token: u32) !u64 { + if (!x.graphed) { + try x.m.forward(x.ops(), &.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{token} }}, false, null, x.logits.ptr); + return x.logits.ptr; + } + const g = &x.steps[l.slot]; + if (g.* == null) g.* = try x.m.captureStep(x.ops(), l.slot); + try x.m.step(x.ops(), &g.*.?, token, l.len); + return g.*.?.logits; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + const x = self(ptr); + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.s.synchronize(); + try x.tok.download(@intCast((h % ring) * 4), std.mem.sliceAsBytes(&out)); + return out[0]; + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + const chains = try x.gpa.alloc(fwd.Chain, windows.len); + defer x.gpa.free(chains); + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (w.rows() > max_window) return error.WindowTooWide; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items }; + total += l.rows.items.len; + } + if (total > round_rows) return error.RoundTooWide; + var logits = x.logits.ptr; + if (total == 1) { + logits = try x.step(x.lanes.getPtr(windows[0].stream).?, chains[0].tokens[0]); // a lone row: the step + } else { + const rows = try x.gpa.alloc(usize, total); + defer x.gpa.free(rows); + for (rows, 0..) |*r, i| r.* = i; + try x.m.forward(x.ops(), chains, false, rows, logits); + } + var at: usize = 0; + for (windows, 0..) |w, i| { + const l = x.lanes.getPtr(w.stream).?; + try x.drawRows(l, logits, at, l.rows.items.len, w.positions[0], ring + at, i); + at += l.rows.items.len; + } + const got = try x.gpa.alloc(u32, total); + defer x.gpa.free(got); + try x.s.synchronize(); + try x.tok.download(ring * 4, std.mem.sliceAsBytes(got)); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + l.len = l.base + n; // a chain keeps every row unless rolled back + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + for (windows, paths) |w, path| { + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote + } + } + + /// No draft head is loaded: absorbing is a no-op, and a request for drafts is refused by name. + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + fn dropLane(x: *Kolibri, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + } + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } +}; diff --git a/zig/src/families/kolibri1/cuda.zig b/zig/src/families/kolibri1/cuda.zig index 85e9250575..dece655e8d 100644 --- a/zig/src/families/kolibri1/cuda.zig +++ b/zig/src/families/kolibri1/cuda.zig @@ -7,6 +7,8 @@ pub const weights = @import("weights.zig"); pub const tree = @import("tree.zig"); pub const kernels = @import("kernels.zig"); pub const forward = @import("forward.zig"); +pub const snap = @import("snap.zig"); +pub const backend = @import("backend.zig"); test { _ = @import("config.zig"); @@ -16,4 +18,6 @@ test { _ = @import("tree.zig"); _ = @import("kernels.zig"); _ = @import("forward.zig"); + _ = @import("snap.zig"); + _ = @import("backend.zig"); } diff --git a/zig/src/families/kolibri1/snap.zig b/zig/src/families/kolibri1/snap.zig new file mode 100644 index 0000000000..0c7961455d --- /dev/null +++ b/zig/src/families/kolibri1/snap.zig @@ -0,0 +1,109 @@ +//! A stream's cache state at a prompt position, for the core prompt cache: full layers' keys, sliding layers' window. +const std = @import("std"); +const cuda = @import("cuda"); +const fwd = @import("forward.zig"); + +/// A kept state: the keys of positions [0, at) of full layers, then the window's last keys of each sliding layer. +pub const Saved = struct { at: u32, buf: cuda.DeviceBuffer }; + +/// Positions a sliding layer keeps at `at`: the next row reads (at - window, at], so window - 1 before it. +fn tail(m: *const fwd.Model, at: usize) usize { + return @min(at, m.c.window - 1); +} + +fn rowBytes(m: *const fwd.Model) usize { + return m.c.kv_heads * m.c.head_dim * 2; +} + +/// Bytes a state at `at` takes. +pub fn bytes(m: *const fwd.Model, at: usize) u64 { + var n: usize = 0; + for (m.c.full[0..m.c.layers]) |full| n += 2 * (if (full) at else tail(m, at)); + return n * rowBytes(m); +} + +/// Each layer's cache, K then V, and positions [first, at) of `slot`: one copy a run of ring rows, in order. +fn each(m: *const fwd.Model, slot: usize, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + const rb = rowBytes(m); + var off: usize = 0; + for (0..m.c.layers) |i| for ([_][]u64{ m.k, m.v }) |cache| { + const full = m.c.full[i]; + const len = if (full) m.context else m.ring; + const base = cache[i] + slot * len * rb; + var p = if (full) 0 else at - tail(m, at); + while (p < at) { + const row = p % len; + const run = @min(at - p, len - row); + try f(ctx, base + row * rb, off, run * rb); + off += run * rb; + p += run; + } + }; +} + +const Copy = struct { d: *const cuda.Driver, s: cuda.Stream, buf: u64, out: bool }; + +fn copy(c: Copy, dev: u64, off: usize, len: usize) anyerror!void { + const src = if (c.out) dev else c.buf + off; + const dst = if (c.out) c.buf + off else dev; + try c.d.check(c.d.api.cuMemcpyDtoDAsync_v2(dst, src, len, c.s.handle), "cuMemcpyDtoDAsync"); +} + +/// Copy `slot`'s state at `at` into new storage, queued on `s` behind the chunk that ends there. +pub fn save(gpa: std.mem.Allocator, m: *const fwd.Model, s: cuda.Stream, slot: usize, at: u32) !*Saved { + if (at == 0 or at > m.context) return error.NotAtMark; + const k = try gpa.create(Saved); + errdefer gpa.destroy(k); + k.* = .{ .at = at, .buf = try cuda.DeviceBuffer.alloc(m.d, bytes(m, at)) }; + errdefer k.buf.free(); + try each(m, slot, at, Copy{ .d = m.d, .s = s, .buf = k.buf.ptr, .out = true }, copy); + return k; +} + +/// Make `slot`'s caches hold `k`'s state, queued on `s` before the next chunk. +pub fn restore(m: *const fwd.Model, s: cuda.Stream, slot: usize, k: *const Saved) !void { + if (k.at > m.context) return error.PastContext; + try each(m, slot, k.at, Copy{ .d = m.d, .s = s, .buf = k.buf.ptr, .out = false }, copy); +} + +/// Free a state once the stream's queued copies of it are done. +pub fn drop(gpa: std.mem.Allocator, s: cuda.Stream, k: *Saved) void { + s.synchronize() catch {}; + k.buf.free(); + gpa.destroy(k); +} + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +test "a state holds full layers' positions and each sliding ring's window, wrapped runs in position order" { + const Config = @import("config.zig").Config; + var c: Config = .{ .layers = 2, .hidden = 128, .heads = 4, .kv_heads = 1, .head_dim = 128, .vocab = 8, .experts = 2, .top_k = 1, .moe_width = 128, .shared_width = 128, .window = 5, .rope_theta = 1e4, .eps = 1e-6 }; + c.full[1] = true; + var k = [_]u64{ 1 << 20, 2 << 20 }; + var v = [_]u64{ 3 << 20, 4 << 20 }; + var m: fwd.Model = undefined; + m.c = c; + m.context = 64; + m.ring = 8; + m.k = &k; + m.v = &v; + const rb = 256; + try testing.expectEqual(@as(u64, 2 * (4 + 10) * rb), bytes(&m, 10)); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&m, 1, 10, Count{ .runs = &runs }, Count.add); + // sliding K: positions 6..9 at ring rows 6, 7 then 0, 1 of slot 1; then V; then full K and V, positions 0..9 + const want = [_][3]usize{ + .{ (1 << 20) + (8 + 6) * rb, 0, 2 * rb }, .{ (1 << 20) + 8 * rb, 2 * rb, 2 * rb }, + .{ (3 << 20) + (8 + 6) * rb, 4 * rb, 2 * rb }, .{ (3 << 20) + 8 * rb, 6 * rb, 2 * rb }, + .{ (2 << 20) + 64 * rb, 8 * rb, 10 * rb }, .{ (4 << 20) + 64 * rb, 18 * rb, 10 * rb }, + }; + try testing.expectEqualSlices([3]usize, &want, runs.items); +} From 6498344d33a19cce0b142af9cc26a8588be0f761 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:53:25 +0200 Subject: [PATCH 27/41] native: kolibri1 in the CUDA family registry, its refusals in words, --parallel streams at open, no-lookahead states --- zig/build/cuda.zig | 9 ++ zig/src/families/kolibri1/cuda.zig | 3 + zig/src/families/kolibri1/native.zig | 142 +++++++++++++++++++++++++++ zig/src/native/cuda.zig | 4 +- zig/src/native/cuda_refusal.zig | 13 ++- zig/src/native/cuda_snaps.zig | 7 +- 6 files changed, 175 insertions(+), 3 deletions(-) create mode 100644 zig/src/families/kolibri1/native.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 7b63281ff6..3ce658e76d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -162,6 +162,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); runner.addImport("glm_cuda", glm_cuda); + runner.addImport("kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes)); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @@ -181,6 +182,11 @@ fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return m; } +/// The Kolibri-1 CUDA family over the runtime, the core and the lane core. +fn kolibri1Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/kolibri1/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); +} + /// LFM2 (d1's decision model) on CUDA: config, weights, kernels and the forward that scores labels. fn lfm2Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module) *std.Build.Module { return b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm2/lfm2.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } }); @@ -207,6 +213,7 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "lfm2", .module = lfm2Module(b, target, optimize, cuda, core) }, .{ .name = "core", .module = core }, .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "kolibri1", .module = kolibri1Module(b, target, optimize, cuda, core, lanes) }, }, }); return .{ .api = api, .engines = mod }; @@ -280,6 +287,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); + // the Kolibri-1 CUDA family's host side + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes) })).step); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = shares })).step); diff --git a/zig/src/families/kolibri1/cuda.zig b/zig/src/families/kolibri1/cuda.zig index dece655e8d..5d2a0a9101 100644 --- a/zig/src/families/kolibri1/cuda.zig +++ b/zig/src/families/kolibri1/cuda.zig @@ -1,5 +1,6 @@ //! Kolibri-1 (block FP8) on CUDA: config, weights, the layer graph on shared operators, the lane backend. +pub const Config = @import("config.zig").Config; pub const config = @import("config.zig"); pub const names = @import("names.zig"); pub const pack = @import("pack.zig"); @@ -9,6 +10,7 @@ pub const kernels = @import("kernels.zig"); pub const forward = @import("forward.zig"); pub const snap = @import("snap.zig"); pub const backend = @import("backend.zig"); +pub const native = @import("native.zig"); test { _ = @import("config.zig"); @@ -20,4 +22,5 @@ test { _ = @import("forward.zig"); _ = @import("snap.zig"); _ = @import("backend.zig"); + _ = @import("native.zig"); } diff --git a/zig/src/families/kolibri1/native.zig b/zig/src/families/kolibri1/native.zig new file mode 100644 index 0000000000..c4756f11b6 --- /dev/null +++ b/zig/src/families/kolibri1/native.zig @@ -0,0 +1,142 @@ +//! Kolibri-1 behind the native CUDA registry (native/cuda.zig): open, the prompt cache's states, refusals in words. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const weights = @import("weights.zig"); +const kern = @import("kernels.zig"); +const fwd = @import("forward.zig"); +const be = @import("backend.zig"); +const snap = @import("snap.zig"); + +pub const model_type = "kolibri1"; +pub const formats: []const []const u8 = &.{"fp8-block"}; +pub const default_context: i64 = 65536; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = be.prompt_rows; +/// A kept state is the caches through its position and nothing past it (no draft head reads a token on). +pub const snap_lookahead: u32 = 0; +/// No drafter yet: --drafter is refused before the load. +pub const takes_drafter = false; +pub const drafter_refusal = "--drafter: Kolibri-1 on CUDA loads no drafter yet; serve it without --drafter"; + +/// `explain` words this family's open refusals too. +pub const explains_open = true; + +/// No `vision` or `precision` field: the server refuses --vision and --precision for this family by name. +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8 }; + +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; + +/// What the native server drives; every stream's caches are allocated at open (`free_streams` of them). +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + learn_bytes: usize = 0, + target: *be.Kolibri, // the prompt cache keeps this backend's states (snap_*) +}; + +const Owned = struct { gpa: std.mem.Allocator, w: weights.Weights, k: kern.Kernels, s: cuda.Stream, m: fwd.Model, kb: be.Kolibri }; + +/// Streams whose caches the open allocates: the server's --parallel, at most a round's rows. +fn streams(lanes_asked: u32) usize { + return std.math.clamp(@as(usize, lanes_asked), 1, be.round_rows); +} + +/// Positions a stream's caches hold: the window, whole 1,024s. +fn positions(context: usize) usize { + return (context + 1023) / 1024 * 1024; +} + +/// The checkpoint, the shared kernels and every stream's caches. +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.w = try weights.load(gpa, io, ctx.d, dir, null); + errdefer own.w.deinit(); + own.k = try kern.Kernels.load(ctx); + errdefer own.k.deinit(); + own.s = try cuda.Stream.init(ctx.d, false); + errdefer own.s.deinit(); + const context = positions(o.context); + const n = streams(o.lanes); + own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, n, be.prompt_rows); + errdefer own.m.deinit(); + own.kb = try be.Kolibri.init(gpa, ctx.d, &own.m, &own.k, own.s); + return .{ + .backend = own.kb.backend(), + .facts = own.kb.facts(), + .rows = be.max_window, + .stream_bytes = fwd.Model.streamBytes(own.w.config, context), + .free_streams = @intCast(n), + .ctx = own, + .deinit = release, + .target = &own.kb, + }; +} + +fn kolibri(ptr: *anyopaque) *be.Kolibri { + return @ptrCast(@alignCast(ptr)); +} + +/// Bytes a state at `at` prompt tokens takes. +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(kolibri(ptr).m, at); +} + +/// The prompt pass stands at `at` (a chunk start): copy the stream's caches there. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + return @ptrCast(try kolibri(ptr).saveState(owner, at)); +} + +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + return kolibri(ptr).restoreState(owner, @ptrCast(@alignCast(saved))); +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + kolibri(ptr).dropState(@ptrCast(@alignCast(saved))); +} + +/// A request or load this engine refuses, in words; null: none of its own. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => "Kolibri-1 on CUDA has no draft head yet: drafts come from the prompt's own text only", + error.StructuresUnsupported => "structured output is not built for Kolibri-1 on CUDA yet", + error.CapturedKernels => "Kolibri-1 runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + own.kb.deinit(); + own.m.deinit(); + own.s.deinit(); + own.k.deinit(); + own.w.deinit(); + own.gpa.destroy(own); +} + +test "the prompt cache's hooks compile against the store's vtable" { + const pc: struct { bytes: *const fn (*anyopaque, u32) u64, save: *const fn (*anyopaque, ?*anyopaque, u32) anyerror!*anyopaque, restore: *const fn (*anyopaque, ?*anyopaque, *anyopaque) anyerror!void, drop: *const fn (*anyopaque, *anyopaque) void } = .{ .bytes = snap_bytes, .save = snap_save, .restore = snap_restore, .drop = snap_drop }; + try std.testing.expect(@intFromPtr(pc.save) != 0); +} + +test "every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.CapturedKernels }) |e| { + try std.testing.expect(explain(null, e) != null); + } + try std.testing.expect(explain(null, error.OutOfMemory) == null); + try std.testing.expectEqual(@as(usize, 65536), positions(65000)); + try std.testing.expectEqual(@as(usize, 1), streams(0)); + try std.testing.expectEqual(@as(usize, be.round_rows), streams(1000)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..331046a9fa 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -6,6 +6,7 @@ const lanes = @import("lanes"); const nemotron = @import("nemotron"); const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); +const kolibri1 = @import("kolibri1"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -15,7 +16,7 @@ const cuda_disk = @import("cuda_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, kolibri1.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; @@ -335,6 +336,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O opts.precision = std.meta.stringToEnum(@typeInfo(@TypeOf(opts.precision)).optional.child, @tagName(p)).?; } if (@hasField(F.Options, "room")) opts.room = .{ .ctx = &probe, .read = Probe.read }; + if (@hasField(F.Options, "lanes")) opts.lanes = o.lanes; // a family that allocates every stream's caches at open if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index 7cdf93f18f..c9ff715f26 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -21,7 +21,10 @@ const lfm2_vision = "--vision: LFM2 reads the checkpoint's images through its ow /// A family takes --drafter only if it says so, and --learn-conversations only as Nemotron; null: `o` may open `F`. pub fn family(comptime F: type, o: api.Open) ?[]const u8 { - if (o.drafter != null and !(@hasDecl(F, "takes_drafter") and F.takes_drafter)) return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; + if (o.drafter != null and !(@hasDecl(F, "takes_drafter") and F.takes_drafter)) { + if (@hasDecl(F, "drafter_refusal")) return F.drafter_refusal; // a family without an MTP head says so itself + return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; + } if (o.learn_conversations and F != nemotron.native) return "--learn-conversations keeps conversation states on disk for Nemotron only on CUDA: serve this model without it"; return null; } @@ -63,6 +66,14 @@ test "a family that explicitly declines a drafter refuses it" { try std.testing.expectStringStartsWith(family(Declined, .{ .dir = "d", .model_type = "m", .drafter = "draft" }) orelse "", "--drafter loads a separate drafter"); } +test "a family with no head of its own words its --drafter refusal" { + const Headless = struct { + pub const takes_drafter = false; + pub const drafter_refusal = "--drafter: none yet"; + }; + try std.testing.expectEqualStrings("--drafter: none yet", family(Headless, .{ .dir = "d", .model_type = "m", .drafter = "x" }) orelse ""); +} + test "--learn-conversations opens Nemotron only, the one CUDA family whose states have a disk image" { const Other = struct {}; try std.testing.expectStringStartsWith(family(Other, .{ .dir = "d", .model_type = "m", .learn_conversations = true }) orelse "", "--learn-conversations keeps"); diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index bbaedc2c57..d02541953d 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -2,9 +2,10 @@ const std = @import("std"); const api = @import("engine_api"); -/// Any row, the head one token on (Nemotron's); `snap_grid`'s chunk starts and `warm_quiet`'s warm turns if declared. +/// Any row, the head one token on (Nemotron's) unless `snap_lookahead` says; `snap_grid` and `warm_quiet` if declared. pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { var rules: api.prompt_cache.Rules = .{ .lookahead = 1, .min_prompt = 0 }; + if (@hasDecl(F, "snap_lookahead")) rules.lookahead = F.snap_lookahead; // a family with no head reads no token on if (@hasDecl(F, "warm_quiet")) rules.warm = .waits; if (@hasDecl(F, "snap_grid")) { rules.planned = true; @@ -39,6 +40,10 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as const Warm = struct { pub const warm_quiet: u64 = 1; }; + const Headless = struct { + pub const snap_lookahead: u32 = 0; + }; + try std.testing.expectEqual(@as(u32, 0), cacheRules(Headless).lookahead); const any = cacheRules(Other); try std.testing.expect(!any.planned and any.min_prompt == 0 and any.lookahead == 1 and any.warm == .off); try std.testing.expect(cacheRules(Warm).warm == .waits and cacheRules(Glm).warm == .off); From 0101e1ff0a2963f171d51e01371899ece2a9d144 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:53:26 +0200 Subject: [PATCH 28/41] tests: kolibri1 on the checkpoint: layer 0 against float64, the window edge, full layers, steps, drafted, shared rounds, cuts, kept states --- zig/tests/cuda/kolibri_layer_tests.zig | 187 +++++++++++++++++ zig/tests/cuda/kolibri_ref.zig | 181 ++++++++++++++++ zig/tests/cuda/kolibri_tests.zig | 275 +++++++++++++++++++++++++ zig/tests/cuda/main.zig | 12 ++ 4 files changed, 655 insertions(+) create mode 100644 zig/tests/cuda/kolibri_layer_tests.zig create mode 100644 zig/tests/cuda/kolibri_ref.zig create mode 100644 zig/tests/cuda/kolibri_tests.zig diff --git a/zig/tests/cuda/kolibri_layer_tests.zig b/zig/tests/cuda/kolibri_layer_tests.zig new file mode 100644 index 0000000000..9e5e6fd600 --- /dev/null +++ b/zig/tests/cuda/kolibri_layer_tests.zig @@ -0,0 +1,187 @@ +//! Kolibri-1's layers on the real checkpoint: layer 0 against float64, the window's edge, the full layers' reach. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const kolibri1 = @import("kolibri1"); +const check = @import("check.zig"); +const ref = @import("kolibri_ref.zig"); +const Gpu = check.Gpu; + +const fwd = kolibri1.forward; + +/// The first `layers` layers of the checkpoint, the shared kernels and `slots` streams' caches. +pub const Rig = struct { + gpu: Gpu, + w: kolibri1.weights.Weights, + k: kolibri1.kernels.Kernels, + s: cuda.Stream, + m: fwd.Model, + logits: cuda.DeviceBuffer, + + pub fn init(r: *Rig, gpu: Gpu, model: []const u8, count: ?usize, context: usize, slots: usize, rows: usize) !void { + r.gpu = gpu; + r.w = try kolibri1.weights.load(gpu.gpa, gpu.io, gpu.d, model, count); + errdefer r.w.deinit(); + r.k = try kolibri1.kernels.Kernels.load(gpu.ctx); + errdefer r.k.deinit(); + r.s = try cuda.Stream.init(gpu.d, false); + errdefer r.s.deinit(); + r.m = try fwd.Model.init(gpu.gpa, gpu.d, r.s, &r.w, &r.k, context, slots, rows); + errdefer r.m.deinit(); + r.logits = try cuda.DeviceBuffer.alloc(gpu.d, rows * r.w.config.vocab * 4); + } + + pub fn deinit(r: *Rig) void { + r.logits.free(); + r.m.deinit(); + r.s.deinit(); + r.k.deinit(); + r.w.deinit(); + } + + pub fn ops(r: *const Rig) kolibri1.kernels.Ops { + return .{ .k = &r.k, .s = r.s }; + } + + /// The prompt in grid chunks into `slot`; the last row's fp32 logits (gpa-owned). + pub fn prompt(r: *Rig, slot: usize, ids: []const u32) ![]f32 { + var at: usize = 0; + while (at < ids.len) : (at += fwd.prompt_rows) { + const end = @min(ids.len, at + fwd.prompt_rows); + try r.m.forward(r.ops(), &.{.{ .slot = slot, .p0 = at, .tokens = ids[at..end] }}, true, null, r.logits.ptr); + } + return r.read(1); + } + + /// One decode row: `token` at position `p0` of `slot`; its logits. + pub fn decode(r: *Rig, slot: usize, p0: usize, token: u32) ![]f32 { + try r.m.forward(r.ops(), &.{.{ .slot = slot, .p0 = p0, .tokens = &.{token} }}, false, null, r.logits.ptr); + return r.read(1); + } + + pub fn read(r: *Rig, rows: usize) ![]f32 { + try r.s.synchronize(); + const out = try r.gpu.gpa.alloc(f32, rows * r.w.config.vocab); + try r.logits.download(0, std.mem.sliceAsBytes(out)); + return out; + } +}; + +/// Prompt ids from a file of integers (commas or whitespace). +pub fn readIds(gpu: Gpu, path: []const u8) ![]u32 { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + var out: std.ArrayList(u32) = .empty; + errdefer out.deinit(gpu.gpa); + var it = std.mem.tokenizeAny(u8, text, ", \n\r\t[]"); + while (it.next()) |t| try out.append(gpu.gpa, try std.fmt.parseInt(u32, t, 10)); + return out.toOwnedSlice(gpu.gpa); +} + +fn argmax(x: []const f32) usize { + var best: usize = 0; + for (x, 0..) |v, i| if (v > x[best]) { + best = i; + }; + return best; +} + +fn same(a: []const f32, b: []const f32) bool { + return std.mem.eql(u8, std.mem.sliceAsBytes(a), std.mem.sliceAsBytes(b)); +} + +/// Layer 0 and the head against float64 (routing bias, the shared expert, FP8 block scales): decode and prompt rows. +pub fn reference(gpu: Gpu, model: []const u8, ids_file: []const u8) !void { + const ids = try readIds(gpu, ids_file); + defer gpu.gpa.free(ids); + var rig: Rig = undefined; + try rig.init(gpu, model, 1, 4096, 2, fwd.prompt_rows); + defer rig.deinit(); + var ck = try core.checkpoint.Checkpoint.openModel(gpu.gpa, gpu.io, model); + defer ck.close(); + var map = try kolibri1.names.map(gpu.gpa, &ck, rig.w.config); + defer map.deinit(); + var c = rig.w.config; + c.layers = 1; + var worst: f64 = 0; + var bad: usize = 0; + const n = 24; + for (0..3) |case| { + const toks = ids[case * 200 ..][0..n]; + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const want = try ref.lastLogits(arena.allocator(), &map, c, toks); + const whole = try rig.prompt(0, toks); + defer gpu.gpa.free(whole); + const head = try rig.prompt(1, toks[0 .. n - 1]); + gpu.gpa.free(head); + const step = try rig.decode(1, n - 1, toks[n - 1]); + defer gpu.gpa.free(step); + var scale: f64 = 0; + for (want) |v| scale = @max(scale, @abs(v)); + for ([_][]const f32{ whole, step }, [_][]const u8{ "prompt", "decode" }) |got, kind| { + var most: f64 = 0; + for (got, want) |g, w| most = @max(most, @abs(@as(f64, g) - w)); + const top = argmax(got) == std.mem.indexOfMax(f64, want); + std.debug.print("RESULT kolibri1 reference case {d} {s}: max |logit - f64| {e:.3} of {d:.3}, top-1 {s}\n", .{ case, kind, most, scale, if (top) "equal" else "DIFFERS" }); + worst = @max(worst, most / scale); + if (!top or most > 0.02 * scale) bad += 1; + } + } + try check.expect(bad == 0, "kolibri1 reference: {d} rows past 2% of the logits' scale or a different top-1", .{bad}); + check.pass("kolibri1 reference: layer 0 and the head within {d:.4} of the logits' scale of float64, top-1 equal", .{worst}); +} + +/// The window is 513 keys in prompt and decode rows (layer 0), and layer 4's full attention reaches key 0. +pub fn layers(gpu: Gpu, model: []const u8, ids_file: []const u8) !void { + const ids = try readIds(gpu, ids_file); + defer gpu.gpa.free(ids); + var rig: Rig = undefined; + try rig.init(gpu, model, 5, 4096, 2, fwd.prompt_rows); + defer rig.deinit(); + const win = rig.w.config.window; + const n = 1200; + var base: [n]u32 = undefined; + @memcpy(&base, ids[0..n]); + // one layer: the last row (position n - 1) reads keys n - win through n - 1, nothing before + var one: Rig = undefined; + try one.init(gpu, model, 1, 4096, 2, fwd.prompt_rows); + defer one.deinit(); + var bad: usize = 0; + for ([_]usize{ n - 1 - win, n - win }, [_]bool{ false, true }) |at, moves| { + var edit = base; + edit[at] = if (edit[at] == 1000) 1001 else 1000; + for ([_]bool{ true, false }) |prompt_rows| { + const a = try lastRow(&one, &base, prompt_rows); + defer gpu.gpa.free(a); + const b = try lastRow(&one, &edit, prompt_rows); + defer gpu.gpa.free(b); + const changed = !same(a, b); + std.debug.print("RESULT kolibri1 window: key {d} back {s} the last row ({s} row)\n", .{ n - 1 - at, if (changed) "moves" else "does not move", if (prompt_rows) "prompt" else "decode" }); + if (changed != moves) bad += 1; + } + } + // five layers, four sliding: they reach 4 * (win - 1) keys back, so key 0 moves row 2,999 only through layer 4 + const far = 3000; + const longer = try gpu.gpa.alloc(u32, far); + defer gpu.gpa.free(longer); + @memcpy(longer, ids[0..far]); + const a = try lastRow(&rig, longer, true); + defer gpu.gpa.free(a); + longer[0] = if (longer[0] == 1000) 1001 else 1000; + const b = try lastRow(&rig, longer, true); + defer gpu.gpa.free(b); + const reach = !same(a, b); + std.debug.print("RESULT kolibri1 full layer: key 0 {s} row {d} through five layers (layer 4 full)\n", .{ if (reach) "moves" else "does not move", far - 1 }); + if (!reach) bad += 1; + try check.expect(bad == 0, "kolibri1 layers: {d} window or full-layer checks wrong", .{bad}); + check.pass("kolibri1 layers: the window holds exactly {d} keys (prompt and decode rows), layer 4 reaches key 0", .{win}); +} + +/// The last row's logits of `ids`: all as prompt rows, or all but the last then the last as a decode row. +fn lastRow(r: *Rig, ids: []const u32, prompt_rows: bool) ![]f32 { + if (prompt_rows) return r.prompt(0, ids); + const head = try r.prompt(1, ids[0 .. ids.len - 1]); + r.gpu.gpa.free(head); + return r.decode(1, ids.len - 1, ids[ids.len - 1]); +} diff --git a/zig/tests/cuda/kolibri_ref.zig b/zig/tests/cuda/kolibri_ref.zig new file mode 100644 index 0000000000..5f66c39064 --- /dev/null +++ b/zig/tests/cuda/kolibri_ref.zig @@ -0,0 +1,181 @@ +//! Kolibri-1's first layer and head in float64 from the checkpoint's own tensors, rounded where the engine rounds. +const std = @import("std"); +const core = @import("core"); +const kolibri1 = @import("kolibri1"); + +const Tensor = core.checkpoint.Tensor; +const names = kolibri1.names; + +fn bf16(t: Tensor, i: usize) f64 { + return @as(f32, @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16)); +} + +fn f32At(t: Tensor, i: usize) f64 { + return @as(f32, @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little))); +} + +/// x rounded to bf16 (nearest even), as the engine's bf16 stores round. +fn roundBf16(x: f64) f64 { + const f: f32 = @floatCast(x); + const b: u32 = @bitCast(f); + const r = b +% 0x7fff +% ((b >> 16) & 1); + return @as(f32, @bitCast(r & 0xffff0000)); +} + +fn roundF32(x: f64) f64 { + return @as(f32, @floatCast(x)); +} + +/// The value of an e4m3 code (no NaN in a checkpoint's weights). +fn e4m3(code: u8) f64 { + const e: i32 = (code >> 3) & 15; + const m: f64 = @floatFromInt(code & 7); + const v = if (e == 0) m * 0x1p-9 else (8 + m) * std.math.ldexp(@as(f64, 1), e - 10); + return if (code & 0x80 != 0) -v else v; +} + +/// out[j] = sum over k of x[k] * dequantized w[j, k] (block-FP8: code times its 128 x 128 block's scale). +fn fp8Rows(p: names.Fp8, x: []const f64, out: []f64) void { + const k = p.w.dim(1); + const kb = (k + 127) / 128; + for (out, 0..) |*o, j| { + var acc: f64 = 0; + for (0..k) |c| acc += x[c] * e4m3(p.w.bytes[j * k + c]) * f32At(p.s, (j / 128) * kb + c / 128); + o.* = acc; + } +} + +/// rms(x) * w over `x`, in float64. +fn rms(x: []const f64, w: Tensor, eps: f64, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const r = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + for (out, x, 0..) |*o, v, i| o.* = v * r * bf16(w, i); +} + +/// The last row's logits of `tokens` at positions 0.. through layer 0, the final norm and the head (decode arithmetic). +pub fn lastLogits(a: std.mem.Allocator, m: *const names.Map, c: kolibri1.Config, tokens: []const u32) ![]f64 { + const n = tokens.len; + const d = c.hidden; + const h = c.heads; + const hk = c.kv_heads; + const dh = c.head_dim; + const L = m.layers[0]; + const res = try a.alloc(f64, n * d); + const x = try a.alloc(f64, n * d); + const q = try a.alloc(f64, n * h * dh); + const kk = try a.alloc(f64, n * hk * dh); + const vv = try a.alloc(f64, n * hk * dh); + const qkv = try a.alloc(f64, c.qkvDim()); + for (tokens, 0..) |t, r| { + for (0..d) |i| res[r * d + i] = bf16(m.embed, t * d + i); + rms(res[r * d ..][0..d], L.input_norm, c.eps, x[r * d ..][0..d]); + for (x[r * d ..][0..d]) |*v| v.* = roundBf16(v.*); + fp8Rows(L.q, x[r * d ..][0..d], qkv[0 .. h * dh]); + fp8Rows(L.k, x[r * d ..][0..d], qkv[h * dh ..][0 .. hk * dh]); + fp8Rows(L.v, x[r * d ..][0..d], qkv[(h + hk) * dh ..][0 .. hk * dh]); + for (qkv) |*v| v.* = roundBf16(v.*); + for (0..h + hk) |hd| { + const src = qkv[hd * dh ..][0..dh]; + const dst = if (hd < h) q[(r * h + hd) * dh ..][0..dh] else kk[(r * hk + hd - h) * dh ..][0..dh]; + var y: [256]f64 = undefined; + rms(src, if (hd < h) L.q_norm else L.k_norm, c.eps, y[0..dh]); + for (0..dh) |i| { // neox halves, sliding layers only (layer 0 is sliding) + const j = i % (dh / 2); + const ang = @as(f64, @floatFromInt(r)) / std.math.pow(f64, c.rope_theta, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(dh))); + const rot = if (i < dh / 2) -y[i + dh / 2] else y[i - dh / 2]; + dst[i] = if (c.full[0]) y[i] else roundBf16(y[i] * roundF32(@cos(ang)) + rot * roundF32(@sin(ang))); + } + if (c.full[0]) for (dst) |*v| { + v.* = roundBf16(v.*); + }; + } + @memcpy(vv[r * hk * dh ..][0 .. hk * dh], qkv[(h + hk) * dh ..][0 .. hk * dh]); + } + // the last row's attention over its window, then o, the residual and the MoE + const r = n - 1; + const att = try a.alloc(f64, h * dh); + const lo = if (r + 1 > c.window) r + 1 - c.window else 0; + const keys = try a.alloc(usize, r + 1 - lo); + for (keys, 0..) |*kx, i| kx.* = lo + i; + const scale = 1.0 / @sqrt(@as(f64, @floatFromInt(dh))); + for (0..h) |hd| { + const g = hd / (h / hk); + var sc = try a.alloc(f64, keys.len); + var most: f64 = -std.math.inf(f64); + for (keys, 0..) |p, i| { + var dot: f64 = 0; + for (0..dh) |e| dot += q[(r * h + hd) * dh + e] * kk[(p * hk + g) * dh + e]; + sc[i] = dot * scale; + most = @max(most, sc[i]); + } + var sum: f64 = 0; + for (sc) |*v| { + v.* = @exp(v.* - most); + sum += v.*; + } + for (0..dh) |e| { + var acc: f64 = 0; + for (keys, 0..) |p, i| acc += sc[i] * vv[(p * hk + g) * dh + e]; + att[hd * dh + e] = roundBf16(acc / sum); + } + } + const o = try a.alloc(f64, d); + fp8Rows(L.o, att, o); + for (o) |*v| v.* = roundBf16(v.*); + const res1 = res[r * d ..][0..d]; + const t1 = try a.alloc(f64, d); + rms(o, L.post_attn_norm, c.eps, t1); + for (res1, t1) |*v, u| v.* = roundF32(v.* + u); + const x2 = try a.alloc(f64, d); + rms(res1, L.pre_moe_norm, c.eps, x2); + for (x2) |*v| v.* = roundBf16(v.*); + const mo = try moe(a, L, c, x2); + rms(mo, L.post_moe_norm, c.eps, t1); + for (res1, t1) |*v, u| v.* = roundF32(v.* + u); + rms(res1, m.norm, c.eps, x2); + for (x2) |*v| v.* = roundBf16(v.*); + const logits = try a.alloc(f64, c.vocab); + for (logits, 0..) |*lg, j| { + var acc: f64 = 0; + for (0..d) |i| acc += x2[i] * bf16(m.head, j * d + i); + lg.* = acc; + } + return logits; +} + +/// Top-k experts by sigmoid logit + bias (lower id on ties), weights sigmoid(logit), the shared expert at weight 1. +fn moe(a: std.mem.Allocator, L: names.Layer, c: kolibri1.Config, x: []const f64) ![]f64 { + const d = c.hidden; + const ne = c.experts; + const logit = try a.alloc(f64, ne); + const key = try a.alloc(f64, ne); + for (logit, key, 0..) |*lg, *ky, e| { + var acc: f64 = 0; + for (0..d) |i| acc += x[i] * bf16(L.router, e * d + i); + lg.* = roundF32(acc); + ky.* = lg.* + bf16(L.bias, e); + } + const out = try a.alloc(f64, d); + @memset(out, 0); + const gate = try a.alloc(f64, c.moe_width); + const up = try a.alloc(f64, c.moe_width); + const y = try a.alloc(f64, d); + for (0..c.top_k + 1) |slot| { + var e: usize = ne; // the shared expert + var w: f64 = 1; + if (slot < c.top_k) { + e = std.mem.indexOfMax(f64, key); + w = roundF32(1.0 / (1.0 + @exp(-logit[e]))); + key[e] = -std.math.inf(f64); + } + const ex = L.experts[e]; + fp8Rows(ex.gate, x, gate); + fp8Rows(ex.up, x, up); + for (gate, up) |*g, u| g.* = roundBf16(g.* / (1.0 + @exp(-g.*)) * u); + fp8Rows(ex.down, gate, y); + for (out, y) |*o, v| o.* += w * roundF32(v); + } + for (out) |*o| o.* = roundBf16(o.*); + return out; +} diff --git a/zig/tests/cuda/kolibri_tests.zig b/zig/tests/cuda/kolibri_tests.zig new file mode 100644 index 0000000000..8817358a46 --- /dev/null +++ b/zig/tests/cuda/kolibri_tests.zig @@ -0,0 +1,275 @@ +//! Kolibri-1 on the real checkpoint: one-row steps, the lane rounds, shared rounds, prompt cuts and kept states. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kolibri1 = @import("kolibri1"); +const check = @import("check.zig"); +const layer_tests = @import("kolibri_layer_tests.zig"); +const Gpu = check.Gpu; + +const fwd = kolibri1.forward; +const Rig = layer_tests.Rig; +const readIds = layer_tests.readIds; + +fn hash(tokens: []const u32) u64 { + return std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(tokens)); +} + +fn argmax(x: []const f32) u32 { + var best: usize = 0; + for (x, 0..) |v, i| if (v > x[best]) { + best = i; + }; + return @intCast(best); +} + +const Case = struct { prompt: []const u32, sampling: ?lanes.Sampling = null, drafts: bool }; +const Run = struct { out: [][]u32, rate: f64, rounds: u64, accepted: u64 }; + +/// Cases admitted together into one lane engine on the backend, run to the end: tokens, rate, rounds, drafts landed. +fn lanesRun(gpu: Gpu, kb: *kolibri1.backend.Kolibri, cases: []const Case, count: u32) !Run { + const gpa = gpu.gpa; + var cfg = try lanes.Config.init(gpa, kb.facts(), kolibri1.backend.max_window, kolibri1.backend.max_window - 1); + defer cfg.deinit(gpa); + var clock: lanes.fake.FixedClock = .{}; + var engine = lanes.Engine.init(gpa, &cfg, kb.backend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(lanes.Stream, cases.len); + defer gpa.free(streams); + const props = try gpa.alloc(lanes.SuffixLookup, cases.len); + defer gpa.free(props); + for (cases, streams, props) |c, *s, *p| { + p.* = try lanes.SuffixLookup.init(gpa, .{ .min_match = 4 }); + s.* = try lanes.Stream.init(gpa, .{ .id = "k", .prompt = c.prompt, .max_new = count, .sampling = c.sampling, .drafts = c.drafts, .proposer = p.proposer() }); + } + defer for (streams, props) |*s, *p| { + s.deinit(gpa); + p.deinit(); + }; + for (streams) |*s| try engine.addStream(s); + const t0 = check.now(gpu.io); + while (engine.activeCount() > 0) try engine.step(); + const secs = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + var run: Run = .{ .out = try gpa.alloc([]u32, cases.len), .rate = 0, .rounds = 0, .accepted = 0 }; + var n: usize = 0; + for (run.out, streams) |*o, *s| { + o.* = try gpa.dupe(u32, s.emitted()); + n += o.len; + run.rounds += s.rounds; + run.accepted += s.accepted; + } + run.rate = @as(f64, @floatFromInt(n)) / secs; + return run; +} + +fn freeRun(gpa: std.mem.Allocator, r: Run) void { + for (r.out) |o| gpa.free(o); + gpa.free(r.out); +} + +/// Greedy tokens by one eager decode row at a time (the plain reference every other path must equal). +fn eager(r: *Rig, prompt: []const u32, count: usize) ![]u32 { + const out = try r.gpu.gpa.alloc(u32, count); + const first = try r.prompt(0, prompt); + out[0] = argmax(first); + r.gpu.gpa.free(first); + for (1..count) |i| { + const row = try r.decode(0, prompt.len + i - 1, out[i - 1]); + defer r.gpu.gpa.free(row); + out[i] = argmax(row); + } + return out; +} + +/// Each prompt (JSON id lists) decoded greedy for `count` tokens: ids and the top-2 margin of each step, as JSON lines. +pub fn greedy(gpu: Gpu, model: []const u8, prompts: []const u8, count_text: []const u8, out_path: []const u8) !void { + const gpa = gpu.gpa; + const count = try std.fmt.parseInt(usize, count_text, 10); + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, prompts, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice([][]u32, gpa, text, .{}); + defer parsed.deinit(); + var rig: Rig = undefined; + try rig.init(gpu, model, null, 8192, 1, fwd.prompt_rows); + defer rig.deinit(); + var out: std.Io.Writer.Allocating = .init(gpa); + defer out.deinit(); + for (parsed.value) |p| { + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpa); + try out.writer.writeAll("{\"steps\": ["); + var row = try rig.prompt(0, p); + for (0..count) |i| { + const best = argmax(row); + var second: f32 = -std.math.inf(f32); + for (row, 0..) |v, j| if (j != best and v > second) { + second = v; + }; + try out.writer.print("{s}[{d}, {d}]", .{ if (i == 0) "" else ", ", best, row[best] - second }); + try ids.append(gpa, best); + gpa.free(row); + row = if (i + 1 < count) try rig.decode(0, p.len + i, best) else &.{}; + } + try out.writer.print("], \"hash\": \"{x}\"}}\n", .{hash(ids.items)}); + } + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = out_path, .data = out.written() }); + std.debug.print("RESULT kolibri1 greedy: {d} prompts, {d} tokens each, to {s}\n", .{ parsed.value.len, count, out_path }); +} + +/// The full model: eager rows == captured steps == lane rounds, drafted == plain, shared rounds == solo, cuts, states. +pub fn full(gpu: Gpu, model: []const u8, ids_file: []const u8, count_text: []const u8) !void { + const gpa = gpu.gpa; + const count = try std.fmt.parseInt(u32, count_text, 10); + const ids = try readIds(gpu, ids_file); + defer gpa.free(ids); + var rig: Rig = undefined; + try rig.init(gpu, model, null, 16384, 4, fwd.prompt_rows); + defer rig.deinit(); + var kb = try kolibri1.backend.Kolibri.init(gpa, gpu.d, &rig.m, &rig.k, rig.s); + defer kb.deinit(); + var bad: usize = 0; + // a short prompt, and a long one whose second half repeats its first (copy drafts land there) + const short = ids[0..200]; + const long = try gpa.alloc(u32, 2 * 1500); + defer gpa.free(long); + @memcpy(long[0..1500], ids[0..1500]); + @memcpy(long[1500..], ids[0..1500]); + for ([_][]const u32{ short, long }, [_][]const u8{ "short", "long" }) |p, name| { + const plain = try eager(&rig, p, count); + defer gpa.free(plain); + const lone = try lanesRun(gpu, &kb, &.{.{ .prompt = p, .drafts = false }}, count); + defer freeRun(gpa, lone); + const drafted = try lanesRun(gpu, &kb, &.{.{ .prompt = p, .drafts = true }}, count); + defer freeRun(gpa, drafted); + const ok = std.mem.eql(u32, plain, lone.out[0]) and std.mem.eql(u32, plain, drafted.out[0]); + std.debug.print("RESULT kolibri1 {s}: eager rows {x}, steps {x} ({d:.2} tok/s), drafted {x} ({d:.2} tok/s, {d} rounds, {d} drafts landed)\n", .{ name, hash(plain), hash(lone.out[0]), lone.rate, hash(drafted.out[0]), drafted.rate, drafted.rounds, drafted.accepted }); + if (!ok) bad += 1; + } + // sampled: drafted == plain at a seed; then four streams in shared rounds against each alone + const sampled: lanes.Sampling = .{ .seed = 7, .temperature = 0.8, .top_k = 20, .top_p = 0.95 }; + const sp = try lanesRun(gpu, &kb, &.{.{ .prompt = long, .sampling = sampled, .drafts = false }}, count); + defer freeRun(gpa, sp); + const sd = try lanesRun(gpu, &kb, &.{.{ .prompt = long, .sampling = sampled, .drafts = true }}, count); + defer freeRun(gpa, sd); + std.debug.print("RESULT kolibri1 sampled: plain {x}, drafted {x}\n", .{ hash(sp.out[0]), hash(sd.out[0]) }); + if (!std.mem.eql(u32, sp.out[0], sd.out[0])) bad += 1; + const four = [_]Case{ .{ .prompt = long, .drafts = true }, .{ .prompt = short, .sampling = sampled, .drafts = true }, .{ .prompt = ids[300..700], .drafts = true }, .{ .prompt = long[0 .. long.len - 7], .sampling = sampled, .drafts = true } }; + const together = try lanesRun(gpu, &kb, &four, count); + defer freeRun(gpa, together); + for (four, 0..) |c, i| { + const alone = try lanesRun(gpu, &kb, &.{c}, count); + defer freeRun(gpa, alone); + const eq = std.mem.eql(u32, together.out[i], alone.out[0]); + std.debug.print("RESULT kolibri1 shared rounds: stream {d} {x} together, {x} alone\n", .{ i, hash(together.out[i]), hash(alone.out[0]) }); + if (!eq) bad += 1; + } + std.debug.print("RESULT kolibri1 shared rounds: {d:.2} tok/s for four streams\n", .{together.rate}); + bad += try cuts(&rig, ids); + bad += try states(&rig, &kb, ids, count); + try check.expect(bad == 0, "kolibri1 full: {d} paths differ from the plain rows", .{bad}); + check.pass("kolibri1 full: steps, drafted, sampled, shared rounds, prompt cuts and kept states equal the plain rows ({d} tokens each)", .{count}); +} + +/// A prompt's last logits cut on the 2,048 grid, in 1,000-row and in 1,536-row chunks: the same bytes? +fn cuts(r: *Rig, ids: []const u32) !usize { + const gpa = r.gpu.gpa; + const p = ids[0..5000]; + const grid = try r.prompt(0, p); + defer gpa.free(grid); + var bad: usize = 0; + for ([_]usize{ 1000, 1536 }) |chunk| { + var at: usize = 0; + while (at < p.len) : (at += chunk) { + const end = @min(p.len, at + chunk); + try r.m.forward(r.ops(), &.{.{ .slot = 1, .p0 = at, .tokens = p[at..end] }}, true, null, r.logits.ptr); + } + const got = try r.read(1); + defer gpa.free(got); + const eq = std.mem.eql(u8, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(grid)); + var most: f32 = 0; + for (got, grid) |g, w| most = @max(most, @abs(g - w)); + std.debug.print("RESULT kolibri1 cuts: {d}-row chunks against the 2,048 grid: {s} (largest difference {e:.3})\n", .{ chunk, if (eq) "same bytes" else "differ", most }); + if (!eq) bad += 1; + } + return bad; +} + +/// States kept at 2,048, 3,001 and 4,096 by one stream, resumed by another: greedy replies against fresh. +fn states(r: *Rig, kb: *kolibri1.backend.Kolibri, ids: []const u32, count: u32) !usize { + const gpa = r.gpu.gpa; + const p = ids[0..5000]; + const fresh = try lanesRun(r.gpu, kb, &.{.{ .prompt = p, .drafts = true }}, count); + defer freeRun(gpa, fresh); + var bad: usize = 0; + for ([_]u32{ 2048, 3001, 4096 }) |at| { + // keep the state through the backend's own path: a pass marked at `at`, its hook saving it + const Keep = struct { + kb: *kolibri1.backend.Kolibri, + saved: ?*kolibri1.snap.Saved = null, + fn hook(ptr: *anyopaque, s: *lanes.Stream, mark: u32) void { + const k: *@This() = @ptrCast(@alignCast(ptr)); + k.saved = k.kb.saveState(s, mark) catch null; + } + }; + var keep: Keep = .{ .kb = kb }; + var marks = [_]u32{at}; + var s = try lanes.Stream.init(gpa, .{ .id = "keep", .prompt = p[0 .. at + 100], .max_new = 1, .drafts = false, .reuse = .{ .marks = &marks, .hook = .{ .ptr = &keep, .at = Keep.hook } } }); + defer s.deinit(gpa); + const be = kb.backend(); + try be.prefill(&s); + be.release(&s); + const saved = keep.saved orelse return error.NothingKept; + defer kb.dropState(saved); + // a second stream resumes from it, prefilling the rest of a longer prompt + var lookup = try lanes.SuffixLookup.init(gpa, .{ .min_match = 4 }); + defer lookup.deinit(); + var t = try lanes.Stream.init(gpa, .{ .id = "resume", .prompt = p, .max_new = count, .drafts = true, .proposer = lookup.proposer(), .reuse = .{ .saved = saved, .at = at } }); + defer t.deinit(gpa); + var cfg = try lanes.Config.init(gpa, kb.facts(), kolibri1.backend.max_window, kolibri1.backend.max_window - 1); + defer cfg.deinit(gpa); + var clock: lanes.fake.FixedClock = .{}; + var engine = lanes.Engine.init(gpa, &cfg, be, clock.clock()); + defer engine.deinit(); + try engine.addStream(&t); + while (engine.activeCount() > 0) try engine.step(); + const eq = std.mem.eql(u32, t.emitted(), fresh.out[0]) and t.cached == at and !t.reuse_failed; + std.debug.print("RESULT kolibri1 kept state at {d}: {d} cached, reply {x}, fresh {x}\n", .{ at, t.cached, hash(t.emitted()), hash(fresh.out[0]) }); + if (!eq) bad += 1; + } + return bad; +} + +/// Teacher-forced fp32 logits of the first 2,048 prompt rows to `out_path`; with `prompts`, each one's last row too. +pub fn logits(gpu: Gpu, model: []const u8, ids_file: []const u8, out_path: []const u8, prompts: ?[]const u8) !void { + const gpa = gpu.gpa; + const ids = try readIds(gpu, ids_file); + defer gpa.free(ids); + const n = @min(ids.len, fwd.prompt_rows); + var rig: Rig = undefined; + try rig.init(gpu, model, null, 4096, 1, fwd.prompt_rows); + defer rig.deinit(); + const rows = try gpa.alloc(usize, n); + defer gpa.free(rows); + for (rows, 0..) |*x, i| x.* = i; + try rig.m.forward(rig.ops(), &.{.{ .slot = 0, .p0 = 0, .tokens = ids[0..n] }}, true, rows, rig.logits.ptr); + const host = try rig.read(n); + defer gpa.free(host); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = out_path, .data = std.mem.sliceAsBytes(host) }); + std.debug.print("RESULT kolibri1 logits: {d} rows x {d} to {s}\n", .{ n, rig.w.config.vocab, out_path }); + const file = prompts orelse return; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, file, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice([][]u32, gpa, text, .{}); + defer parsed.deinit(); + var last: std.ArrayList(f32) = .empty; + defer last.deinit(gpa); + for (parsed.value) |p| { + const row = try rig.prompt(0, p); + defer gpa.free(row); + try last.appendSlice(gpa, row); + } + const path = try std.fmt.allocPrint(gpa, "{s}.last", .{out_path}); + defer gpa.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = std.mem.sliceAsBytes(last.items) }); + std.debug.print("RESULT kolibri1 logits: the last rows of {d} prompts to {s}\n", .{ parsed.value.len, path }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2892d646ba..e8d36c149a 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,8 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const kolibri_tests = @import("kolibri_tests.zig"); +const kolibri_layer_tests = @import("kolibri_layer_tests.zig"); const tree_tests = @import("tree_tests.zig"); const block_tests = @import("block_tests.zig"); const window_tests = @import("window_tests.zig"); @@ -110,6 +112,11 @@ const usage = \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ kolibri-reference MODEL IDS Kolibri-1's layer 0 and head against float64 from the checkpoint (prompt and decode rows) + \\ kolibri-layers MODEL IDS Kolibri-1's window edge (513 keys) and its fifth layer's full reach, on the checkpoint + \\ kolibri-full MODEL IDS N Kolibri-1: eager rows, steps, drafted, sampled, shared rounds, cuts and kept states alike + \\ kolibri-greedy MODEL PROMPTS N OUT Kolibri-1 greedy from each prompt (JSON id lists): ids and top-2 margins + \\ kolibri-logits MODEL IDS OUT [PROMPTS] Kolibri-1's teacher-forced fp32 logits of 2,048 prompt rows; each prompt's last \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ ; @@ -216,6 +223,11 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "kolibri-reference")) return kolibri_layer_tests.reference(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "kolibri-layers")) return kolibri_layer_tests.layers(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "kolibri-full")) return kolibri_tests.full(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); + if (std.mem.eql(u8, cmd, "kolibri-greedy")) return kolibri_tests.greedy(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); + if (std.mem.eql(u8, cmd, "kolibri-logits")) return kolibri_tests.logits(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), if (rest.len > 3) rest[3] else null); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; From 60d5bdd8fe65f8635e685dbda7808d3d42734399 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 10:07:33 +0200 Subject: [PATCH 29/41] tests: kolibri1's float64 check bounds a row at 0.3 % of the logits' scale (a dropped routing bias or swapped expert scales pass 0.7 %) --- zig/tests/cuda/kolibri_layer_tests.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/tests/cuda/kolibri_layer_tests.zig b/zig/tests/cuda/kolibri_layer_tests.zig index 9e5e6fd600..9799153054 100644 --- a/zig/tests/cuda/kolibri_layer_tests.zig +++ b/zig/tests/cuda/kolibri_layer_tests.zig @@ -125,10 +125,10 @@ pub fn reference(gpu: Gpu, model: []const u8, ids_file: []const u8) !void { const top = argmax(got) == std.mem.indexOfMax(f64, want); std.debug.print("RESULT kolibri1 reference case {d} {s}: max |logit - f64| {e:.3} of {d:.3}, top-1 {s}\n", .{ case, kind, most, scale, if (top) "equal" else "DIFFERS" }); worst = @max(worst, most / scale); - if (!top or most > 0.02 * scale) bad += 1; + if (!top or most > 0.003 * scale) bad += 1; // clean rows stay near 0.1 %, a dropped bias or swapped scale passes 0.7 % } } - try check.expect(bad == 0, "kolibri1 reference: {d} rows past 2% of the logits' scale or a different top-1", .{bad}); + try check.expect(bad == 0, "kolibri1 reference: {d} rows past 0.3% of the logits' scale or a different top-1", .{bad}); check.pass("kolibri1 reference: layer 0 and the head within {d:.4} of the logits' scale of float64, top-1 equal", .{worst}); } From 577555190427e7512a98790aa36686ecd5587cb4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 10:33:39 +0200 Subject: [PATCH 30/41] docs: Kolibri-1 on native CUDA (research checkpoint): the pinned revision, the shared operators it runs on, checks, speed --- docs/recipes/README.md | 1 + docs/recipes/kolibri-1.md | 85 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 docs/recipes/kolibri-1.md diff --git a/docs/recipes/README.md b/docs/recipes/README.md index a635ce9602..8522cc3581 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -17,6 +17,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) | | DeepSeek-V4-Flash | [MLX on a 256 GB Mac, DSpark and MTP drafts](deepseek-v4-flash.md) | | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) | +| Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) | ## Capability floor diff --git a/docs/recipes/kolibri-1.md b/docs/recipes/kolibri-1.md new file mode 100644 index 0000000000..dd10bad56c --- /dev/null +++ b/docs/recipes/kolibri-1.md @@ -0,0 +1,85 @@ +# Kolibri-1 + +This page describes the native engine. `kolibri1` serves Aleph Alpha's Kolibri-1, a research checkpoint, on one +NVIDIA GPU through the lane host. It is a draft family: plain decoding and copy drafts from the prompt's own text are +in; the EAGLE-3 drafter is the next step and is not loaded here. + +## Checkpoint + +Tested revision: `Aleph-Alpha/Kolibri-1` at `e52eb4627d11516b0c01de49210ab5a4e4061444` (32 shards, 74 GB). The +weights are block FP8: e4m3 codes in 128 x 128 blocks with fp32 `weight_scale_inv`, the routers, norms, embeddings and +head in bf16. Publisher sha256 of `model.safetensors.index.json`: +`c6769a34820c9057b404b2a5cae1ac24fd7b4a9ca08b84eb0a06e7b5bad3005e`. + +```bash +tensorfold-native serve /path/to/Kolibri-1 --name kolibri-1 --context 32768 --parallel 8 +``` + +The model has 50 layers, 2,560 wide, 48 query heads over 4 key heads of 128. Every fifth layer (4, 9, ..., 49) attends +to every key without positions; the other 40 attend to a 513-key window with RoPE. Each layer routes by sigmoid with an +expert bias to 6 of 384 experts, weighted by the unnormalised sigmoid, plus one ungated shared expert. + +## What the family adds + +It brings no kernels of its own. Every launch is a shared CUDA operator: + +- `cuda/fp8.zig` and `qlinear` for the attention projections (block FP8, `fp8g` lane tiles); +- `cuda/experts.zig` for the routed and shared experts (`fp8g`, the shared expert stacked last at weight 1); +- `qk_rows` (QK norm, RoPE, the cache write), `norms`, `embed_rows`, `moe_glue` (sigmoid top-k with bias), + `dense_rows` (router and head); +- `window_attention` for prompt rows of every layer (the full layers with no window) and for sliding decode rows; +- `tree_attention` for the full layers' decode rows over committed keys and each window's chain. + +The family supplies its config check, the tensor map (116,303 tensors, none left over), the packing into those +layouts, the layer graph and the lane backend. One-row steps replay a CUDA graph a stream; wider windows and shared +rounds run eager. Both give a row the same bits. + +## Serving + +- Prompts run in passes of up to 2,048 rows, one a step, so other streams' rounds run between them. A row's bits do not + depend on its chunk: 1,000- and 1,536-row cuts give the 2,048-row cut's logits byte for byte. +- Concurrent streams share rounds of up to 64 rows. Each stream's caches are allocated at load, one per `--parallel` + stream: at a 32k window 0.83 GiB a stream, at 64k 1.45 GiB. +- The prompt cache keeps a stream's state at any row: full layers' keys up to it and each sliding layer's last 512. +- Drafts: the lane core's copy drafts verify in windows of up to 16 rows. There is no draft head yet, so a request + for head drafts and `--drafter` are refused by name. +- Refused at load: other quantizations, `norm_topk_prob`, other layer kinds, a config with a vision tower, shapes the + shared kernels do not take, `--vision`, `--precision` and a captured kernel set (`TENSORFOLD_CUDA_KERNELS`). + +## Checks + +```bash +zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/bin/tf-cuda-test kolibri-reference MODEL IDS # layer 0 and the head against float64 from the checkpoint +zig-out/bin/tf-cuda-test kolibri-layers MODEL IDS # the 513-key window edge, layer 4's full reach +zig-out/bin/tf-cuda-test kolibri-full MODEL IDS 128 # steps, drafted, sampled, shared rounds, cuts, kept states +zig-out/bin/tf-cuda-test kolibri-greedy MODEL PROMPTS N OUT # greedy ids and top-2 margins a step +zig-out/bin/tf-cuda-test kolibri-logits MODEL IDS OUT # teacher-forced logits of 2,048 rows, for quality +``` + +`IDS` is a file of real-text token ids (5,000 or more for `kolibri-full`, 3,000 for `kolibri-layers`). Random ids give +flat logits and say nothing about quality. + +The native rows are not byte-equal to the Python engine's: its glue was Triton, these are the shared CUDA operators, +and their fp32 roundings differ by an ulp here and there. Kolibri-1's top-6-of-384 routing turns such differences into +different experts on near-tied rows, so greedy replies part where two tokens are nearly tied. On 2,048 rows of real +text, teacher-forced: next-token top-1 agrees on 92.3% of rows, KL 0.06; NLL 3.574 here against 3.578 for the +Python engine, top-1 35.96% against 35.47%, mean entropy 2.354 against 2.358. + +## Speed + +GB10, greedy, served by `tensorfold-native` with copy drafts on: + +| | tok/s | +| --- | ---: | +| One stream, code / chat (`bench_openai.py`, 256 tokens) | 51.9 / 49.8 | +| 4 streams, aggregate | 124.0 / 126.0 | +| 8 streams, aggregate | 175.5 / 183.1 | + +Cold prompts (`bench_cache.py --require-cold`, `--prompt-cache-gib 0`): 2,048 tokens 1.00 s, 8,192 tokens 5.04 s, +32,768 tokens 21.5 s to the first token. + +## Next + +- The EAGLE-3 drafter through the lane core's drafter interface. +- Prompt passes of 8,192 rows, as the Python engine ran them: each pass reads every expert once. From 7d599181b328dcddadfd553e89aee666404c3bf0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 10:53:24 +0200 Subject: [PATCH 31/41] cuda/window_attention: the tensor-core decode's per-head max needs no flag (a head with no key keeps m, exp2 of +0 is exactly 1), the old kernel's speed --- zig/kernels/cuda/window_attention.cu | 8 ++++---- zig/src/core/contract/registry.zig | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index e8619c872d..3c6da22a58 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -469,10 +469,10 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); } - const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; // each head's own: a NaN head moves only its row - const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; - const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; - const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + // each head's own max (no key for it: m stays, exp2(+0) is exactly 1), so a NaN head moves only its row + const float n0 = fmaxf(m0, tm0), n1 = fmaxf(m1, tm1); + const float a0 = m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0)); + const float a1 = m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1)); float c0 = 0.0f, c1 = 0.0f; #pragma unroll for (int n = 0; n < 4; ++n) { diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index ce90a057ea..73f84a45a4 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -48,7 +48,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "window attention tf_window_part", .symbol = "tf_window_part", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 51, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 } } }, .{ .launch = .{ .name = "window attention tf_window_merge", .symbol = "tf_window_merge", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "window attention tf_window_decode_tc", .symbol = "tf_window_decode_tc", .image = "window_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 166, .static = 0, .bound = 128 }, .{ .regs = 166, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "window attention tf_window_decode_tc", .symbol = "tf_window_decode_tc", .image = "window_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 4, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{4}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 8, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{8}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "moe tf_moe_topk_sigmoid", .symbol = "tf_moe_topk_sigmoid", .image = "moe_glue", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 } } }, From 7633646ab1cf11c17bbdaa12322eab05141911f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:06:23 +0200 Subject: [PATCH 32/41] kolibri1: prompt passes of 8,192 rows, as the Python engine ran them (sliding rings hold a pass and the window, 8,768 keys); cuts of 5,000 and 20,000 rows check against them --- docs/recipes/kolibri-1.md | 15 +++++----- zig/src/families/kolibri1/forward.zig | 4 +-- zig/tests/cuda/kolibri_tests.zig | 43 +++++++++++++++------------ 3 files changed, 34 insertions(+), 28 deletions(-) diff --git a/docs/recipes/kolibri-1.md b/docs/recipes/kolibri-1.md index dd10bad56c..9601290198 100644 --- a/docs/recipes/kolibri-1.md +++ b/docs/recipes/kolibri-1.md @@ -36,10 +36,12 @@ rounds run eager. Both give a row the same bits. ## Serving -- Prompts run in passes of up to 2,048 rows, one a step, so other streams' rounds run between them. A row's bits do not - depend on its chunk: 1,000- and 1,536-row cuts give the 2,048-row cut's logits byte for byte. +- Prompts run in passes of up to 8,192 rows, as the Python engine ran them, one a step, so other streams' rounds run + between them (a full pass holds them for about 3.7 s). A row's bits do not depend on its chunk: 1,000-, 1,536- and + 2,048-row cuts of a 5,000- and a 20,000-token prompt give the 8,192-row passes' logits byte for byte. - Concurrent streams share rounds of up to 64 rows. Each stream's caches are allocated at load, one per `--parallel` - stream: at a 32k window 0.83 GiB a stream, at 64k 1.45 GiB. + stream: at a 32k window 1.29 GiB a stream, at 64k 1.92 GiB. Each sliding layer's ring holds a pass and its window + (8,768 keys). - The prompt cache keeps a stream's state at any row: full layers' keys up to it and each sliding layer's last 512. - Drafts: the lane core's copy drafts verify in windows of up to 16 rows. There is no draft head yet, so a request for head drafts and `--drafter` are refused by name. @@ -57,7 +59,7 @@ zig-out/bin/tf-cuda-test kolibri-greedy MODEL PROMPTS N OUT # greedy ids and t zig-out/bin/tf-cuda-test kolibri-logits MODEL IDS OUT # teacher-forced logits of 2,048 rows, for quality ``` -`IDS` is a file of real-text token ids (5,000 or more for `kolibri-full`, 3,000 for `kolibri-layers`). Random ids give +`IDS` is a file of real-text token ids (20,000 or more for `kolibri-full`, 3,000 for `kolibri-layers`). Random ids give flat logits and say nothing about quality. The native rows are not byte-equal to the Python engine's: its glue was Triton, these are the shared CUDA operators, @@ -76,10 +78,9 @@ GB10, greedy, served by `tensorfold-native` with copy drafts on: | 4 streams, aggregate | 124.0 / 126.0 | | 8 streams, aggregate | 175.5 / 183.1 | -Cold prompts (`bench_cache.py --require-cold`, `--prompt-cache-gib 0`): 2,048 tokens 1.00 s, 8,192 tokens 5.04 s, -32,768 tokens 21.5 s to the first token. +Cold prompts (`bench_cache.py --require-cold --repeats 3`, `--prompt-cache-gib 0`, median of the cold requests): +2,048 tokens 0.97 s, 8,192 tokens 3.72 s, 32,768 tokens 16.8 s, 64,000 tokens 37.8 s to the first token. ## Next - The EAGLE-3 drafter through the lane core's drafter interface. -- Prompt passes of 8,192 rows, as the Python engine ran them: each pass reads every expert once. diff --git a/zig/src/families/kolibri1/forward.zig b/zig/src/families/kolibri1/forward.zig index ad842336ed..f555a2339c 100644 --- a/zig/src/families/kolibri1/forward.zig +++ b/zig/src/families/kolibri1/forward.zig @@ -12,7 +12,7 @@ const experts = cuda.experts; const TreeAttention = cuda.tree_attention.TreeAttention; /// Prompt rows a forward takes: the chunk grid, from position 0. -pub const prompt_rows = 2048; +pub const prompt_rows = 8192; /// Decode rows a forward takes over every stream's window. pub const round_rows = 64; const key_block = 64; // the window kernels' key tile @@ -377,6 +377,6 @@ pub const Model = struct { }; test "the ring holds a prompt chunk and a window, in whole key tiles" { - try std.testing.expectEqual(@as(usize, 2624), ringSize(prompt_rows, 513)); + try std.testing.expectEqual(@as(usize, 8768), ringSize(prompt_rows, 513)); try std.testing.expect(ringSize(prompt_rows, 513) - 513 >= prompt_rows); } diff --git a/zig/tests/cuda/kolibri_tests.zig b/zig/tests/cuda/kolibri_tests.zig index 8817358a46..f2c7c2d694 100644 --- a/zig/tests/cuda/kolibri_tests.zig +++ b/zig/tests/cuda/kolibri_tests.zig @@ -10,6 +10,8 @@ const Gpu = check.Gpu; const fwd = kolibri1.forward; const Rig = layer_tests.Rig; const readIds = layer_tests.readIds; +/// Rows `logits` writes: 2,048, as the quality runs read them. +const quality_rows = 2048; fn hash(tokens: []const u32) u64 { return std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(tokens)); @@ -123,7 +125,7 @@ pub fn full(gpu: Gpu, model: []const u8, ids_file: []const u8, count_text: []con const ids = try readIds(gpu, ids_file); defer gpa.free(ids); var rig: Rig = undefined; - try rig.init(gpu, model, null, 16384, 4, fwd.prompt_rows); + try rig.init(gpu, model, null, 20480, 4, fwd.prompt_rows); defer rig.deinit(); var kb = try kolibri1.backend.Kolibri.init(gpa, gpu.d, &rig.m, &rig.k, rig.s); defer kb.deinit(); @@ -170,26 +172,29 @@ pub fn full(gpu: Gpu, model: []const u8, ids_file: []const u8, count_text: []con check.pass("kolibri1 full: steps, drafted, sampled, shared rounds, prompt cuts and kept states equal the plain rows ({d} tokens each)", .{count}); } -/// A prompt's last logits cut on the 2,048 grid, in 1,000-row and in 1,536-row chunks: the same bytes? +/// 5,000 and 20,000 rows' last logits on the pass grid against 1,000-, 1,536- and 2,048-row chunks: same bytes? fn cuts(r: *Rig, ids: []const u32) !usize { const gpa = r.gpu.gpa; - const p = ids[0..5000]; - const grid = try r.prompt(0, p); - defer gpa.free(grid); + if (ids.len < 20000) return error.TooFewIds; var bad: usize = 0; - for ([_]usize{ 1000, 1536 }) |chunk| { - var at: usize = 0; - while (at < p.len) : (at += chunk) { - const end = @min(p.len, at + chunk); - try r.m.forward(r.ops(), &.{.{ .slot = 1, .p0 = at, .tokens = p[at..end] }}, true, null, r.logits.ptr); + for ([_]usize{ 5000, 20000 }) |len| { + const p = ids[0..len]; + const grid = try r.prompt(0, p); + defer gpa.free(grid); + for ([_]usize{ 1000, 1536, 2048 }) |chunk| { + var at: usize = 0; + while (at < p.len) : (at += chunk) { + const end = @min(p.len, at + chunk); + try r.m.forward(r.ops(), &.{.{ .slot = 1, .p0 = at, .tokens = p[at..end] }}, true, null, r.logits.ptr); + } + const got = try r.read(1); + defer gpa.free(got); + const eq = std.mem.eql(u8, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(grid)); + var most: f32 = 0; + for (got, grid) |g, w| most = @max(most, @abs(g - w)); + std.debug.print("RESULT kolibri1 cuts: {d} rows in {d}-row chunks against {d}-row passes: {s} (largest difference {e:.3})\n", .{ len, chunk, fwd.prompt_rows, if (eq) "same bytes" else "differ", most }); + if (!eq) bad += 1; } - const got = try r.read(1); - defer gpa.free(got); - const eq = std.mem.eql(u8, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(grid)); - var most: f32 = 0; - for (got, grid) |g, w| most = @max(most, @abs(g - w)); - std.debug.print("RESULT kolibri1 cuts: {d}-row chunks against the 2,048 grid: {s} (largest difference {e:.3})\n", .{ chunk, if (eq) "same bytes" else "differ", most }); - if (!eq) bad += 1; } return bad; } @@ -239,12 +244,12 @@ fn states(r: *Rig, kb: *kolibri1.backend.Kolibri, ids: []const u32, count: u32) return bad; } -/// Teacher-forced fp32 logits of the first 2,048 prompt rows to `out_path`; with `prompts`, each one's last row too. +/// Teacher-forced fp32 logits of the first `quality_rows` rows to `out_path`; with `prompts`, each one's last row too. pub fn logits(gpu: Gpu, model: []const u8, ids_file: []const u8, out_path: []const u8, prompts: ?[]const u8) !void { const gpa = gpu.gpa; const ids = try readIds(gpu, ids_file); defer gpa.free(ids); - const n = @min(ids.len, fwd.prompt_rows); + const n = @min(ids.len, quality_rows); var rig: Rig = undefined; try rig.init(gpu, model, null, 4096, 1, fwd.prompt_rows); defer rig.deinit(); From 179648432b561e986a2ec83951360b49749e114a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:30:23 +0200 Subject: [PATCH 33/41] kolibri1: plan the streams inside the room the weights leave before allocating them; refuse a fixed --parallel past the memory or a round's 64 rows; admission adds no lane to a family that allocated its streams at open Model.openBytes counts init's allocations by the same path without allocating; the open reads the server's room, allocates as many streams as asked when they fit, the most that fit under auto, and refuses a fixed count that does not fit or exceeds 64 before any cache is allocated. cuda_memory.growRoom gives a family whose Options carry `lanes` no room to grow into, so the host admits only the streams it holds. --- zig/src/families/kolibri1/backend.zig | 5 ++ zig/src/families/kolibri1/forward.zig | 43 ++++++++++++-- zig/src/families/kolibri1/native.zig | 82 ++++++++++++++++++++++++--- zig/src/native/cuda.zig | 10 +++- zig/src/native/cuda_memory.zig | 17 ++++++ 5 files changed, 142 insertions(+), 15 deletions(-) diff --git a/zig/src/families/kolibri1/backend.zig b/zig/src/families/kolibri1/backend.zig index 1f5154330f..029e25536b 100644 --- a/zig/src/families/kolibri1/backend.zig +++ b/zig/src/families/kolibri1/backend.zig @@ -46,6 +46,11 @@ pub const Kolibri = struct { steps: []?fwd.Model.Step, // a slot's one-row step, captured at its first use (null: not yet) graphed: bool, // TENSORFOLD_KOLIBRI_GRAPHS=0 runs every step eager + /// The device bytes `init` allocates: a round's logits, the token ring, and each slot's last row and rule. + pub fn openBytes(vocab: usize, slots: usize) usize { + return round_rows * vocab * 6 + (ring + round_rows) * 4 + round_rows * 16 + slots * (vocab * 4 + @sizeOf(cuda.sample.Rule)); + } + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *fwd.Model, k: *const kern.Kernels, s: cuda.Stream) !Kolibri { const v = m.c.vocab; const free = try gpa.alloc(bool, m.slots); diff --git a/zig/src/families/kolibri1/forward.zig b/zig/src/families/kolibri1/forward.zig index f555a2339c..1e89b606e3 100644 --- a/zig/src/families/kolibri1/forward.zig +++ b/zig/src/families/kolibri1/forward.zig @@ -43,10 +43,15 @@ pub const Model = struct { origin: u64, // the bf16 base tree attention's cache offsets count from s: Scratch, buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated const Scratch = struct { ids: u64, pos: u64, slot: u64, res: u64, x: u64, qkv: u64, q: u64, kn: u64, vn: u64, att: u64, o: u64, part: u64, logits_e: u64, pick: u64, wts: u64, plan: grouped.Plan, act: u64, y: u64, mo: u64, flat: u64, paths: u64, depths: u64, offs: u64, sel: u64, win: u64, tree: u64 }; fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += @max(bytes, 256); + return 0; + } const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); errdefer { var x = b; @@ -58,6 +63,7 @@ pub const Model = struct { fn zeros(m: *Model, s: cuda.Stream, bytes: usize) !u64 { const p = try m.alloc(bytes); + if (m.planned != null) return p; try m.d.check(m.d.api.cuMemsetD8Async(p, 0, @max(bytes, 256), s.handle), "cuMemsetD8Async"); return p; } @@ -71,10 +77,34 @@ pub const Model = struct { /// Caches for `slots` streams of `context` positions, RoPE tables, and scratch for `max_rows` rows a forward. pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, max_rows: usize) !Model { + var m = blank(gpa, d, w, context, slots, max_rows); + errdefer m.deinit(); + try m.build(s, k); + return m; + } + + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const kern.Kernels, context: usize, slots: usize, max_rows: usize) !usize { + var m = blank(gpa, d, w, context, slots, max_rows); + defer m.deinit(); + m.planned = 0; + try m.build(s, k); + return m.planned.?; + } + + fn blank(gpa: std.mem.Allocator, d: *const cuda.Driver, w: *const wts.Weights, context: usize, slots: usize, max_rows: usize) Model { var c = w.config; c.layers = w.layers.len; // a model loaded with fewer layers (checks) runs those alone - var m: Model = .{ .gpa = gpa, .d = d, .w = w, .c = c, .context = context, .slots = slots, .ring = ringSize(prompt_rows, c.window), .max_rows = max_rows, .groups = TreeAttention.groups(context + tree.max_rows), .k = &.{}, .v = &.{}, .norms = &.{}, .cos = 0, .sin = 0, .origin = 0, .s = undefined }; - errdefer m.deinit(); + return .{ .gpa = gpa, .d = d, .w = w, .c = c, .context = context, .slots = slots, .ring = ringSize(prompt_rows, c.window), .max_rows = max_rows, .groups = TreeAttention.groups(context + tree.max_rows), .k = &.{}, .v = &.{}, .norms = &.{}, .cos = 0, .sin = 0, .origin = 0, .s = undefined }; + } + + fn build(m: *Model, s: cuda.Stream, k: *const kern.Kernels) !void { + const gpa = m.gpa; + const d = m.d; + const c = m.c; + const w = m.w; + const slots = m.slots; + const context = m.context; const dh = c.head_dim; m.k = try gpa.alloc(u64, c.layers); m.v = try gpa.alloc(u64, c.layers); @@ -85,6 +115,7 @@ pub const Model = struct { m.k[i] = try m.zeros(s, bytes); m.v[i] = try m.zeros(s, bytes); m.norms[i] = try m.alloc(2 * dh * 4); + if (m.planned != null) continue; const L = w.layers[i]; try d.check(d.api.cuMemcpyDtoDAsync_v2(m.norms[i], L.q_norm, dh * 4, s.handle), "cuMemcpyDtoDAsync"); try d.check(d.api.cuMemcpyDtoDAsync_v2(m.norms[i] + dh * 4, L.k_norm, dh * 4, s.handle), "cuMemcpyDtoDAsync"); @@ -92,8 +123,7 @@ pub const Model = struct { try m.rope(s); m.origin = try m.zeros(s, 256); try m.scratch(k); - try s.synchronize(); - return m; + if (m.planned == null) try s.synchronize(); } pub fn deinit(m: *Model) void { @@ -107,6 +137,11 @@ pub const Model = struct { /// RoPE tables: angle p / theta^(2j / D) in float64, cos and sin rounded to fp32, the halves repeated (neox). fn rope(m: *Model, s: cuda.Stream) !void { const dh = m.c.head_dim; + if (m.planned != null) { + m.cos = try m.alloc(m.context * dh * 4); + m.sin = try m.alloc(m.context * dh * 4); + return; + } const cos = try m.gpa.alloc(f32, m.context * dh); defer m.gpa.free(cos); const sin = try m.gpa.alloc(f32, m.context * dh); diff --git a/zig/src/families/kolibri1/native.zig b/zig/src/families/kolibri1/native.zig index c4756f11b6..a9c94d72ca 100644 --- a/zig/src/families/kolibri1/native.zig +++ b/zig/src/families/kolibri1/native.zig @@ -22,8 +22,10 @@ pub const drafter_refusal = "--drafter: Kolibri-1 on CUDA loads no drafter yet; /// `explain` words this family's open refusals too. pub const explains_open = true; +/// The server's reading of the room left after the weights; null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; /// No `vision` or `precision` field: the server refuses --vision and --precision for this family by name. -pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8 }; +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; @@ -43,11 +45,29 @@ pub const Loaded = struct { const Owned = struct { gpa: std.mem.Allocator, w: weights.Weights, k: kern.Kernels, s: cuda.Stream, m: fwd.Model, kb: be.Kolibri }; -/// Streams whose caches the open allocates: the server's --parallel, at most a round's rows. -fn streams(lanes_asked: u32) usize { - return std.math.clamp(@as(usize, lanes_asked), 1, be.round_rows); +/// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses. +fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + if (fixed and asked > be.round_rows) return error.ParallelPastRound; + var n = std.math.clamp(@as(usize, asked), 1, be.round_rows); + while (try plan.bytes(n) > room) : (n -= 1) { + if (fixed) return error.ParallelPastMemory; + if (n == 1) return error.NoStreamFits; + } + return n; } +/// The device bytes `n` streams take: their caches, the forward's scratch and the backend's buffers. +const Plan = struct { + own: *const Owned, + d: *const cuda.Driver, + context: usize, + + fn bytes(p: Plan, n: usize) !u64 { + const o = p.own; + return try fwd.Model.openBytes(o.gpa, p.d, o.s, &o.w, &o.k, p.context, n, be.prompt_rows) + be.Kolibri.openBytes(o.w.config.vocab, n); + } +}; + /// Positions a stream's caches hold: the window, whole 1,024s. fn positions(context: usize) usize { return (context + 1023) / 1024 * 1024; @@ -56,6 +76,7 @@ fn positions(context: usize) usize { /// The checkpoint, the shared kernels and every stream's caches. pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.CapturedKernels; + if (o.lanes_fixed and o.lanes > be.round_rows) return error.ParallelPastRound; // before the weights load const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; @@ -66,7 +87,16 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.s = try cuda.Stream.init(ctx.d, false); errdefer own.s.deinit(); const context = positions(o.context); - const n = streams(o.lanes); + // every stream's caches are allocated here and none later: plan them inside the room the weights leave + const plan: Plan = .{ .own = own, .d = ctx.d, .context = context }; + const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + if (e == error.ParallelPastMemory) { + const most = fit(room, o.lanes, false, plan) catch 0; + std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, gib(room), most, if (most == 1) "" else "s" }); + } + return e; + }; own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, n, be.prompt_rows); errdefer own.m.deinit(); own.kb = try be.Kolibri.init(gpa, ctx.d, &own.m, &own.k, own.s); @@ -82,6 +112,10 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ }; } +fn gib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + fn kolibri(ptr: *anyopaque) *be.Kolibri { return @ptrCast(@alignCast(ptr)); } @@ -112,6 +146,9 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { error.NoDraftHead => "Kolibri-1 on CUDA has no draft head yet: drafts come from the prompt's own text only", error.StructuresUnsupported => "structured output is not built for Kolibri-1 on CUDA yet", error.CapturedKernels => "Kolibri-1 runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: Kolibri-1 on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer Kolibri-1 streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.NoStreamFits => "the CUDA memory budget fits no Kolibri-1 stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", else => null, }; } @@ -132,11 +169,40 @@ test "the prompt cache's hooks compile against the store's vtable" { } test "every refusal of its own is worded, the rest left to the server" { - for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.CapturedKernels }) |e| { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.CapturedKernels, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits }) |e| { try std.testing.expect(explain(null, e) != null); } try std.testing.expect(explain(null, error.OutOfMemory) == null); try std.testing.expectEqual(@as(usize, 65536), positions(65000)); - try std.testing.expectEqual(@as(usize, 1), streams(0)); - try std.testing.expectEqual(@as(usize, be.round_rows), streams(1000)); +} + +/// A synthetic plan: `stream` bytes a stream plus `fixed` bytes of scratch. +const Fake = struct { + stream: u64, + fixed: u64, + fn bytes(f: Fake, n: usize) !u64 { + return f.fixed + n * f.stream; + } +}; + +test "the streams allocated at open fit the room the weights leave: auto takes what fits, a fixed count refuses" { + const g: u64 = 1 << 30; + const one: Fake = .{ .stream = g, .fixed = 0 }; + // 70 GiB of weights under a 72 GiB cap leave 2 GiB: eight 1 GiB streams would hold 78 GiB + try std.testing.expectEqual(@as(usize, 2), try fit(72 * g - 70 * g, 8, false, one)); + try std.testing.expectError(error.ParallelPastMemory, fit(72 * g - 70 * g, 8, true, one)); + // the reserve instead of a cap: 10 GiB free less an 8 GiB reserve + try std.testing.expectEqual(@as(usize, 2), try fit(10 * g - 8 * g, 8, false, one)); + // scratch counts: 2.5 GiB of room, 1 GiB of it scratch, holds one stream + try std.testing.expectEqual(@as(usize, 1), try fit(5 * g / 2, 8, false, Fake{ .stream = g, .fixed = g })); + try std.testing.expectError(error.NoStreamFits, fit(g - 1, 8, false, one)); + try std.testing.expectEqual(@as(usize, 8), try fit(8 * g, 8, true, one)); + try std.testing.expectEqual(@as(usize, 8), try fit(std.math.maxInt(u64), 8, false, one)); // no room hook: as asked +} + +test "a fixed --parallel past a round's rows refuses, whatever the room" { + const g: u64 = 1 << 30; + const one: Fake = .{ .stream = g, .fixed = 0 }; + try std.testing.expectError(error.ParallelPastRound, fit(1024 * g, be.round_rows + 1, true, one)); + try std.testing.expectEqual(@as(usize, be.round_rows), try fit(1024 * g, be.round_rows, true, one)); } diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 331046a9fa..df06e172dd 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -336,7 +336,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O opts.precision = std.meta.stringToEnum(@typeInfo(@TypeOf(opts.precision)).optional.child, @tagName(p)).?; } if (@hasField(F.Options, "room")) opts.room = .{ .ctx = &probe, .read = Probe.read }; - if (@hasField(F.Options, "lanes")) opts.lanes = o.lanes; // a family that allocates every stream's caches at open + if (@hasField(F.Options, "lanes")) { // a family that allocates every stream's caches at open + opts.lanes = o.lanes; + opts.lanes_fixed = o.lanes_fixed; + } if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; @@ -381,7 +384,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); }; const free = loaded.free_streams; - const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { + const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); + const streams = budget.admit(grow, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { error.NoStream => try std.fmt.allocPrint(a, no_stream, .{ window, @@ -397,7 +401,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O window, toGib(room), low, - free + room / loaded.stream_bytes, + free + grow / loaded.stream_bytes, }), }; return null; diff --git a/zig/src/native/cuda_memory.zig b/zig/src/native/cuda_memory.zig index f5da4bf761..c3c0c8af92 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -110,6 +110,11 @@ pub fn admit(room: u64, stream: u64, free: u32, asked: u32, fixed: bool) error{ return if (fixed) error.TooMany else fits; } +/// The room admission adds streams in: none for a family that allocated at open every stream it can serve. +pub fn growRoom(room: u64, grows: bool) u64 { + return if (grows) room else 0; +} + /// Kept prompt states' budget: --prompt-cache-gib (0: none), else 16 GiB, never past what the streams leave of `room`. pub fn cacheBytes(asked: ?f64, room: u64, streams: u64) u64 { const want: u64 = if (asked) |g| (if (g > 0) std.math.lossyCast(u64, g * gib) else 0) else 16 << 30; @@ -149,6 +154,18 @@ test "the memory plan: reserve, cap, and admission that fails closed" { try std.testing.expectError(error.TooMany, admit(2 * g, g, 1, 4, true)); } +test "a family that allocates every stream at open is admitted only the streams it holds" { + const g: u64 = 1 << 30; + // 64 slots allocated, room for one more: a fixed 65 refuses instead of adopting a lane with no slot + try std.testing.expectError(error.TooMany, admit(growRoom(g, false), g, 64, 65, true)); + // auto planned 3 that fit: room left over for another stream adds no lane + try std.testing.expectEqual(@as(u32, 3), try admit(growRoom(g, false), g, 3, 8, false)); + // streams paid for at open stay admitted when nothing is left + try std.testing.expectEqual(@as(u32, 8), try admit(growRoom(0, false), g, 8, 8, true)); + // a family that grows streams after its open keeps the room + try std.testing.expectEqual(@as(u32, 65), try admit(growRoom(g, true), g, 64, 65, true)); +} + test "the prompt cache: 16 GiB unless the flag sizes it, inside what the streams leave" { const g: u64 = 1 << 30; try std.testing.expectEqual(16 * g, cacheBytes(null, 60 * g, 8 * g)); From 0ad4f170c80ed435194f056a18c9d328b9c3b57b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 01:37:58 +0200 Subject: [PATCH 34/41] tests/cuda/contract_attn: the reference's constructor frees what it took when a later allocation fails, with a host test that fails each allocation in turn --- zig/build/cuda.zig | 3 +++ zig/tests/cuda/contract_attn.zig | 22 +++++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index f3d50c01de..45f0ca293b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -274,6 +274,9 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + // the contract's attention reference (tests/cuda/contract_attn.zig): its allocations under failure + const attn_ref = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_attn.zig"), .target = host, .optimize = .debug, .imports = &.{.{ .name = "core", .module = mods.core }} }); + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = attn_ref })).step); // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); diff --git a/zig/tests/cuda/contract_attn.zig b/zig/tests/cuda/contract_attn.zig index cec73374a3..afa284d55c 100644 --- a/zig/tests/cuda/contract_attn.zig +++ b/zig/tests/cuda/contract_attn.zig @@ -27,8 +27,19 @@ pub const Ref = struct { pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref { const keys = try a.alloc(usize, most); + errdefer a.free(keys); for (keys, 0..) |*k, i| k.* = i; - return .{ .a = a, .d = d, .head = .{ .q = try a.alloc(f64, d), .k = try a.alloc(f64, most * d), .v = try a.alloc(f64, most * d) }, .keys = keys, .p = try a.alloc(f64, most), .want = try a.alloc(f64, d), .tol = try a.alloc(f64, d) }; + const q = try a.alloc(f64, d); + errdefer a.free(q); + const k = try a.alloc(f64, most * d); + errdefer a.free(k); + const v = try a.alloc(f64, most * d); + errdefer a.free(v); + const p = try a.alloc(f64, most); + errdefer a.free(p); + const want = try a.alloc(f64, d); + errdefer a.free(want); + return .{ .a = a, .d = d, .head = .{ .q = q, .k = k, .v = v }, .keys = keys, .p = p, .want = want, .tol = try a.alloc(f64, d) }; } pub fn deinit(r: *Ref) void { @@ -122,3 +133,12 @@ pub fn mirrored(rng: std.Random, out: []f32) void { pub fn normal(rng: std.Random) u16 { return c.toBf16(rng.floatNorm(f32)); } + +fn initAndFree(a: std.mem.Allocator) !void { + var r = try Ref.init(a, 8, 5); + r.deinit(); +} + +test "Ref.init frees what it took when a later allocation fails" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, initAndFree, .{}); +} From 2ae3b6952ee24a608fa686fe5276e869fd1a5864 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:01:32 +0200 Subject: [PATCH 35/41] cuda/block_attention: the replay takes every NaN and infinite value out of the mma and adds it back at its key's float64 weight, so an infinite value gives NaN where float64 does --- zig/kernels/cuda/attn_odd.cuh | 63 ++++++++++++++++++++++---- zig/kernels/cuda/block_attention.cu | 25 ++++++---- zig/src/core/contract/registry.zig | 2 +- zig/tests/cuda/contract_block_attn.zig | 30 ++++++++++++ 4 files changed, 103 insertions(+), 17 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 596ee3a95e..7501de7b7c 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -3,6 +3,7 @@ #include #include +#include #include namespace attn_odd { @@ -34,27 +35,73 @@ __device__ __forceinline__ void clean(uint4* vs, int threads) { } } -// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them. -template -__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0, - float l1) { +// The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. +__device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } + +// A query row's float64 score against a key row, attention_ref's order: sum q_e k_e over e, then times scale. +template +__device__ __forceinline__ double score64(const __nv_bfloat16* q, const __nv_bfloat16* k, double scale) { + double s = 0.0; +#pragma unroll 1 + for (int e = 0; e < D; ++e) + s = __dadd_rn(s, __dmul_rn((double)__bfloat162float(q[e]), (double)__bfloat162float(k[e]))); + return __dmul_rn(s, scale); +} + +// Rows g, g + 8: w(r, K row) is attention_ref's exp(score - max) > 0; each(r, f) walks r's K rows (max on first use). +template +struct Weights { + Q q; + Each each; + double scale; + int g; + double m[2] = {-INFINITY, -INFINITY}; + bool have[2] = {false, false}; + + __device__ __forceinline__ bool operator()(int r, const __nv_bfloat16* k) { + const int i = r == g ? 0 : 1; + if (!have[i]) { + double mx = -INFINITY; + each(r, [&](const __nv_bfloat16* kr) { mx = fmax(mx, score64(q(r), kr, scale)); }); + m[i] = mx; + have[i] = true; + } + return __dsub_rn(score64(q(r), k, scale), m[i]) >= least(); // NaN: no weight + } +}; + +// q(r): row r's query. +template +__device__ __forceinline__ Weights weights(Q q, Each each, float scale, int lane) { + return Weights{q, each, (double)scale, lane / 4}; +} + +// Keys k0..k1 clean() emptied (V vrow(k), K krow(k)): NaN and infinities to rows that see them, inf x 0 at weight 0. +template +__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, KRow krow, int k0, int k1, int lane, Vis vis, + float l0, float l1, W& w) { const int g = lane / 4, qd = lane % 4; for (int k = k0; k < k1; ++k) { const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN if (!v0 && !v1) continue; const uint32_t* v = reinterpret_cast(vrow(k)); + bool odd = false; +#pragma unroll + for (int i = 0; i < D / 8; ++i) odd |= finite2(v[4 * i + qd]) != v[4 * i + qd]; + if (!odd) continue; + const bool w0 = v0 && w(g, krow(k)), w1 = v1 && w(g + 8, krow(k)); #pragma unroll for (int i = 0; i < D / 8; ++i) { const uint32_t x = v[4 * i + qd]; if (finite2(x) == x) continue; const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u); if ((x & 0x7f80u) == 0x7f80u) { - if (v0) o[i][0] = __fadd_rn(o[i][0], lo); - if (v1) o[i][2] = __fadd_rn(o[i][2], lo); + if (v0) o[i][0] = __fadd_rn(o[i][0], w0 ? lo : __fmul_rn(0.0f, lo)); + if (v1) o[i][2] = __fadd_rn(o[i][2], w1 ? lo : __fmul_rn(0.0f, lo)); } if ((x & 0x7f800000u) == 0x7f800000u) { - if (v0) o[i][1] = __fadd_rn(o[i][1], hi); - if (v1) o[i][3] = __fadd_rn(o[i][3], hi); + if (v0) o[i][1] = __fadd_rn(o[i][1], w0 ? hi : __fmul_rn(0.0f, hi)); + if (v1) o[i][3] = __fadd_rn(o[i][3], w1 ? hi : __fmul_rn(0.0f, hi)); } } } diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index f5b1e2b18a..84cd949888 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -168,9 +168,8 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( Fold f; const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; - // the context's edge tiles, where some row cannot see some key: e0..e1 - const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT); - auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below) + const int e0 = lo / KT * KT; // the first context tile + auto context = [&](bool odd) { // odd: the tiles' NaN and infinite values leave the mma (added back below) for (int k0 = e0; k0 < s; k0 += KT) { for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { const int key = k0 + i / C; @@ -180,7 +179,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( } commit(); wait_all(); - if (odd && k0 < e1) attn_odd::clean(vs, blockDim.x); + if (odd) attn_odd::clean(vs, blockDim.x); __syncthreads(); f.tile(qa, ks, vs, scale, lane, k0, vis); __syncthreads(); @@ -202,17 +201,27 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( block_keys(); __syncthreads(); f.tile(qa, ks, vs, scale, lane, 0, own); - // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits) + // a NaN or infinite output: again, NaN and infinite values out of the mma, added back (finite values: same bits) if (__syncthreads_or(attn_odd::nonfinite_out(f.o, g8 < L, g8 + 8 < L))) { f = Fold(); context(true); - attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1); + const auto ck = [&](int k) { return kc + (int64_t)k * D; }; + const auto bk = [&](int c) { return kb + ((int64_t)g * rows + j * L + c) * D; }; + const auto qrow = [&](int r) { return q + ((int64_t)head * rows + j * L + r) * D; }; + const auto each = [&](int r, auto f) { // row r's keys: its window of the context, then the block's + for (int k = lo; k < s; ++k) + if (vis(r, k)) f(ck(k)); + for (int c = 0; c < L; ++c) + if (own(r, c)) f(bk(c)); + }; + auto w = attn_odd::weights(qrow, each, scale, lane); + attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, ck, e0, s, lane, vis, f.l0, f.l1, w); block_keys(); - if (causal) attn_odd::clean(vs, blockDim.x); // row r cannot see keys past r + attn_odd::clean(vs, blockDim.x); __syncthreads(); f.tile(qa, ks, vs, scale, lane, 0, own); const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; - if (causal) attn_odd::add_back(f.o, row, 0, L, lane, own, f.l0, f.l1); + attn_odd::add_back(f.o, row, bk, 0, L, lane, own, f.l0, f.l1, w); } if (g8 < L) #pragma unroll diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 73f84a45a4..7519f9bcdf 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -72,7 +72,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 155374b86d..4d599951f4 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -140,6 +140,36 @@ pub const Block = struct { } if (s0 > window) v0[(4 * s0 + s0 - window) * D + 13] = inf; // a key block row 0 alone sees self.ctx.items[self.ns - 1][(3 * self.ctxLen(self.ns - 1) + self.ctxLen(self.ns - 1) - 1) * D] = nan; + // +inf values row 0 sees at a float64 weight of 0 (0 x inf: NaN), and at weights below float32's least (inf) + const k0 = self.ctx.items[0][0 .. hk * s0 * D]; + self.reach(k0, v0, 6, 24, null, 19); // the block's first key + if (s0 > window) self.reach(k0, v0, 5, 20, s0 - window + 8, 17); // the window's edge tile, rows 0 to 8 see it + self.under(k0, v0); + } + + /// Head `head`'s row 0: query all 8, key s0 - 1 all 8, key `at` (null: block key 0) all -8 with +inf at e. + fn reach(self: *Block, k0: []u16, v0: []u16, kh: usize, head: usize, at: ?usize, e: usize) void { + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + @memset(self.q.items[(head * rows) * D ..][0..D], c.toBf16(8)); + @memset(k0[(kh * s0 + s0 - 1) * D ..][0..D], c.toBf16(8)); + const far = if (at) |k| k0[(kh * s0 + k) * D ..][0..D] else self.kb.items[(kh * rows) * D ..][0..D]; + @memset(far, c.toBf16(-8)); + if (at) |k| v0[(kh * s0 + k) * D + e] = c.toBf16(std.math.inf(f32)) else self.vb.items[(kh * rows) * D + e] = c.toBf16(std.math.inf(f32)); + } + + /// Key head 7, stream 0: head 28's row 0 query all 2, key s0 - 1 all 8, every other key all -8 with +inf at 21. + fn under(self: *Block, k0: []u16, v0: []u16) void { + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + const kh = 7; + @memset(self.q.items[(28 * rows) * D ..][0..D], c.toBf16(2)); + for (0..s0 + self.len()) |k| { + const key = if (k < s0) k0[(kh * s0 + k) * D ..][0..D] else self.kb.items[(kh * rows + k - s0) * D ..][0..D]; + const val = if (k < s0) v0[(kh * s0 + k) * D ..][0..D] else self.vb.items[(kh * rows + k - s0) * D ..][0..D]; + @memset(key, c.toBf16(if (k == s0 - 1) 8 else -8)); + if (k != s0 - 1) val[21] = c.toBf16(std.math.inf(f32)); + } } /// Gate 4: every key alike (each p is 1); values mirrored over the context keys every row sees, else 0. From 462d952c246cb9116d1c2e7bd7d01d08e9ad29eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:58:23 +0200 Subject: [PATCH 36/41] cuda/window_attention: the prompt replay takes every NaN and infinite value out of the mma and adds it back at its key's float64 weight (attn_odd), so an infinite value gives NaN where float64 does --- zig/kernels/cuda/window_attention.cu | 18 +++++++++--------- zig/tests/cuda/contract_window.zig | 26 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index 3c6da22a58..e2333e52e9 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ float rowsum(const float (&c)[8]) { __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); } -// ODD: the edge tiles' NaN and infinite values leave the mma, and the rows that see them add them back. +// ODD: every NaN and infinite value leaves the mma, and the rows that see them add them back at float64's weights. template __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, @@ -275,9 +275,6 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, mma(o[i + 1], pa1, b[2], b[3]); } }; - const auto edge = [&](int t) { // tile t, where some row of the block cannot see some key - return (t0 + t) * BN < (win > 0 ? p0 + last - win + 1 : 0) || (t0 + t) * BN + BN - 1 > p0 + row0; - }; for (int t = 0; t < tiles; ++t) { const int it = 4 * t; scores(*reinterpret_cast(&sc[0]), take(it, false)); @@ -336,9 +333,9 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]); } } - values(pa[0], pa[1], take(it + 2, ODD && edge(t))); + values(pa[0], pa[1], take(it + 2, ODD)); done(it + 2); - values(pa[2], pa[3], take(it + 3, ODD && edge(t))); + values(pa[2], pa[3], take(it + 3, ODD)); done(it + 3); } wait_group<0>(); @@ -347,9 +344,12 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, if (__syncthreads_or(attn_odd::nonfinite_out(o, r0 < W, r1 < W))) return true; } else { const auto row = [&](int k) { return vc + ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D; }; - const auto vis = [&](int r, int k) { return k >= (win > 0 ? pw + r - win + 1 : 0) && k <= pw + r; }; - for (int t = 0; t < tiles; ++t) - if (edge(t)) attn_odd::add_back(o, row, (t0 + t) * BN, min((t0 + t + 1) * BN, limit), lane, vis, l0, l1); + const auto krow = [&](int k) { return kc + ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D; }; + const auto lo = [&](int r) { return win > 0 ? max(pw + r - win + 1, 0) : 0; }; + const auto vis = [&](int r, int k) { return wrow + r < W && k >= lo(r) && k <= pw + r; }; + const auto each = [&](int r, auto f) { for (int k = lo(r); k <= pw + r; ++k) f(krow(k)); }; // row r's keys + auto w = attn_odd::weights([&](int r) { return q + ((int64_t)(wrow + r) * H + head) * D; }, each, scale, lane); + attn_odd::add_back(o, row, krow, t0 * BN, limit, lane, vis, l0, l1, w); } #pragma unroll for (int i = 0; i < NB; ++i) { diff --git a/zig/tests/cuda/contract_window.zig b/zig/tests/cuda/contract_window.zig index 9108b3feb0..d3157476f8 100644 --- a/zig/tests/cuda/contract_window.zig +++ b/zig/tests/cuda/contract_window.zig @@ -150,6 +150,32 @@ pub const Window = struct { } if (n > 3) self.cell(self.kc.items, 3, self.at(3).p, (3 % hk), 0).* = nan; if (n > 4) self.cell(self.vc.items, 4, self.first(self.at(4).p), (4 % hk), 13).* = inf; + // prompt rows: +inf at float64 weight 0 (row 5's own key, a causal edge: NaN), and under float32's least (inf) + if (self.f.prompt and n > 6) { + if (hk > 2) self.reach(5, 1); + self.under(n - 1, 2 % hk); + } + } + + /// Row r, key head kh's first head: query all 8, key p - 1 all 8, key p all -8 with +inf at element 17. + fn reach(self: *Window, r: usize, kh: usize) void { + const f = self.f; + const p = self.at(r).p; + @memset(self.q.items[(r * f.h + kh * (f.h / f.hk)) * f.d ..][0..f.d], c.toBf16(8)); + @memset(self.kc.items[(self.cacheRow(0, p - 1) * f.hk + kh) * f.d ..][0..f.d], c.toBf16(8)); + @memset(self.kc.items[(self.cacheRow(0, p) * f.hk + kh) * f.d ..][0..f.d], c.toBf16(-8)); + self.cell(self.vc.items, r, p, kh, 17).* = c.toBf16(std.math.inf(f32)); + } + + /// Row r, key head kh's first head: query all 2, key p - 1 all 8, its other keys all -8 with +inf at element 21. + fn under(self: *Window, r: usize, kh: usize) void { + const f = self.f; + const p = self.at(r).p; + @memset(self.q.items[(r * f.h + kh * (f.h / f.hk)) * f.d ..][0..f.d], c.toBf16(2)); + for (self.first(p)..p + 1) |key| { + @memset(self.kc.items[(self.cacheRow(0, key) * f.hk + kh) * f.d ..][0..f.d], c.toBf16(if (key == p - 1) 8 else -8)); + if (key != p - 1) self.cell(self.vc.items, r, key, kh, 21).* = c.toBf16(std.math.inf(f32)); + } } /// Gate 4: every key alike (each p is 1); values mirrored over the keys every row sees, 0 elsewhere. From 3c82eb3dc6b694272418b68eb7c4439c698b131b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:37:54 +0200 Subject: [PATCH 37/41] cuda/block_attention: the replay's block tile keeps the infinities added back before it The replay adds the context's NaN and infinite values back before the block's own tile, whose rescale is exp(m - m') in float32. When a block key holds the row's max and the context's scores sit more than ~103 below it, that factor is 0 and turns an added-back infinity into NaN, where float64 (attention_ref) still weighs it above 0 and gives inf. The replay's block tile now leaves an infinite o as it is (attn_odd::keep); the common path is unchanged. Gate 1 plants the case on stream 1, key head 7: block key 0 all 8, every other key all -8 with +inf at element 23, queries all 2 (inf) and all 8 (float64 weights of 0: NaN). --- zig/kernels/cuda/attn_odd.cuh | 3 +++ zig/kernels/cuda/block_attention.cu | 13 +++++++------ zig/tests/cuda/contract_block_attn.zig | 22 ++++++++++++++++++++++ 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 7501de7b7c..1a72f5534c 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -35,6 +35,9 @@ __device__ __forceinline__ void clean(uint4* vs, int threads) { } } +// x a, an infinite x (the replay's, already weighed in float64) left as it is whatever the rescale. +__device__ __forceinline__ float keep(float x, float a) { return isinf(x) ? x : __fmul_rn(x, a); } + // The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. __device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index 84cd949888..8af5fa4852 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -42,6 +42,7 @@ __device__ __forceinline__ uint32_t pack(float lo, float hi) { __device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); } // One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V. +// KEEP: the replay's infinities already in o stay (attn_odd::keep), whatever the rescale. struct Fold { float o[NB][4]; float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f; @@ -51,7 +52,7 @@ struct Fold { for (int i = 0; i < NB; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f; } - template + template __device__ __forceinline__ void tile(const uint32_t (&qa)[K16][4], const uint4* ks, const uint4* vs, float scale, int lane, int k0, Vis vis) { const int g8 = lane / 4, qd = lane % 4, mt = lane / 8; @@ -109,10 +110,10 @@ struct Fold { m1 = n1; #pragma unroll for (int i = 0; i < NB; ++i) { - o[i][0] = __fmul_rn(o[i][0], a0); - o[i][1] = __fmul_rn(o[i][1], a0); - o[i][2] = __fmul_rn(o[i][2], a1); - o[i][3] = __fmul_rn(o[i][3], a1); + o[i][0] = KEEP ? attn_odd::keep(o[i][0], a0) : __fmul_rn(o[i][0], a0); + o[i][1] = KEEP ? attn_odd::keep(o[i][1], a0) : __fmul_rn(o[i][1], a0); + o[i][2] = KEEP ? attn_odd::keep(o[i][2], a1) : __fmul_rn(o[i][2], a1); + o[i][3] = KEEP ? attn_odd::keep(o[i][3], a1) : __fmul_rn(o[i][3], a1); } uint32_t pa[2][4]; #pragma unroll @@ -219,7 +220,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( block_keys(); attn_odd::clean(vs, blockDim.x); __syncthreads(); - f.tile(qa, ks, vs, scale, lane, 0, own); + f.tile(qa, ks, vs, scale, lane, 0, own); // its rescale (0 when it holds the max) keeps the infinities const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; attn_odd::add_back(f.o, row, bk, 0, L, lane, own, f.l0, f.l1, w); } diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 4d599951f4..1197339202 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -145,6 +145,28 @@ pub const Block = struct { self.reach(k0, v0, 6, 24, null, 19); // the block's first key if (s0 > window) self.reach(k0, v0, 5, 20, s0 - window + 8, 17); // the window's edge tile, rows 0 to 8 see it self.under(k0, v0); + if (self.ns > 1) self.ownMax(1, 7, 23); + } + + /// Stream j, key head kh: block key 0 all 8, every other key all -8 with +inf at e, the block's row 0 querying all 2 + /// (head 4 kh: weights below float32's least, inf) and all 8 (head 4 kh + 1: float64 weights of 0, NaN). + fn ownMax(self: *Block, j: usize, kh: usize, e: usize) void { + const L = self.len(); + const rows = self.ns * L; + const s = self.ctxLen(j); + const x = self.ctx.items[j]; + const inf = c.toBf16(std.math.inf(f32)); + const G = h / hk; + @memset(self.q.items[((kh * G) * rows + j * L) * D ..][0..D], c.toBf16(2)); + @memset(self.q.items[((kh * G + 1) * rows + j * L) * D ..][0..D], c.toBf16(8)); + for (0..s) |k| { + @memset(x[(kh * s + k) * D ..][0..D], c.toBf16(-8)); + x[hk * s * D + (kh * s + k) * D + e] = inf; + } + for (0..L) |col| { + @memset(self.kb.items[(kh * rows + j * L + col) * D ..][0..D], c.toBf16(if (col == 0) 8 else -8)); + if (col != 0) self.vb.items[(kh * rows + j * L + col) * D + e] = inf; + } } /// Head `head`'s row 0: query all 8, key s0 - 1 all 8, key `at` (null: block key 0) all -8 with +inf at e. From 0facb0f104d5c7a000726c2586380cdf70e2ecfc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:38:04 +0200 Subject: [PATCH 38/41] cuda/tree_attention: a NaN or infinite partial runs again with those values out of the mma, added back at float64's weights A visible infinite value whose weight underflows in float32 (more than ~103 below the max) but not in float64 (745) went through the mma as 0 x inf, or reached a merge at a factor of 0: NaN, where attention_ref gives inf. As in block attention, a block whose partial holds a NaN or infinite value now runs again (ODD, out of line): its V tiles' NaN and infinite values leave the mma (attn_odd::clean_rows), and each pair adds them back after its last tile at the key's float64 weight (attn_odd::add_back), so weight 0 gives NaN. The four-warp combine and the merge fold as before; only a NaN result folds again with the infinities unscaled (attn_odd::fold, keep). tf_tree_shared takes the window's keys (kn) for the path keys' weights. Finite outputs keep their bits. Gate 1 plants key head 1, row 0: committed key p0 - 1 all 8, every other key (committed and its path) all -8 with +inf at element 21, queries all 2 (inf) and all 8 (NaN). --- zig/kernels/cuda/attn_odd.cuh | 19 ++++-- zig/kernels/cuda/tree_attention.cu | 91 +++++++++++++++++++++++---- zig/src/core/contract/registry.zig | 26 ++++---- zig/src/cuda/tree_attention.zig | 1 + zig/tests/cuda/contract_tree_attn.zig | 24 ++++++- 5 files changed, 131 insertions(+), 30 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 1a72f5534c..a7c4354f75 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -25,19 +25,30 @@ __device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, return (r0 && a != a) || (r1 && b != b); } -// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. -template -__device__ __forceinline__ void clean(uint4* vs, int threads) { - for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) { +// A KEYS-key tile staged [key][D] (xor swizzle): chunks tid, tid + threads, ... (the ones tid copied) made finite. +template +__device__ __forceinline__ void clean_rows(uint4* vs, int tid, int threads) { + for (int i = tid; i < KEYS * (D / 8); i += threads) { uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7)); const uint4 x = *p; *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w)); } } +// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. +template +__device__ __forceinline__ void clean(uint4* vs, int threads) { + clean_rows(vs, threadIdx.x, threads); +} + // x a, an infinite x (the replay's, already weighed in float64) left as it is whatever the rescale. __device__ __forceinline__ float keep(float x, float a) { return isinf(x) ? x : __fmul_rn(x, a); } +// x a + y folding partials, an infinite x left unscaled (see keep). +__device__ __forceinline__ float fold(float x, float a, float y) { + return isinf(x) ? __fadd_rn(x, y) : __fmaf_rn(x, a, y); +} + // The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. __device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 4e50f442e9..2673273642 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -5,6 +5,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_tree { @@ -58,8 +59,10 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ } // Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -template -__device__ __forceinline__ void tree_group( +// ODD: every NaN and infinite value leaves the mma, and the pairs add them back at float64's weights. +// !ODD: true (nothing written) when some warp's partial is NaN or infinite: the block runs again as ODD. +template +__device__ __forceinline__ bool tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, @@ -75,7 +78,7 @@ __device__ __forceinline__ void tree_group( const int s = row_stream[r]; const int p = streams[4 * s + 2], depth = depths[r]; const int total = p + depth, k0 = gi * SPAN, k1 = min(total, k0 + SPAN); - if (k0 >= total) return; // the merge reads only the row's own groups + if (k0 >= total) return false; // the merge reads only the row's own groups const __nv_bfloat16* kc = origin + offs[2 * s]; const __nv_bfloat16* vc = origin + offs[2 * s + 1]; // ST tile stages: a one-warp block fetches ST - 1 tiles ahead while it computes one (stages never change bits) @@ -149,7 +152,7 @@ __device__ __forceinline__ void tree_group( ks = kv0 + st * 2 * HB * C; vs = ks + HB * C; st = (st + 1) % ST; - + if constexpr (ODD) attn_odd::clean_rows(vs, lane, 32); // the chunks this lane copied __syncwarp(); float sc[NT8][4]; #pragma unroll @@ -238,6 +241,25 @@ __device__ __forceinline__ void tree_group( } __syncwarp(); } + if constexpr (!ODD) { // a NaN or infinite partial: the block runs again as ODD (finite values: the same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(o, pf + g8 < npairs, pf + g8 + 8 < npairs))) return true; + } else { + // key k of the group's row: committed, then the row's path; pair m: row rq + (pf + m) / G, head (pf + m) % G + const auto at = [&](const __nv_bfloat16* c, const __nv_bfloat16* n, int k) { + return k < p ? c + ((int64_t)k * hk + kh) * D : n + ((int64_t)paths[r * MAXD + (k - p)] * hk + kh) * D; + }; + const auto qrow = [&](int m) { return q + ((rq + (pf + m) / G) * h + kh * G + (pf + m) % G) * D; }; + const auto each = [&](int m, auto f) { // pair m's keys: its stream's committed ones, then its own row's path + const int64_t pr = rq + (pf + m) / G; + for (int k = 0; k < p; ++k) f(kc + ((int64_t)k * hk + kh) * D); + for (int j = 0; j < depths[pr]; ++j) f(kn + ((int64_t)paths[pr * MAXD + j] * hk + kh) * D); + }; + auto w = attn_odd::weights(qrow, each, scale, lane); + const auto vis = [&](int m, int) { return pf + m < npairs; }; // a pair sees every key of its group + for (int t = t0; t * HB < k1; t += 4) // after the last tile: no rescale turns an infinity into NaN + attn_odd::add_back(o, [&](int k) { return at(vc, vn, k); }, [&](int k) { return at(kc, kn, k); }, + t * HB, min(k1, t * HB + HB), lane, vis, l0, l1, w); + } if constexpr (W == 1) { // one warp a block: each slice's partial (o, max, sum) of its 16 pairs; the merge combines the 4 in order #pragma unroll @@ -254,7 +276,7 @@ __device__ __forceinline__ void tree_group( } if (qd == 0) pm[part * G + head] = half ? m1 : m0, pl[part * G + head] = half ? l1 : l0; } - return; + return false; } else { __syncthreads(); float* os = reinterpret_cast(smem); // [4 warps][16 heads][D], then max and sum [4][16] @@ -288,19 +310,37 @@ __device__ __forceinline__ void tree_group( acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); l = __fmaf_rn(a, ls[w * 16 + hd], l); } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a weight of 0; again, kept + acc = 0.0f; + for (int w = 0; w < 4; ++w) + if (ms[w * 16 + hd] != -INFINITY) + acc = attn_odd::fold(os[(w * 16 + hd) * D + e], texp(__fsub_rn(ms[w * 16 + hd], mm)), acc); + } po[(part * G + head) * D + e] = acc; if (e == 0) pm[part * G + head] = mm, pl[part * G + head] = l; } + return false; } } +// A block whose partial is NaN or infinite, run again as ODD (out of the common path). +template +__device__ __noinline__ void tree_odd( + const __nv_bfloat16* q, const __nv_bfloat16* kn, const __nv_bfloat16* vn, const __nv_bfloat16* origin, + const long long* offs, const int* row_stream, const int* streams, const int* paths, const int* depths, float* po, + float* pm, float* pl, int h, int hk, float scale, int gi, int ng, int64_t rq, int pf, int npairs, int64_t row, + int slice) { + tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, + gi, ng, rq, pf, npairs, row, slice); +} + // A stream's full committed groups once for its rows, 16 (row, head) pairs a block: each partial a row block's bits. template __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, - float scale, int ng, int n) { + float scale, int ng, int n, const __nv_bfloat16* __restrict__ kn) { const int G = h / hk; const int per = W == 4 ? 1 : 4; // slices an item takes a block for (int item = blockIdx.x; item < n * ng * per; item += gridDim.x) { // (row, group[, slice]): a fixed grid @@ -311,8 +351,11 @@ __device__ __forceinline__ void tree_shared( const int pairs = streams[4 * st + 1] * G, p = streams[4 * st + 2]; const int tile = static_cast(r - r0); // a stream's j-th row hosts its j-th pair tile (tiles <= rows) if (tile * 16 >= pairs || (gi + 1) * SPAN > p) continue; - tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, - r0, tile * 16, pairs, r, item % per); + // vn: never read, a full committed group has no path keys (the replay weighs them, from kn) + if (tree_group(q, kn, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, + scale, gi, ng, r0, tile * 16, pairs, r, item % per)) + tree_odd(q, kn, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + ng, r0, tile * 16, pairs, r, item % per); __syncthreads(); } } @@ -322,8 +365,10 @@ __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ - float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { \ - tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ + float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n, \ + const __nv_bfloat16* __restrict__ kn) { \ + tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n, \ + kn); \ } TF_TREE_SHARED(tf_tree_shared, 128, 32, 4, 1) @@ -344,8 +389,10 @@ __device__ __forceinline__ void tree_part( const int p = streams[4 * row_stream[r] + 2], total = p + depths[r]; const int per = W == 4 ? 1 : 4, slot = blockIdx.z / per, slots = gridDim.z / per; for (int gi = slot + p / SPAN; gi < ng && gi * SPAN < total; gi += slots) { // full committed: shared - tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, - ng, r, 0, h / hk, r, blockIdx.z % per); + if (tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, + scale, gi, ng, r, 0, h / hk, r, blockIdx.z % per)) + tree_odd(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + ng, r, 0, h / hk, r, blockIdx.z % per); __syncthreads(); // the next group's tiles reuse the partials' shared memory } } @@ -417,6 +464,26 @@ extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( m = next; } } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a weight of 0 here; again, kept + float mm = -INFINITY; + acc = 0.0f; + for (int s0 = 0; s0 < slots; s0 += per) { + float cm = -INFINITY, co = 0.0f, cl = 0.0f; + for (int w = 0; w < per; ++w) cm = fmaxf(cm, pm[(base + s0 + w) * G + hd]); + for (int w = 0; w < per; ++w) { + const int64_t q = (base + s0 + w) * G + hd; + if (pm[q] == -INFINITY) continue; + const float aw = texp(__fsub_rn(pm[q], cm)); + co = attn_odd::fold(po[q * d + e], aw, co); + cl = __fmaf_rn(aw, pl[q], cl); + } + if (!(cl > 0.0f)) continue; + const float next = fmaxf(mm, cm); + acc = attn_odd::fold(acc, mm == -INFINITY ? 0.0f : texp(__fsub_rn(mm, next)), + attn_odd::keep(co, texp(__fsub_rn(cm, next)))); + mm = next; + } + } out[(r * h + kh * G + hd) * d + e] = __float2bfloat16_rn(__fdiv_rn(acc, l)); } } diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7519f9bcdf..778a714ce0 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -59,19 +59,19 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "dense rows tf_dense_rows_16", .symbol = "tf_dense_rows_16", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 79, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "dense rows tf_dense_rows_32", .symbol = "tf_dense_rows_32", .image = "dense_rows", .threads = 256 }, .usage = .{ .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 127, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 }, .{ .regs = 180, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "tree attention tf_tree_paths", .symbol = "tf_tree_paths", .image = "tree_attention", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 }, .{ .regs = 167, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared_w1s3", .symbol = "tf_tree_shared_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part", .symbol = "tf_tree_part", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part_w1", .symbol = "tf_tree_part_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part_w1s3", .symbol = "tf_tree_part_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 }, .{ .regs = 167, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared_256", .symbol = "tf_tree_shared_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1", .symbol = "tf_tree_shared_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 }, .{ .regs = 237, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1s3", .symbol = "tf_tree_shared_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 212, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 }, .{ .regs = 236, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part_256", .symbol = "tf_tree_part_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 230, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 }, .{ .regs = 231, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1s3", .symbol = "tf_tree_shared_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part", .symbol = "tf_tree_part", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1", .symbol = "tf_tree_part_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1s3", .symbol = "tf_tree_part_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256", .symbol = "tf_tree_shared_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1", .symbol = "tf_tree_shared_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1s3", .symbol = "tf_tree_shared_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256", .symbol = "tf_tree_part_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 } } }, }; diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index 69b2ad598f..d4fc29afb2 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -109,6 +109,7 @@ pub const TreeAttention = struct { c.add(scale); c.add(@as(i32, @intCast(ng))); c.add(@as(i32, @intCast(n))); + c.add(kn); // the replay of a NaN or infinite partial weighs a row's path keys too const per: usize = if (one) slices else 1; const block: u32 = if (one) 32 else 128; // three stages where the one-warp grid fits one wave of the SMs (a block each), else two (two blocks an SM) diff --git a/zig/tests/cuda/contract_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig index 3dd57f47e6..466526de7a 100644 --- a/zig/tests/cuda/contract_tree_attn.zig +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -153,7 +153,8 @@ pub const Tree = struct { return true; } - /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key. + /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key; + /// +inf at float64 weights of 0 and below float32's least. fn plant(self: *Tree) void { const d = self.v.d; const kv = self.v.hk * d; @@ -171,6 +172,27 @@ pub const Tree = struct { self.cache.items[vo + (p0 - 3) * kv + 2 * d + 11] = ninf; } self.cache.items[ko + (p0 - 1) * kv + 3 * d + 0] = nan; + self.under(1, 21); + } + + /// Key head kh, row 0: committed key p0 - 1 all 8, its other keys all -8 with +inf at e. Head kh's first query is + /// all 2 (weights below float32's least: inf), its second all 8 (float64 weights of 0, 0 x inf: NaN). + fn under(self: *Tree, kh: usize, e: usize) void { + const d = self.v.d; + const kv = self.v.hk * d; + const G = self.v.h / self.v.hk; + const p0 = self.committed.items[0]; + const ko: usize = @intCast(self.offs.items[0]); + const vo: usize = @intCast(self.offs.items[1]); + const inf = c.toBf16(std.math.inf(f32)); + @memset(self.q.items[(kh * G) * d ..][0..d], c.toBf16(2)); + @memset(self.q.items[(kh * G + 1) * d ..][0..d], c.toBf16(8)); + for (0..p0) |j| { + @memset(self.cache.items[ko + j * kv + kh * d ..][0..d], c.toBf16(if (j == p0 - 1) 8 else -8)); + if (j != p0 - 1) self.cache.items[vo + j * kv + kh * d + e] = inf; + } + @memset(self.kn.items[kh * d ..][0..d], c.toBf16(-8)); // row 0's path: itself + self.vn.items[kh * d + e] = inf; } /// Gate 4: every key alike (each p is 1), committed values mirrored per column, path values 0. From 8cba69e1472ca7b57a8c712fadfdda4164cf4c4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:05:08 +0200 Subject: [PATCH 39/41] cuda/window_attention: the decode kernels run a NaN or infinite tile again with those values out of the sums, added back at float64's weights A visible infinite value whose weight underflows in float32 (more than ~103 below the max) but not in float64 (745) went into a sum as 0 x inf, or reached a fold of partials at a factor of 0: NaN, where attention_ref gives inf. The prompt rows had the replay (462d952c); the decode kernels now have it too, out of line: - tf_window_decode_tc: a NaN output (the one wrong class: an infinite output holds a weight above float32's least, so above float64's) runs the block, or the one warp that saw keys, again as ODD. Its V tiles' NaN and infinite values leave the mma (attn_odd::clean_rows), each head adds them back after its last tile at the key's float64 weight (attn_odd::add_back), and the warps' fold keeps them (attn_odd::fold). - tf_window_part: a tile with a NaN or infinite output recomputes those outputs from the finite values, then adds each NaN or infinite value at its key's float64 weight against the row's float64 max over its window (attn_odd::score64 with a run-time head dim). tf_window_merge folds a NaN result again with the infinities unscaled. Finite outputs keep their bits. Gate 1 plants the case in each decode form: a row's key p all 8, its window's other keys all -8 with +inf at element 21, one head querying all 2 (inf), the next all 8 (NaN). --- zig/kernels/cuda/attn_odd.cuh | 9 +++ zig/kernels/cuda/window_attention.cu | 113 ++++++++++++++++++++++++--- zig/src/core/contract/registry.zig | 4 +- zig/tests/cuda/contract_window.zig | 20 ++++- 4 files changed, 134 insertions(+), 12 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index a7c4354f75..5ef139d587 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -62,6 +62,15 @@ __device__ __forceinline__ double score64(const __nv_bfloat16* q, const __nv_bfl return __dmul_rn(s, scale); } +// The same with a head dim known only at run time. +__device__ __forceinline__ double score64(const __nv_bfloat16* q, const __nv_bfloat16* k, int d, double scale) { + double s = 0.0; +#pragma unroll 1 + for (int e = 0; e < d; ++e) + s = __dadd_rn(s, __dmul_rn((double)__bfloat162float(q[e]), (double)__bfloat162float(k[e]))); + return __dmul_rn(s, scale); +} + // Rows g, g + 8: w(r, K row) is attention_ref's exp(score - max) > 0; each(r, f) walks r's K rows (max on first use). template struct Weights { diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index e2333e52e9..1b8c93028f 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -22,6 +22,43 @@ constexpr int DMAX = 128; // head dim (static shared memory: q, a K and a V til using namespace tf_window; +// tf_window_part's tile with a NaN or infinite output (out of the common path): those outputs again, finite values +// summed as before, then each NaN or infinite value the row sees at its key's float64 weight (0: inf x 0 is NaN). +__device__ __noinline__ void window_part_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const float* sc, + const __nv_bfloat16* vs, double* m64, float* po, int64_t r, int kh, + int g, int h, int hk, int d, int t, int lo, int p, int64_t base, + long long ring, float scale) { + const auto krow = [&](int key) { return kc + ((base + key % ring) * hk + kh) * d; }; + const auto qrow = [&](int hd) { return q + (r * h + kh * g + hd) * d; }; + { // each head's float64 max over the row's window, attention_ref's scores: a warp a head, lanes over keys + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + for (int hd = warp; hd < g; hd += NT / 32) { + double mx = -INFINITY; + for (int key = lo + lane; key <= p; key += 32) mx = fmax(mx, attn_odd::score64(qrow(hd), krow(key), d, scale)); + for (int k = 16; k > 0; k >>= 1) mx = fmax(mx, __shfl_xor_sync(0xFFFFFFFFu, mx, k)); + if (lane == 0) m64[hd] = mx; + } + } + __syncthreads(); + for (int o = threadIdx.x; o < g * d; o += NT) { + if (isfinite(po[o])) continue; + const int hd = o / d, e = o % d; + float a = 0.f; + for (int jj = 0; jj < BN; ++jj) { + const float v = __bfloat162float(vs[jj * d + e]); + if (isfinite(v)) a = __fmaf_rn(sc[hd * BN + jj], v, a); + } + for (int jj = 0; jj < BN; ++jj) { + const float v = __bfloat162float(vs[jj * d + e]); + const int key = t * BN + jj; + if (isfinite(v) || key < lo || key > p) continue; + const bool w = __dsub_rn(attn_odd::score64(qrow(hd), krow(key), d, scale), m64[hd]) >= attn_odd::least(); + a = __fadd_rn(a, w ? v : __fmul_rn(0.f, v)); // NaN score: no weight + } + po[o] = a; + } +} + // Decode split by key tile (block: row, key head, tile); tf_window_merge folds the partials in tile order. extern "C" __global__ void __launch_bounds__(NT) tf_window_part(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, @@ -80,13 +117,19 @@ extern "C" __global__ void __launch_bounds__(NT) tf_window_part(const __nv_bfloa } } __syncthreads(); + bool odd = false; for (int o = threadIdx.x; o < g * d; o += NT) { const int hd = o / d, e = o % d; float a = 0.f; for (int jj = 0; jj < BN; ++jj) a = __fmaf_rn(sc[hd * BN + jj], __bfloat162float(vs[jj * d + e]), a); po[part * g * d + o] = a; + odd |= !isfinite(a); } if (threadIdx.x < g) pm[part * g + threadIdx.x] = mx[threadIdx.x], pl[part * g + threadIdx.x] = ls[threadIdx.x]; + if (__syncthreads_or(odd)) { // a NaN or infinite output: those again (finite outputs keep their bits) + __shared__ double m64[GMAX]; + window_part_odd(q, kc, sc, vs, m64, po + part * g * d, r, kh, g, h, hk, d, t, lo, p, base, ring, scale); + } } // out [W, H, D] bf16: a row's partials folded in tile order, o = sum_j exp(m_j - M) o_j / sum_j exp(m_j - M) l_j. @@ -110,6 +153,13 @@ extern "C" __global__ void __launch_bounds__(NT) tf_window_merge(const float* __ acc = __fmaf_rn(a, po[(part0 + j) * g * d + o], acc); l = __fmaf_rn(a, pl[(part0 + j) * g + hd], l); } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a factor of 0; again, kept + acc = 0.f; + for (int j = 0; j < tiles; ++j) { + const float mj = pm[(part0 + j) * g + hd]; + if (mj != -INFINITY) acc = attn_odd::fold(po[(part0 + j) * g * d + o], expf(__fsub_rn(mj, m)), acc); + } + } out[(r * h + kh * g + hd) * d + o % d] = __float2bfloat16_rn(__fdiv_rn(acc, l)); } } @@ -382,7 +432,11 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric } // Decode on tensor cores (d 128): a row's heads one mma tile; warps take every 4th 32-key tile, fold in warp order. -extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( +// ODD: every NaN and infinite value leaves the mma, and the heads add them back at float64's weights. +// !ODD: true when some output is NaN, the one wrong class a 0 x inf can make (an infinite output holds a weight +// above float32's least, so above float64's): the block (one warp: its warp) runs again as ODD over it. +template +__device__ __forceinline__ bool decode_tc_rows( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, const int* __restrict__ pos, const int* __restrict__ slot, int h, int hk, int win, long long ring, float scale) { @@ -499,6 +553,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( o[i][3] = __fmul_rn(o[i][3], a1); } wait_group<1>(); // tile t's V (the next K may still be in flight) + if constexpr (ODD) attn_odd::clean_rows(vs, lane, 32); // the chunks this lane copied __syncwarp(); uint32_t pa[2][4]; #pragma unroll @@ -527,20 +582,31 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( } wait_group<0>(); const int na = min(4, p / HB - lo / HB + 1); // warps that saw keys + if constexpr (ODD) { + const auto krow = [&](int k) { return kc + ((base + k % ring) * hk + kh) * D; }; + const auto vrow = [&](int k) { return vc + ((base + k % ring) * hk + kh) * D; }; + const auto qrow = [&](int m) { return q + ((int64_t)r * h + kh * G + m) * D; }; + const auto each = [&](int, auto f) { for (int k = lo; k <= p; ++k) f(krow(k)); }; // the row's window + auto w = attn_odd::weights(qrow, each, scale, lane); + const auto vis = [&](int m, int) { return m < G; }; + for (int t = lo / HB + warp; t <= p / HB; t += 4) // after the last tile: no rescale turns an infinity into NaN + attn_odd::add_back(o, vrow, krow, max(lo, t * HB), min(p + 1, t * HB + HB), lane, vis, l0, l1, w); + } if (na == 1) { // one warp's fold is the result (the merge below would scale it by exp(0) = 1) if (warp == 0) { + bool nan = false; #pragma unroll for (int i = 0; i < NB; ++i) { const int d = 8 * i + 2 * qd; - if (g8 < G) - *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8) * D + d) = - pack(tdiv(o[i][0], l0), tdiv(o[i][1], l0)); - if (g8 + 8 < G) - *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8 + 8) * D + d) = - pack(tdiv(o[i][2], l1), tdiv(o[i][3], l1)); + const float x0 = tdiv(o[i][0], l0), x1 = tdiv(o[i][1], l0), x2 = tdiv(o[i][2], l1), x3 = tdiv(o[i][3], l1); + if (g8 < G) *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8) * D + d) = pack(x0, x1); + if (g8 + 8 < G) *reinterpret_cast(out + ((int64_t)r * h + kh * G + g8 + 8) * D + d) = pack(x2, x3); + nan |= (g8 < G && (x0 != x0 || x1 != x1)) || (g8 + 8 < G && (x2 != x2 || x3 != x3)); } + // a NaN output (a 0 x inf the fold could not avoid; an infinity here is float64's): the warp again as ODD + if (!ODD && __any_sync(0xffffffffu, nan)) return true; } - return; + return false; } __syncthreads(); float* os = reinterpret_cast(smem); // [4 warps][16 heads][D] @@ -575,6 +641,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( il[hd] = l; } __syncthreads(); + bool nan = false; for (int idx = threadIdx.x; idx < G * D / 2; idx += 128) { const int hd = idx / (D / 2), e = 2 * (idx % (D / 2)); float acc0 = 0.0f, acc1 = 0.0f; @@ -586,9 +653,37 @@ extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( acc0 = __fmaf_rn(a, v.x, acc0); acc1 = __fmaf_rn(a, v.y, acc1); } + if (ODD && (acc0 != acc0 || acc1 != acc1)) { // NaN: a replayed infinity (float64's) met a factor of 0; kept + acc0 = acc1 = 0.0f; + for (int w = 0; w < 4; ++w) { + if (ms[w * 16 + hd] == -INFINITY) continue; + const float a = as[w * 16 + hd]; + acc0 = attn_odd::fold(os[(w * 16 + hd) * D + e], a, acc0); + acc1 = attn_odd::fold(os[(w * 16 + hd) * D + e + 1], a, acc1); + } + } const float l = il[hd]; - *reinterpret_cast(out + ((int64_t)r * h + kh * G + hd) * D + e) = pack(tdiv(acc0, l), tdiv(acc1, l)); + const float x0 = tdiv(acc0, l), x1 = tdiv(acc1, l); + *reinterpret_cast(out + ((int64_t)r * h + kh * G + hd) * D + e) = pack(x0, x1); + nan |= x0 != x0 || x1 != x1; } + // a NaN output (a 0 x inf in a fold or here; an infinity is float64's then): the block again as ODD + return !ODD && __syncthreads_or(nan); +} + +// A block (or the one warp that saw keys) with a NaN output, run again with ODD on (out of the common path). +__device__ __noinline__ void decode_tc_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, + __nv_bfloat16* out, const int* pos, const int* slot, int h, int hk, int win, + long long ring, float scale) { + decode_tc_rows(q, kc, vc, out, pos, slot, h, hk, win, ring, scale); +} + +extern "C" __global__ void __launch_bounds__(128) tf_window_decode_tc( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, + __nv_bfloat16* __restrict__ out, const int* __restrict__ pos, const int* __restrict__ slot, int h, int hk, int win, + long long ring, float scale) { + if (decode_tc_rows(q, kc, vc, out, pos, slot, h, hk, win, ring, scale)) + decode_tc_odd(q, kc, vc, out, pos, slot, h, hk, win, ring, scale); } } // namespace tf_window_prompt diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 778a714ce0..fc84b24fa1 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -46,9 +46,9 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "window attention tf_window_part", .symbol = "tf_window_part", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 51, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 44, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 }, .{ .regs = 48, .static = 45696, .bound = 256 } } }, + .{ .launch = .{ .name = "window attention tf_window_part", .symbol = "tf_window_part", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 55, .static = 45824, .bound = 256 }, .{ .regs = 50, .static = 45824, .bound = 256 }, .{ .regs = 50, .static = 45824, .bound = 256 }, .{ .regs = 56, .static = 45824, .bound = 256 }, .{ .regs = 56, .static = 45824, .bound = 256 } } }, .{ .launch = .{ .name = "window attention tf_window_merge", .symbol = "tf_window_merge", .image = "window_attention", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 }, .{ .regs = 48, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "window attention tf_window_decode_tc", .symbol = "tf_window_decode_tc", .image = "window_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "window attention tf_window_decode_tc", .symbol = "tf_window_decode_tc", .image = "window_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 176, .static = 0, .bound = 128 }, .{ .regs = 176, .static = 0, .bound = 128 }, .{ .regs = 176, .static = 0, .bound = 128 }, .{ .regs = 178, .static = 0, .bound = 128 }, .{ .regs = 178, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 4, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{4}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "window attention wattn_kernel<128, 8, 8, 8>", .symbol = std.fmt.comptimePrint(wattn_symbol, .{8}), .image = "window_attention", .threads = 256, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 223, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 }, .{ .regs = 222, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "moe tf_moe_topk_sigmoid", .symbol = "tf_moe_topk_sigmoid", .image = "moe_glue", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 }, .{ .regs = 48, .static = 16704, .bound = 256 } } }, diff --git a/zig/tests/cuda/contract_window.zig b/zig/tests/cuda/contract_window.zig index d3157476f8..d9cf2e0863 100644 --- a/zig/tests/cuda/contract_window.zig +++ b/zig/tests/cuda/contract_window.zig @@ -134,7 +134,8 @@ pub const Window = struct { return &x[(self.cacheRow(w.base, key) * self.f.hk + kh) * self.f.d + e]; } - /// NaN in a query; +inf and -inf in values; both in one column; NaN in a key; +inf in a window's first value. + /// NaN in a query; +inf and -inf in values; both in one column; NaN in a key; +inf in a window's first value; + /// +inf at float64 weights of 0 and below float32's least. fn plant(self: *Window) void { const n = self.shape.rows; const nan = c.toBf16(std.math.nan(f32)); @@ -155,6 +156,23 @@ pub const Window = struct { if (hk > 2) self.reach(5, 1); self.under(n - 1, 2 % hk); } + // decode rows: the same over a row's window (unwindowed: row 4, whose keys stop short of row 3's NaN key) + if (!self.f.prompt and n > 7) self.split(if (self.f.win == 0) 4 else n - 1, if (hk > 2) 2 else 0, 21); + } + + /// Row r, key head kh: key p all 8, the window's other keys all -8 with +inf at e; kh's first head queries all 2 + /// (weights below float32's least: inf), its second all 8 (float64 weights of 0: NaN). Tiles split the keys. + fn split(self: *Window, r: usize, kh: usize, e: usize) void { + const f = self.f; + const p = self.at(r).p; + const G = f.h / f.hk; + @memset(self.q.items[(r * f.h + kh * G) * f.d ..][0..f.d], c.toBf16(2)); + @memset(self.q.items[(r * f.h + kh * G + 1) * f.d ..][0..f.d], c.toBf16(8)); + for (self.first(p)..p + 1) |key| { + const k = self.cell(self.kc.items, r, key, kh, 0); + @memset(@as([*]u16, @ptrCast(k))[0..f.d], c.toBf16(if (key == p) 8 else -8)); + if (key != p) self.cell(self.vc.items, r, key, kh, e).* = c.toBf16(std.math.inf(f32)); + } } /// Row r, key head kh's first head: query all 8, key p - 1 all 8, key p all -8 with +inf at element 17. From 965c8cd10cd8a6d062c65aa67cb1c5a9414c1021 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:25:26 +0200 Subject: [PATCH 40/41] cuda/tree_attention, block_attention: the replay's comments on one line each --- zig/kernels/cuda/block_attention.cu | 3 +-- zig/kernels/cuda/tree_attention.cu | 6 ++---- zig/tests/cuda/contract_block_attn.zig | 3 +-- zig/tests/cuda/contract_tree_attn.zig | 6 ++---- 4 files changed, 6 insertions(+), 12 deletions(-) diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index 8af5fa4852..b8675f80d0 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -41,8 +41,7 @@ __device__ __forceinline__ uint32_t pack(float lo, float hi) { __device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); } -// One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V. -// KEEP: the replay's infinities already in o stay (attn_odd::keep), whatever the rescale. +// One warp's fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax, P x V (KEEP: o's inf stays). struct Fold { float o[NB][4]; float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f; diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 2673273642..5ab4bdf8f8 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -58,9 +58,7 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ for (int slot = depth - 1; slot >= 0; --slot, cur = parents[cur]) paths[node * MAXD + slot] = cur; } -// Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -// ODD: every NaN and infinite value leaves the mma, and the pairs add them back at float64's weights. -// !ODD: true (nothing written) when some warp's partial is NaN or infinite: the block runs again as ODD. +// Group gi of key head kh for 16 (row, head) pairs from pf (m: row rq + m / G, head m % G); !ODD true: run tree_odd. template __device__ __forceinline__ bool tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, @@ -323,7 +321,7 @@ __device__ __forceinline__ bool tree_group( } } -// A block whose partial is NaN or infinite, run again as ODD (out of the common path). +// A block with a NaN or infinite partial, again as ODD: those values out of the mma, added back at float64's weights. template __device__ __noinline__ void tree_odd( const __nv_bfloat16* q, const __nv_bfloat16* kn, const __nv_bfloat16* vn, const __nv_bfloat16* origin, diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 1197339202..bc343b9938 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -148,8 +148,7 @@ pub const Block = struct { if (self.ns > 1) self.ownMax(1, 7, 23); } - /// Stream j, key head kh: block key 0 all 8, every other key all -8 with +inf at e, the block's row 0 querying all 2 - /// (head 4 kh: weights below float32's least, inf) and all 8 (head 4 kh + 1: float64 weights of 0, NaN). + /// Stream j, key head kh: block key 0 all 8, other keys -8 with +inf at e; row 0 queries 2 (inf) and 8 (NaN). fn ownMax(self: *Block, j: usize, kh: usize, e: usize) void { const L = self.len(); const rows = self.ns * L; diff --git a/zig/tests/cuda/contract_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig index 466526de7a..a81f6493df 100644 --- a/zig/tests/cuda/contract_tree_attn.zig +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -153,8 +153,7 @@ pub const Tree = struct { return true; } - /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key; - /// +inf at float64 weights of 0 and below float32's least. + /// NaN in a query, +inf and -inf in values (committed, on a path, in one column), a NaN key, +inf at tiny weights. fn plant(self: *Tree) void { const d = self.v.d; const kv = self.v.hk * d; @@ -175,8 +174,7 @@ pub const Tree = struct { self.under(1, 21); } - /// Key head kh, row 0: committed key p0 - 1 all 8, its other keys all -8 with +inf at e. Head kh's first query is - /// all 2 (weights below float32's least: inf), its second all 8 (float64 weights of 0, 0 x inf: NaN). + /// Key head kh, row 0: key p0 - 1 all 8, its other keys -8 with +inf at e; queries 2 (inf) and 8 (weight 0: NaN). fn under(self: *Tree, kh: usize, e: usize) void { const d = self.v.d; const kv = self.v.hk * d; From 74b749fa04ebb233a51529c86ff9eff626d1f05e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:31:41 +0200 Subject: [PATCH 41/41] cuda/window_attention: the split decode and its merge move to window_split.cuh (the file was past 600 lines), comments on one line --- zig/kernels/cuda/window_attention.cu | 162 +------------------------- zig/kernels/cuda/window_split.cuh | 164 +++++++++++++++++++++++++++ zig/tests/cuda/contract_window.zig | 6 +- 3 files changed, 168 insertions(+), 164 deletions(-) create mode 100644 zig/kernels/cuda/window_split.cuh diff --git a/zig/kernels/cuda/window_attention.cu b/zig/kernels/cuda/window_attention.cu index 1b8c93028f..a39dd892f1 100644 --- a/zig/kernels/cuda/window_attention.cu +++ b/zig/kernels/cuda/window_attention.cu @@ -1,7 +1,5 @@ // Sliding-window attention for every CUDA family (float64 references: core/attention_ref.zig): decode and prompt rows. -// Decode: a block a (row, key head); keys walk 64-key tiles by absolute position, so a row's bits are its own. - #include #include #include @@ -9,160 +7,7 @@ #include "attn_odd.cuh" #include "qmm_frag.cuh" - -namespace tf_window { - -constexpr int NT = 256; -constexpr int BN = 64; // keys a tile -constexpr int GMAX = 16; // query heads a key head -constexpr int KT = BN + 2; // kt's row stride: a bank further each dim -constexpr int DMAX = 128; // head dim (static shared memory: q, a K and a V tile, scores within 48 KiB) - -} // namespace tf_window - -using namespace tf_window; - -// tf_window_part's tile with a NaN or infinite output (out of the common path): those outputs again, finite values -// summed as before, then each NaN or infinite value the row sees at its key's float64 weight (0: inf x 0 is NaN). -__device__ __noinline__ void window_part_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const float* sc, - const __nv_bfloat16* vs, double* m64, float* po, int64_t r, int kh, - int g, int h, int hk, int d, int t, int lo, int p, int64_t base, - long long ring, float scale) { - const auto krow = [&](int key) { return kc + ((base + key % ring) * hk + kh) * d; }; - const auto qrow = [&](int hd) { return q + (r * h + kh * g + hd) * d; }; - { // each head's float64 max over the row's window, attention_ref's scores: a warp a head, lanes over keys - const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; - for (int hd = warp; hd < g; hd += NT / 32) { - double mx = -INFINITY; - for (int key = lo + lane; key <= p; key += 32) mx = fmax(mx, attn_odd::score64(qrow(hd), krow(key), d, scale)); - for (int k = 16; k > 0; k >>= 1) mx = fmax(mx, __shfl_xor_sync(0xFFFFFFFFu, mx, k)); - if (lane == 0) m64[hd] = mx; - } - } - __syncthreads(); - for (int o = threadIdx.x; o < g * d; o += NT) { - if (isfinite(po[o])) continue; - const int hd = o / d, e = o % d; - float a = 0.f; - for (int jj = 0; jj < BN; ++jj) { - const float v = __bfloat162float(vs[jj * d + e]); - if (isfinite(v)) a = __fmaf_rn(sc[hd * BN + jj], v, a); - } - for (int jj = 0; jj < BN; ++jj) { - const float v = __bfloat162float(vs[jj * d + e]); - const int key = t * BN + jj; - if (isfinite(v) || key < lo || key > p) continue; - const bool w = __dsub_rn(attn_odd::score64(qrow(hd), krow(key), d, scale), m64[hd]) >= attn_odd::least(); - a = __fadd_rn(a, w ? v : __fmul_rn(0.f, v)); // NaN score: no weight - } - po[o] = a; - } -} - -// Decode split by key tile (block: row, key head, tile); tf_window_merge folds the partials in tile order. -extern "C" __global__ void __launch_bounds__(NT) tf_window_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 d, int win, long long ring, float scale) { - __shared__ float qs[GMAX * DMAX]; - __shared__ __nv_bfloat16 kt[DMAX * KT]; // K transposed, [d][key], padded: a warp reads 32 keys of one dim - __shared__ __nv_bfloat16 vs[BN * DMAX]; - __shared__ float sc[GMAX * BN]; - __shared__ float mx[GMAX], ls[GMAX]; - const int64_t r = blockIdx.x; - const int kh = blockIdx.y, tiles = gridDim.z, j = blockIdx.z; - const int g = h / hk; - 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 + 1, 0) : 0; - const int t = lo / BN + j; - const int64_t part = (r * hk + kh) * tiles + j; // this block's partial, g heads of d values - if (t > p / BN) { // past the row's last tile: an empty partial - for (int i = threadIdx.x; i < g; i += NT) pm[part * g + i] = -INFINITY, pl[part * g + i] = 0.f; - return; - } - for (int i = threadIdx.x; i < g * d; i += NT) qs[i] = __bfloat162float(q[(r * h + kh * g + i / d) * d + i % d]); - for (int i = threadIdx.x; i < BN * d; i += NT) { - const int key = t * BN + i / d; - const bool in = key >= lo && key <= p; - const int64_t at = ((base + (in ? key : lo) % ring) * hk + kh) * d + i % d; - kt[(i % d) * KT + i / d] = in ? kc[at] : __float2bfloat16_rn(0.0f); // outside the window: zeros - vs[i] = in ? vc[at] : __float2bfloat16_rn(0.0f); - } - __syncthreads(); - for (int i = threadIdx.x; i < g * BN; i += NT) { // score (head i / BN, key i % BN) - const int hd = i / BN, jj = i % BN, key = t * BN + jj; - float s = 0.f; -#pragma unroll 8 - for (int e = 0; e < d; ++e) s = __fmaf_rn(qs[hd * d + e], __bfloat162float(kt[e * KT + jj]), s); - sc[i] = (key >= lo && key <= p) ? __fmul_rn(s, scale) : -INFINITY; - } - __syncthreads(); - { // a warp a head: the tile's max, the weights exp(s - max) and their sum, lanes in a fixed tree - const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; - for (int hd = warp; hd < g; hd += NT / 32) { - float a = sc[hd * BN + lane], b = sc[hd * BN + 32 + lane]; - float m = fmaxf(a, b); - for (int k = 16; k > 0; k >>= 1) m = fmaxf(m, __shfl_xor_sync(0xFFFFFFFFu, m, k)); - a = a == -INFINITY ? 0.f : expf(__fsub_rn(a, m)); - b = b == -INFINITY ? 0.f : expf(__fsub_rn(b, m)); - sc[hd * BN + lane] = a; - sc[hd * BN + 32 + lane] = b; - float l = __fadd_rn(a, b); - for (int k = 16; k > 0; k >>= 1) l = __fadd_rn(l, __shfl_xor_sync(0xFFFFFFFFu, l, k)); - if (lane == 0) mx[hd] = m, ls[hd] = l; - } - } - __syncthreads(); - bool odd = false; - for (int o = threadIdx.x; o < g * d; o += NT) { - const int hd = o / d, e = o % d; - float a = 0.f; - for (int jj = 0; jj < BN; ++jj) a = __fmaf_rn(sc[hd * BN + jj], __bfloat162float(vs[jj * d + e]), a); - po[part * g * d + o] = a; - odd |= !isfinite(a); - } - if (threadIdx.x < g) pm[part * g + threadIdx.x] = mx[threadIdx.x], pl[part * g + threadIdx.x] = ls[threadIdx.x]; - if (__syncthreads_or(odd)) { // a NaN or infinite output: those again (finite outputs keep their bits) - __shared__ double m64[GMAX]; - window_part_odd(q, kc, sc, vs, m64, po + part * g * d, r, kh, g, h, hk, d, t, lo, p, base, ring, scale); - } -} - -// out [W, H, D] bf16: a row's partials folded in tile order, o = sum_j exp(m_j - M) o_j / sum_j exp(m_j - M) l_j. -extern "C" __global__ void __launch_bounds__(NT) tf_window_merge(const float* __restrict__ po, - const float* __restrict__ pm, - const float* __restrict__ pl, - __nv_bfloat16* __restrict__ out, int h, int hk, int d, - int tiles) { - const int64_t r = blockIdx.x; - const int kh = blockIdx.y, g = h / hk; - const int64_t part0 = (r * hk + kh) * tiles; - for (int o = threadIdx.x; o < g * d; o += NT) { - const int hd = o / d; - float m = -INFINITY; - for (int j = 0; j < tiles; ++j) m = fmaxf(m, pm[(part0 + j) * g + hd]); - float acc = 0.f, l = 0.f; - for (int j = 0; j < tiles; ++j) { - const float mj = pm[(part0 + j) * g + hd]; - if (mj == -INFINITY) continue; - const float a = expf(__fsub_rn(mj, m)); - acc = __fmaf_rn(a, po[(part0 + j) * g * d + o], acc); - l = __fmaf_rn(a, pl[(part0 + j) * g + hd], l); - } - if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a factor of 0; again, kept - acc = 0.f; - for (int j = 0; j < tiles; ++j) { - const float mj = pm[(part0 + j) * g + hd]; - if (mj != -INFINITY) acc = attn_odd::fold(po[(part0 + j) * g * d + o], expf(__fsub_rn(mj, m)), acc); - } - } - out[(r * h + kh * g + hd) * d + o % d] = __float2bfloat16_rn(__fdiv_rn(acc, l)); - } -} +#include "window_split.cuh" // Prompt rows: prefill_attention.cu's tensor-core kernel with a window (keys (pos - win, pos]) and a ring of keys. namespace tf_window_prompt { @@ -432,9 +277,6 @@ wattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restric } // Decode on tensor cores (d 128): a row's heads one mma tile; warps take every 4th 32-key tile, fold in warp order. -// ODD: every NaN and infinite value leaves the mma, and the heads add them back at float64's weights. -// !ODD: true when some output is NaN, the one wrong class a 0 x inf can make (an infinite output holds a weight -// above float32's least, so above float64's): the block (one warp: its warp) runs again as ODD over it. template __device__ __forceinline__ bool decode_tc_rows( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc, @@ -671,7 +513,7 @@ __device__ __forceinline__ bool decode_tc_rows( return !ODD && __syncthreads_or(nan); } -// A block (or the one warp that saw keys) with a NaN output, run again with ODD on (out of the common path). +// A NaN output, the block (or its one warp with keys) again: NaN and inf out of the mma, added at float64's weights. __device__ __noinline__ void decode_tc_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, __nv_bfloat16* out, const int* pos, const int* slot, int h, int hk, int win, long long ring, float scale) { diff --git a/zig/kernels/cuda/window_split.cuh b/zig/kernels/cuda/window_split.cuh new file mode 100644 index 0000000000..0c7cd8c0db --- /dev/null +++ b/zig/kernels/cuda/window_split.cuh @@ -0,0 +1,164 @@ +// Sliding-window decode split by key tile (float64 references: core/attention_ref.zig), and its merge. +#pragma once + +#include +#include +#include +#include + +#include "attn_odd.cuh" + +// Decode: a block a (row, key head); keys walk 64-key tiles by absolute position, so a row's bits are its own. + +namespace tf_window { + +constexpr int NT = 256; +constexpr int BN = 64; // keys a tile +constexpr int GMAX = 16; // query heads a key head +constexpr int KT = BN + 2; // kt's row stride: a bank further each dim +constexpr int DMAX = 128; // head dim (static shared memory: q, a K and a V tile, scores within 48 KiB) + +} // namespace tf_window + +using namespace tf_window; + +// tf_window_part's NaN or infinite outputs again: the finite values' sum, then each odd value at its float64 weight. +__device__ __noinline__ void window_part_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const float* sc, + const __nv_bfloat16* vs, double* m64, float* po, int64_t r, int kh, + int g, int h, int hk, int d, int t, int lo, int p, int64_t base, + long long ring, float scale) { + const auto krow = [&](int key) { return kc + ((base + key % ring) * hk + kh) * d; }; + const auto qrow = [&](int hd) { return q + (r * h + kh * g + hd) * d; }; + { // each head's float64 max over the row's window, attention_ref's scores: a warp a head, lanes over keys + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + for (int hd = warp; hd < g; hd += NT / 32) { + double mx = -INFINITY; + for (int key = lo + lane; key <= p; key += 32) mx = fmax(mx, attn_odd::score64(qrow(hd), krow(key), d, scale)); + for (int k = 16; k > 0; k >>= 1) mx = fmax(mx, __shfl_xor_sync(0xFFFFFFFFu, mx, k)); + if (lane == 0) m64[hd] = mx; + } + } + __syncthreads(); + for (int o = threadIdx.x; o < g * d; o += NT) { + if (isfinite(po[o])) continue; + const int hd = o / d, e = o % d; + float a = 0.f; + for (int jj = 0; jj < BN; ++jj) { + const float v = __bfloat162float(vs[jj * d + e]); + if (isfinite(v)) a = __fmaf_rn(sc[hd * BN + jj], v, a); + } + for (int jj = 0; jj < BN; ++jj) { + const float v = __bfloat162float(vs[jj * d + e]); + const int key = t * BN + jj; + if (isfinite(v) || key < lo || key > p) continue; + const bool w = __dsub_rn(attn_odd::score64(qrow(hd), krow(key), d, scale), m64[hd]) >= attn_odd::least(); + a = __fadd_rn(a, w ? v : __fmul_rn(0.f, v)); // NaN score: no weight + } + po[o] = a; + } +} + +// Decode split by key tile (block: row, key head, tile); tf_window_merge folds the partials in tile order. +extern "C" __global__ void __launch_bounds__(NT) tf_window_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 d, int win, long long ring, float scale) { + __shared__ float qs[GMAX * DMAX]; + __shared__ __nv_bfloat16 kt[DMAX * KT]; // K transposed, [d][key], padded: a warp reads 32 keys of one dim + __shared__ __nv_bfloat16 vs[BN * DMAX]; + __shared__ float sc[GMAX * BN]; + __shared__ float mx[GMAX], ls[GMAX]; + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, tiles = gridDim.z, j = blockIdx.z; + const int g = h / hk; + 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 + 1, 0) : 0; + const int t = lo / BN + j; + const int64_t part = (r * hk + kh) * tiles + j; // this block's partial, g heads of d values + if (t > p / BN) { // past the row's last tile: an empty partial + for (int i = threadIdx.x; i < g; i += NT) pm[part * g + i] = -INFINITY, pl[part * g + i] = 0.f; + return; + } + for (int i = threadIdx.x; i < g * d; i += NT) qs[i] = __bfloat162float(q[(r * h + kh * g + i / d) * d + i % d]); + for (int i = threadIdx.x; i < BN * d; i += NT) { + const int key = t * BN + i / d; + const bool in = key >= lo && key <= p; + const int64_t at = ((base + (in ? key : lo) % ring) * hk + kh) * d + i % d; + kt[(i % d) * KT + i / d] = in ? kc[at] : __float2bfloat16_rn(0.0f); // outside the window: zeros + vs[i] = in ? vc[at] : __float2bfloat16_rn(0.0f); + } + __syncthreads(); + for (int i = threadIdx.x; i < g * BN; i += NT) { // score (head i / BN, key i % BN) + const int hd = i / BN, jj = i % BN, key = t * BN + jj; + float s = 0.f; +#pragma unroll 8 + for (int e = 0; e < d; ++e) s = __fmaf_rn(qs[hd * d + e], __bfloat162float(kt[e * KT + jj]), s); + sc[i] = (key >= lo && key <= p) ? __fmul_rn(s, scale) : -INFINITY; + } + __syncthreads(); + { // a warp a head: the tile's max, the weights exp(s - max) and their sum, lanes in a fixed tree + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + for (int hd = warp; hd < g; hd += NT / 32) { + float a = sc[hd * BN + lane], b = sc[hd * BN + 32 + lane]; + float m = fmaxf(a, b); + for (int k = 16; k > 0; k >>= 1) m = fmaxf(m, __shfl_xor_sync(0xFFFFFFFFu, m, k)); + a = a == -INFINITY ? 0.f : expf(__fsub_rn(a, m)); + b = b == -INFINITY ? 0.f : expf(__fsub_rn(b, m)); + sc[hd * BN + lane] = a; + sc[hd * BN + 32 + lane] = b; + float l = __fadd_rn(a, b); + for (int k = 16; k > 0; k >>= 1) l = __fadd_rn(l, __shfl_xor_sync(0xFFFFFFFFu, l, k)); + if (lane == 0) mx[hd] = m, ls[hd] = l; + } + } + __syncthreads(); + bool odd = false; + for (int o = threadIdx.x; o < g * d; o += NT) { + const int hd = o / d, e = o % d; + float a = 0.f; + for (int jj = 0; jj < BN; ++jj) a = __fmaf_rn(sc[hd * BN + jj], __bfloat162float(vs[jj * d + e]), a); + po[part * g * d + o] = a; + odd |= !isfinite(a); + } + if (threadIdx.x < g) pm[part * g + threadIdx.x] = mx[threadIdx.x], pl[part * g + threadIdx.x] = ls[threadIdx.x]; + if (__syncthreads_or(odd)) { // a NaN or infinite output: those again (finite outputs keep their bits) + __shared__ double m64[GMAX]; + window_part_odd(q, kc, sc, vs, m64, po + part * g * d, r, kh, g, h, hk, d, t, lo, p, base, ring, scale); + } +} + +// out [W, H, D] bf16: a row's partials folded in tile order, o = sum_j exp(m_j - M) o_j / sum_j exp(m_j - M) l_j. +extern "C" __global__ void __launch_bounds__(NT) tf_window_merge(const float* __restrict__ po, + const float* __restrict__ pm, + const float* __restrict__ pl, + __nv_bfloat16* __restrict__ out, int h, int hk, int d, + int tiles) { + const int64_t r = blockIdx.x; + const int kh = blockIdx.y, g = h / hk; + const int64_t part0 = (r * hk + kh) * tiles; + for (int o = threadIdx.x; o < g * d; o += NT) { + const int hd = o / d; + float m = -INFINITY; + for (int j = 0; j < tiles; ++j) m = fmaxf(m, pm[(part0 + j) * g + hd]); + float acc = 0.f, l = 0.f; + for (int j = 0; j < tiles; ++j) { + const float mj = pm[(part0 + j) * g + hd]; + if (mj == -INFINITY) continue; + const float a = expf(__fsub_rn(mj, m)); + acc = __fmaf_rn(a, po[(part0 + j) * g * d + o], acc); + l = __fmaf_rn(a, pl[(part0 + j) * g + hd], l); + } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a factor of 0; again, kept + acc = 0.f; + for (int j = 0; j < tiles; ++j) { + const float mj = pm[(part0 + j) * g + hd]; + if (mj != -INFINITY) acc = attn_odd::fold(po[(part0 + j) * g * d + o], expf(__fsub_rn(mj, m)), acc); + } + } + out[(r * h + kh * g + hd) * d + o % d] = __float2bfloat16_rn(__fdiv_rn(acc, l)); + } +} diff --git a/zig/tests/cuda/contract_window.zig b/zig/tests/cuda/contract_window.zig index d9cf2e0863..a69b91c818 100644 --- a/zig/tests/cuda/contract_window.zig +++ b/zig/tests/cuda/contract_window.zig @@ -134,8 +134,7 @@ pub const Window = struct { return &x[(self.cacheRow(w.base, key) * self.f.hk + kh) * self.f.d + e]; } - /// NaN in a query; +inf and -inf in values; both in one column; NaN in a key; +inf in a window's first value; - /// +inf at float64 weights of 0 and below float32's least. + /// NaN query; +inf and -inf values, both in one column; NaN key; +inf at a window's first key and at tiny weights. fn plant(self: *Window) void { const n = self.shape.rows; const nan = c.toBf16(std.math.nan(f32)); @@ -160,8 +159,7 @@ pub const Window = struct { if (!self.f.prompt and n > 7) self.split(if (self.f.win == 0) 4 else n - 1, if (hk > 2) 2 else 0, 21); } - /// Row r, key head kh: key p all 8, the window's other keys all -8 with +inf at e; kh's first head queries all 2 - /// (weights below float32's least: inf), its second all 8 (float64 weights of 0: NaN). Tiles split the keys. + /// Row r, key head kh: key p all 8, its other keys -8 with +inf at e; queries 2 (inf) and 8 (weight 0: NaN). fn split(self: *Window, r: usize, kh: usize, e: usize) void { const f = self.f; const p = self.at(r).p;