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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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 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 07/15] 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 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 08/15] 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 09/15] 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 10/15] 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 11/15] 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 12/15] 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 13/15] 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 14/15] 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 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 15/15] 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 } } },