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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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 90aebfcb1bcdfef21feaa4760a84d8b5282d99b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 07/14] cuda/tree_attention: each row's own max flag, so a NaN head leaves other rows' keys --- zig/kernels/cuda/tree_attention.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index eea5d2645d..4e50f442e9 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -190,10 +190,10 @@ __device__ __forceinline__ void tree_group( tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x)); tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x)); } - const bool act = tm0 != -INFINITY; - const float n0 = act ? fmaxf(m0, tm0) : m0, n1 = act ? fmaxf(m1, tm1) : m1; - const float a0 = act ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; - const float a1 = act ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; + const bool act0 = tm0 != -INFINITY, act1 = tm1 != -INFINITY; // each row's own: a NaN head moves only its row + const float n0 = act0 ? fmaxf(m0, tm0) : m0, n1 = act1 ? fmaxf(m1, tm1) : m1; + const float a0 = act0 ? (m0 == -INFINITY ? 0.0f : texp(__fsub_rn(m0, n0))) : 1.0f; + const float a1 = act1 ? (m1 == -INFINITY ? 0.0f : texp(__fsub_rn(m1, n1))) : 1.0f; float c0 = 0.0f, c1 = 0.0f; #pragma unroll for (int n = 0; n < NT8; ++n) { From fc48c41c42cae20882996ecb48a62ef9ff7e7055 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 08/14] cuda/block_attention: keys a row cannot see add nothing, so an infinite value there stays out of its sum A block whose output holds a NaN or infinity runs again with the values of the tiles some row cannot see all of taken out of the mma; each row adds back the ones it sees. On finite values the first pass is the kernel as it was, byte for byte. --- zig/kernels/cuda/attn_odd.cuh | 63 +++++++++++++++++++++++++++++ zig/kernels/cuda/block_attention.cu | 60 ++++++++++++++++++--------- 2 files changed, 105 insertions(+), 18 deletions(-) create mode 100644 zig/kernels/cuda/attn_odd.cuh diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh new file mode 100644 index 0000000000..596ee3a95e --- /dev/null +++ b/zig/kernels/cuda/attn_odd.cuh @@ -0,0 +1,63 @@ +// NaN and infinite values in a tensor-core attention tile's V: a key a row cannot see adds nothing to that row's sum. +#pragma once + +#include +#include +#include + +namespace attn_odd { + +// Each bf16 half of x that is NaN or an infinity (all-ones exponent) as +0, the finite halves as they are. +__device__ __forceinline__ uint32_t finite2(uint32_t x) { + return x & (((x & 0x7f80u) == 0x7f80u ? 0u : 0xffffu) | ((x & 0x7f800000u) == 0x7f800000u ? 0u : 0xffff0000u)); +} + +// Some NaN or infinite value in a warp's mma rows g (r0: a real row) and g + 8 (r1) of its fold. +template +__device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, bool r1) { + float a = 0.0f, b = 0.0f; // x - x is 0 for a finite x, NaN for NaN and the infinities +#pragma unroll + for (int i = 0; i < NB; ++i) { + a = __fadd_rn(a, __fadd_rn(__fsub_rn(o[i][0], o[i][0]), __fsub_rn(o[i][1], o[i][1]))); + b = __fadd_rn(b, __fadd_rn(__fsub_rn(o[i][2], o[i][2]), __fsub_rn(o[i][3], o[i][3]))); + } + return (r0 && a != a) || (r1 && b != b); +} + +// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. +template +__device__ __forceinline__ void clean(uint4* vs, int threads) { + for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) { + uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7)); + const uint4 x = *p; + *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w)); + } +} + +// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them. +template +__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0, + float l1) { + const int g = lane / 4, qd = lane % 4; + for (int k = k0; k < k1; ++k) { + const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN + if (!v0 && !v1) continue; + const uint32_t* v = reinterpret_cast(vrow(k)); +#pragma unroll + for (int i = 0; i < D / 8; ++i) { + const uint32_t x = v[4 * i + qd]; + if (finite2(x) == x) continue; + const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u); + if ((x & 0x7f80u) == 0x7f80u) { + if (v0) o[i][0] = __fadd_rn(o[i][0], lo); + if (v1) o[i][2] = __fadd_rn(o[i][2], lo); + } + if ((x & 0x7f800000u) == 0x7f800000u) { + if (v0) o[i][1] = __fadd_rn(o[i][1], hi); + if (v1) o[i][3] = __fadd_rn(o[i][3], hi); + } + } + } +} + +} // namespace attn_odd diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index e3cd06b121..f5b1e2b18a 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -7,6 +7,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_block { @@ -166,30 +167,53 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( } Fold f; const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) - for (int k0 = lo / KT * KT; k0 < s; k0 += KT) { + const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; + // the context's edge tiles, where some row cannot see some key: e0..e1 + const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT); + auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below) + for (int k0 = e0; k0 < s; k0 += KT) { + for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { + const int key = k0 + i / C; + const bool in = key < s; + cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + } + commit(); + wait_all(); + if (odd && k0 < e1) attn_odd::clean(vs, blockDim.x); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, k0, vis); + __syncthreads(); + } + }; + auto block_keys = [&] { // the block's own keys for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { - const int key = k0 + i / C; - const bool in = key < s; - cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); - cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in); + const int c = i / C; + const bool in = c < L; + const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; + cp16z(ks + sw(c, i % C), kb + at, in); + cp16z(vs + sw(c, i % C), vb + at, in); } commit(); wait_all(); + }; + const auto own = [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }; + context(false); + block_keys(); + __syncthreads(); + f.tile(qa, ks, vs, scale, lane, 0, own); + // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(f.o, g8 < L, g8 + 8 < L))) { + f = Fold(); + context(true); + attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1); + block_keys(); + if (causal) attn_odd::clean(vs, blockDim.x); // row r cannot see keys past r __syncthreads(); - f.tile(qa, ks, vs, scale, lane, k0, [&](int r, int key) { return r < L && key < s && s + r - key <= window; }); - __syncthreads(); - } - for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { // the block's own keys - const int c = i / C; - const bool in = c < L; - const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8; - cp16z(ks + sw(c, i % C), kb + at, in); - cp16z(vs + sw(c, i % C), vb + at, in); + f.tile(qa, ks, vs, scale, lane, 0, own); + const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; + if (causal) attn_odd::add_back(f.o, row, 0, L, lane, own, f.l0, f.l1); } - commit(); - wait_all(); - __syncthreads(); - f.tile(qa, ks, vs, scale, lane, 0, [&](int r, int c) { return r < L && c < L && (!causal || c <= r); }); if (g8 < L) #pragma unroll for (int i = 0; i < NB; ++i) From ee81000d8ace5925c4dca89cfb3d75eff92d8f1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:51:57 +0200 Subject: [PATCH 09/14] 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 f4de6c3bf786079428a014c5b267d9ea70b05fd7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 01:37:58 +0200 Subject: [PATCH 10/14] tests/cuda/contract_attn: the reference's constructor frees what it took when a later allocation fails, with a host test that fails each allocation in turn --- zig/build/cuda.zig | 3 +++ zig/tests/cuda/contract_attn.zig | 22 +++++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 4ca82f73c0..b36accc65b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -271,6 +271,9 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + // the contract's attention reference (tests/cuda/contract_attn.zig): its allocations under failure + const attn_ref = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_attn.zig"), .target = host, .optimize = .debug, .imports = &.{.{ .name = "core", .module = mods.core }} }); + step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = attn_ref })).step); // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); diff --git a/zig/tests/cuda/contract_attn.zig b/zig/tests/cuda/contract_attn.zig index cec73374a3..afa284d55c 100644 --- a/zig/tests/cuda/contract_attn.zig +++ b/zig/tests/cuda/contract_attn.zig @@ -27,8 +27,19 @@ pub const Ref = struct { pub fn init(a: std.mem.Allocator, d: usize, most: usize) !Ref { const keys = try a.alloc(usize, most); + errdefer a.free(keys); for (keys, 0..) |*k, i| k.* = i; - return .{ .a = a, .d = d, .head = .{ .q = try a.alloc(f64, d), .k = try a.alloc(f64, most * d), .v = try a.alloc(f64, most * d) }, .keys = keys, .p = try a.alloc(f64, most), .want = try a.alloc(f64, d), .tol = try a.alloc(f64, d) }; + const q = try a.alloc(f64, d); + errdefer a.free(q); + const k = try a.alloc(f64, most * d); + errdefer a.free(k); + const v = try a.alloc(f64, most * d); + errdefer a.free(v); + const p = try a.alloc(f64, most); + errdefer a.free(p); + const want = try a.alloc(f64, d); + errdefer a.free(want); + return .{ .a = a, .d = d, .head = .{ .q = q, .k = k, .v = v }, .keys = keys, .p = p, .want = want, .tol = try a.alloc(f64, d) }; } pub fn deinit(r: *Ref) void { @@ -122,3 +133,12 @@ pub fn mirrored(rng: std.Random, out: []f32) void { pub fn normal(rng: std.Random) u16 { return c.toBf16(rng.floatNorm(f32)); } + +fn initAndFree(a: std.mem.Allocator) !void { + var r = try Ref.init(a, 8, 5); + r.deinit(); +} + +test "Ref.init frees what it took when a later allocation fails" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, initAndFree, .{}); +} From ff6e5c0ea1b1a3933ca7026ff341ef62be62b235 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:01:32 +0200 Subject: [PATCH 11/14] cuda/block_attention: the replay takes every NaN and infinite value out of the mma and adds it back at its key's float64 weight, so an infinite value gives NaN where float64 does --- zig/kernels/cuda/attn_odd.cuh | 63 ++++++++++++++++++++++---- zig/kernels/cuda/block_attention.cu | 25 ++++++---- zig/src/core/contract/registry.zig | 2 +- zig/tests/cuda/contract_block_attn.zig | 30 ++++++++++++ 4 files changed, 103 insertions(+), 17 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 596ee3a95e..7501de7b7c 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -3,6 +3,7 @@ #include #include +#include #include namespace attn_odd { @@ -34,27 +35,73 @@ __device__ __forceinline__ void clean(uint4* vs, int threads) { } } -// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them. -template -__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0, - float l1) { +// The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. +__device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } + +// A query row's float64 score against a key row, attention_ref's order: sum q_e k_e over e, then times scale. +template +__device__ __forceinline__ double score64(const __nv_bfloat16* q, const __nv_bfloat16* k, double scale) { + double s = 0.0; +#pragma unroll 1 + for (int e = 0; e < D; ++e) + s = __dadd_rn(s, __dmul_rn((double)__bfloat162float(q[e]), (double)__bfloat162float(k[e]))); + return __dmul_rn(s, scale); +} + +// Rows g, g + 8: w(r, K row) is attention_ref's exp(score - max) > 0; each(r, f) walks r's K rows (max on first use). +template +struct Weights { + Q q; + Each each; + double scale; + int g; + double m[2] = {-INFINITY, -INFINITY}; + bool have[2] = {false, false}; + + __device__ __forceinline__ bool operator()(int r, const __nv_bfloat16* k) { + const int i = r == g ? 0 : 1; + if (!have[i]) { + double mx = -INFINITY; + each(r, [&](const __nv_bfloat16* kr) { mx = fmax(mx, score64(q(r), kr, scale)); }); + m[i] = mx; + have[i] = true; + } + return __dsub_rn(score64(q(r), k, scale), m[i]) >= least(); // NaN: no weight + } +}; + +// q(r): row r's query. +template +__device__ __forceinline__ Weights weights(Q q, Each each, float scale, int lane) { + return Weights{q, each, (double)scale, lane / 4}; +} + +// Keys k0..k1 clean() emptied (V vrow(k), K krow(k)): NaN and infinities to rows that see them, inf x 0 at weight 0. +template +__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, KRow krow, int k0, int k1, int lane, Vis vis, + float l0, float l1, W& w) { const int g = lane / 4, qd = lane % 4; for (int k = k0; k < k1; ++k) { const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN if (!v0 && !v1) continue; const uint32_t* v = reinterpret_cast(vrow(k)); + bool odd = false; +#pragma unroll + for (int i = 0; i < D / 8; ++i) odd |= finite2(v[4 * i + qd]) != v[4 * i + qd]; + if (!odd) continue; + const bool w0 = v0 && w(g, krow(k)), w1 = v1 && w(g + 8, krow(k)); #pragma unroll for (int i = 0; i < D / 8; ++i) { const uint32_t x = v[4 * i + qd]; if (finite2(x) == x) continue; const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u); if ((x & 0x7f80u) == 0x7f80u) { - if (v0) o[i][0] = __fadd_rn(o[i][0], lo); - if (v1) o[i][2] = __fadd_rn(o[i][2], lo); + if (v0) o[i][0] = __fadd_rn(o[i][0], w0 ? lo : __fmul_rn(0.0f, lo)); + if (v1) o[i][2] = __fadd_rn(o[i][2], w1 ? lo : __fmul_rn(0.0f, lo)); } if ((x & 0x7f800000u) == 0x7f800000u) { - if (v0) o[i][1] = __fadd_rn(o[i][1], hi); - if (v1) o[i][3] = __fadd_rn(o[i][3], hi); + if (v0) o[i][1] = __fadd_rn(o[i][1], w0 ? hi : __fmul_rn(0.0f, hi)); + if (v1) o[i][3] = __fadd_rn(o[i][3], w1 ? hi : __fmul_rn(0.0f, hi)); } } } diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index f5b1e2b18a..84cd949888 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -168,9 +168,8 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( Fold f; const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window) const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; - // the context's edge tiles, where some row cannot see some key: e0..e1 - const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT); - auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below) + const int e0 = lo / KT * KT; // the first context tile + auto context = [&](bool odd) { // odd: the tiles' NaN and infinite values leave the mma (added back below) for (int k0 = e0; k0 < s; k0 += KT) { for (int i = threadIdx.x; i < KT * C; i += blockDim.x) { const int key = k0 + i / C; @@ -180,7 +179,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( } commit(); wait_all(); - if (odd && k0 < e1) attn_odd::clean(vs, blockDim.x); + if (odd) attn_odd::clean(vs, blockDim.x); __syncthreads(); f.tile(qa, ks, vs, scale, lane, k0, vis); __syncthreads(); @@ -202,17 +201,27 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( block_keys(); __syncthreads(); f.tile(qa, ks, vs, scale, lane, 0, own); - // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits) + // a NaN or infinite output: again, NaN and infinite values out of the mma, added back (finite values: same bits) if (__syncthreads_or(attn_odd::nonfinite_out(f.o, g8 < L, g8 + 8 < L))) { f = Fold(); context(true); - attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1); + const auto ck = [&](int k) { return kc + (int64_t)k * D; }; + const auto bk = [&](int c) { return kb + ((int64_t)g * rows + j * L + c) * D; }; + const auto qrow = [&](int r) { return q + ((int64_t)head * rows + j * L + r) * D; }; + const auto each = [&](int r, auto f) { // row r's keys: its window of the context, then the block's + for (int k = lo; k < s; ++k) + if (vis(r, k)) f(ck(k)); + for (int c = 0; c < L; ++c) + if (own(r, c)) f(bk(c)); + }; + auto w = attn_odd::weights(qrow, each, scale, lane); + attn_odd::add_back(f.o, [&](int k) { return vc + (int64_t)k * D; }, ck, e0, s, lane, vis, f.l0, f.l1, w); block_keys(); - if (causal) attn_odd::clean(vs, blockDim.x); // row r cannot see keys past r + attn_odd::clean(vs, blockDim.x); __syncthreads(); f.tile(qa, ks, vs, scale, lane, 0, own); const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; - if (causal) attn_odd::add_back(f.o, row, 0, L, lane, own, f.l0, f.l1); + attn_odd::add_back(f.o, row, bk, 0, L, lane, own, f.l0, f.l1, w); } if (g8 < L) #pragma unroll diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b2cfbc0e18..3fe36e2e50 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -59,7 +59,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 220, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 }, .{ .regs = 233, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 208, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 }, .{ .regs = 231, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 97, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 215, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 }, .{ .regs = 178, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 155374b86d..4d599951f4 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -140,6 +140,36 @@ pub const Block = struct { } if (s0 > window) v0[(4 * s0 + s0 - window) * D + 13] = inf; // a key block row 0 alone sees self.ctx.items[self.ns - 1][(3 * self.ctxLen(self.ns - 1) + self.ctxLen(self.ns - 1) - 1) * D] = nan; + // +inf values row 0 sees at a float64 weight of 0 (0 x inf: NaN), and at weights below float32's least (inf) + const k0 = self.ctx.items[0][0 .. hk * s0 * D]; + self.reach(k0, v0, 6, 24, null, 19); // the block's first key + if (s0 > window) self.reach(k0, v0, 5, 20, s0 - window + 8, 17); // the window's edge tile, rows 0 to 8 see it + self.under(k0, v0); + } + + /// Head `head`'s row 0: query all 8, key s0 - 1 all 8, key `at` (null: block key 0) all -8 with +inf at e. + fn reach(self: *Block, k0: []u16, v0: []u16, kh: usize, head: usize, at: ?usize, e: usize) void { + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + @memset(self.q.items[(head * rows) * D ..][0..D], c.toBf16(8)); + @memset(k0[(kh * s0 + s0 - 1) * D ..][0..D], c.toBf16(8)); + const far = if (at) |k| k0[(kh * s0 + k) * D ..][0..D] else self.kb.items[(kh * rows) * D ..][0..D]; + @memset(far, c.toBf16(-8)); + if (at) |k| v0[(kh * s0 + k) * D + e] = c.toBf16(std.math.inf(f32)) else self.vb.items[(kh * rows) * D + e] = c.toBf16(std.math.inf(f32)); + } + + /// Key head 7, stream 0: head 28's row 0 query all 2, key s0 - 1 all 8, every other key all -8 with +inf at 21. + fn under(self: *Block, k0: []u16, v0: []u16) void { + const rows = self.ns * self.len(); + const s0 = self.ctxLen(0); + const kh = 7; + @memset(self.q.items[(28 * rows) * D ..][0..D], c.toBf16(2)); + for (0..s0 + self.len()) |k| { + const key = if (k < s0) k0[(kh * s0 + k) * D ..][0..D] else self.kb.items[(kh * rows + k - s0) * D ..][0..D]; + const val = if (k < s0) v0[(kh * s0 + k) * D ..][0..D] else self.vb.items[(kh * rows + k - s0) * D ..][0..D]; + @memset(key, c.toBf16(if (k == s0 - 1) 8 else -8)); + if (k != s0 - 1) val[21] = c.toBf16(std.math.inf(f32)); + } } /// Gate 4: every key alike (each p is 1); values mirrored over the context keys every row sees, else 0. From e0726dc18589cf8e35bfc3388d655fc24b56928d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:37:54 +0200 Subject: [PATCH 12/14] cuda/block_attention: the replay's block tile keeps the infinities added back before it The replay adds the context's NaN and infinite values back before the block's own tile, whose rescale is exp(m - m') in float32. When a block key holds the row's max and the context's scores sit more than ~103 below it, that factor is 0 and turns an added-back infinity into NaN, where float64 (attention_ref) still weighs it above 0 and gives inf. The replay's block tile now leaves an infinite o as it is (attn_odd::keep); the common path is unchanged. Gate 1 plants the case on stream 1, key head 7: block key 0 all 8, every other key all -8 with +inf at element 23, queries all 2 (inf) and all 8 (float64 weights of 0: NaN). --- zig/kernels/cuda/attn_odd.cuh | 3 +++ zig/kernels/cuda/block_attention.cu | 13 +++++++------ zig/tests/cuda/contract_block_attn.zig | 22 ++++++++++++++++++++++ 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 7501de7b7c..1a72f5534c 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -35,6 +35,9 @@ __device__ __forceinline__ void clean(uint4* vs, int threads) { } } +// x a, an infinite x (the replay's, already weighed in float64) left as it is whatever the rescale. +__device__ __forceinline__ float keep(float x, float a) { return isinf(x) ? x : __fmul_rn(x, a); } + // The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. __device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index 84cd949888..8af5fa4852 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -42,6 +42,7 @@ __device__ __forceinline__ uint32_t pack(float lo, float hi) { __device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); } // One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V. +// KEEP: the replay's infinities already in o stay (attn_odd::keep), whatever the rescale. struct Fold { float o[NB][4]; float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f; @@ -51,7 +52,7 @@ struct Fold { for (int i = 0; i < NB; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f; } - template + template __device__ __forceinline__ void tile(const uint32_t (&qa)[K16][4], const uint4* ks, const uint4* vs, float scale, int lane, int k0, Vis vis) { const int g8 = lane / 4, qd = lane % 4, mt = lane / 8; @@ -109,10 +110,10 @@ struct Fold { m1 = n1; #pragma unroll for (int i = 0; i < NB; ++i) { - o[i][0] = __fmul_rn(o[i][0], a0); - o[i][1] = __fmul_rn(o[i][1], a0); - o[i][2] = __fmul_rn(o[i][2], a1); - o[i][3] = __fmul_rn(o[i][3], a1); + o[i][0] = KEEP ? attn_odd::keep(o[i][0], a0) : __fmul_rn(o[i][0], a0); + o[i][1] = KEEP ? attn_odd::keep(o[i][1], a0) : __fmul_rn(o[i][1], a0); + o[i][2] = KEEP ? attn_odd::keep(o[i][2], a1) : __fmul_rn(o[i][2], a1); + o[i][3] = KEEP ? attn_odd::keep(o[i][3], a1) : __fmul_rn(o[i][3], a1); } uint32_t pa[2][4]; #pragma unroll @@ -219,7 +220,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_block_attention( block_keys(); attn_odd::clean(vs, blockDim.x); __syncthreads(); - f.tile(qa, ks, vs, scale, lane, 0, own); + f.tile(qa, ks, vs, scale, lane, 0, own); // its rescale (0 when it holds the max) keeps the infinities const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; }; attn_odd::add_back(f.o, row, bk, 0, L, lane, own, f.l0, f.l1, w); } diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 4d599951f4..1197339202 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -145,6 +145,28 @@ pub const Block = struct { self.reach(k0, v0, 6, 24, null, 19); // the block's first key if (s0 > window) self.reach(k0, v0, 5, 20, s0 - window + 8, 17); // the window's edge tile, rows 0 to 8 see it self.under(k0, v0); + if (self.ns > 1) self.ownMax(1, 7, 23); + } + + /// Stream j, key head kh: block key 0 all 8, every other key all -8 with +inf at e, the block's row 0 querying all 2 + /// (head 4 kh: weights below float32's least, inf) and all 8 (head 4 kh + 1: float64 weights of 0, NaN). + fn ownMax(self: *Block, j: usize, kh: usize, e: usize) void { + const L = self.len(); + const rows = self.ns * L; + const s = self.ctxLen(j); + const x = self.ctx.items[j]; + const inf = c.toBf16(std.math.inf(f32)); + const G = h / hk; + @memset(self.q.items[((kh * G) * rows + j * L) * D ..][0..D], c.toBf16(2)); + @memset(self.q.items[((kh * G + 1) * rows + j * L) * D ..][0..D], c.toBf16(8)); + for (0..s) |k| { + @memset(x[(kh * s + k) * D ..][0..D], c.toBf16(-8)); + x[hk * s * D + (kh * s + k) * D + e] = inf; + } + for (0..L) |col| { + @memset(self.kb.items[(kh * rows + j * L + col) * D ..][0..D], c.toBf16(if (col == 0) 8 else -8)); + if (col != 0) self.vb.items[(kh * rows + j * L + col) * D + e] = inf; + } } /// Head `head`'s row 0: query all 8, key s0 - 1 all 8, key `at` (null: block key 0) all -8 with +inf at e. From 942e2fe614bfed34b2c3ee5e8d53fbea2f0cc1d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:38:04 +0200 Subject: [PATCH 13/14] cuda/tree_attention: a NaN or infinite partial runs again with those values out of the mma, added back at float64's weights A visible infinite value whose weight underflows in float32 (more than ~103 below the max) but not in float64 (745) went through the mma as 0 x inf, or reached a merge at a factor of 0: NaN, where attention_ref gives inf. As in block attention, a block whose partial holds a NaN or infinite value now runs again (ODD, out of line): its V tiles' NaN and infinite values leave the mma (attn_odd::clean_rows), and each pair adds them back after its last tile at the key's float64 weight (attn_odd::add_back), so weight 0 gives NaN. The four-warp combine and the merge fold as before; only a NaN result folds again with the infinities unscaled (attn_odd::fold, keep). tf_tree_shared takes the window's keys (kn) for the path keys' weights. Finite outputs keep their bits. Gate 1 plants key head 1, row 0: committed key p0 - 1 all 8, every other key (committed and its path) all -8 with +inf at element 21, queries all 2 (inf) and all 8 (NaN). --- zig/kernels/cuda/attn_odd.cuh | 19 ++++-- zig/kernels/cuda/tree_attention.cu | 91 +++++++++++++++++++++++---- zig/src/core/contract/registry.zig | 26 ++++---- zig/src/cuda/tree_attention.zig | 1 + zig/tests/cuda/contract_tree_attn.zig | 24 ++++++- 5 files changed, 131 insertions(+), 30 deletions(-) diff --git a/zig/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh index 1a72f5534c..a7c4354f75 100644 --- a/zig/kernels/cuda/attn_odd.cuh +++ b/zig/kernels/cuda/attn_odd.cuh @@ -25,19 +25,30 @@ __device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, return (r0 && a != a) || (r1 && b != b); } -// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. -template -__device__ __forceinline__ void clean(uint4* vs, int threads) { - for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) { +// A KEYS-key tile staged [key][D] (xor swizzle): chunks tid, tid + threads, ... (the ones tid copied) made finite. +template +__device__ __forceinline__ void clean_rows(uint4* vs, int tid, int threads) { + for (int i = tid; i < KEYS * (D / 8); i += threads) { uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7)); const uint4 x = *p; *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w)); } } +// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0. +template +__device__ __forceinline__ void clean(uint4* vs, int threads) { + clean_rows(vs, threadIdx.x, threads); +} + // x a, an infinite x (the replay's, already weighed in float64) left as it is whatever the rescale. __device__ __forceinline__ float keep(float x, float a) { return isinf(x) ? x : __fmul_rn(x, a); } +// x a + y folding partials, an infinite x left unscaled (see keep). +__device__ __forceinline__ float fold(float x, float a, float y) { + return isinf(x) ? __fadd_rn(x, y) : __fmaf_rn(x, a, y); +} + // The least x whose float64 exp is not 0: a key whose score is further below its row's max weighs 0 in attention_ref. __device__ __forceinline__ double least() { return __longlong_as_double(0xc0874910d52d3051ll); } diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 4e50f442e9..2673273642 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -5,6 +5,7 @@ #include #include +#include "attn_odd.cuh" #include "qmm_frag.cuh" namespace tf_tree { @@ -58,8 +59,10 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ } // Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -template -__device__ __forceinline__ void tree_group( +// ODD: every NaN and infinite value leaves the mma, and the pairs add them back at float64's weights. +// !ODD: true (nothing written) when some warp's partial is NaN or infinite: the block runs again as ODD. +template +__device__ __forceinline__ bool tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, @@ -75,7 +78,7 @@ __device__ __forceinline__ void tree_group( const int s = row_stream[r]; const int p = streams[4 * s + 2], depth = depths[r]; const int total = p + depth, k0 = gi * SPAN, k1 = min(total, k0 + SPAN); - if (k0 >= total) return; // the merge reads only the row's own groups + if (k0 >= total) return false; // the merge reads only the row's own groups const __nv_bfloat16* kc = origin + offs[2 * s]; const __nv_bfloat16* vc = origin + offs[2 * s + 1]; // ST tile stages: a one-warp block fetches ST - 1 tiles ahead while it computes one (stages never change bits) @@ -149,7 +152,7 @@ __device__ __forceinline__ void tree_group( ks = kv0 + st * 2 * HB * C; vs = ks + HB * C; st = (st + 1) % ST; - + if constexpr (ODD) attn_odd::clean_rows(vs, lane, 32); // the chunks this lane copied __syncwarp(); float sc[NT8][4]; #pragma unroll @@ -238,6 +241,25 @@ __device__ __forceinline__ void tree_group( } __syncwarp(); } + if constexpr (!ODD) { // a NaN or infinite partial: the block runs again as ODD (finite values: the same bits) + if (__syncthreads_or(attn_odd::nonfinite_out(o, pf + g8 < npairs, pf + g8 + 8 < npairs))) return true; + } else { + // key k of the group's row: committed, then the row's path; pair m: row rq + (pf + m) / G, head (pf + m) % G + const auto at = [&](const __nv_bfloat16* c, const __nv_bfloat16* n, int k) { + return k < p ? c + ((int64_t)k * hk + kh) * D : n + ((int64_t)paths[r * MAXD + (k - p)] * hk + kh) * D; + }; + const auto qrow = [&](int m) { return q + ((rq + (pf + m) / G) * h + kh * G + (pf + m) % G) * D; }; + const auto each = [&](int m, auto f) { // pair m's keys: its stream's committed ones, then its own row's path + const int64_t pr = rq + (pf + m) / G; + for (int k = 0; k < p; ++k) f(kc + ((int64_t)k * hk + kh) * D); + for (int j = 0; j < depths[pr]; ++j) f(kn + ((int64_t)paths[pr * MAXD + j] * hk + kh) * D); + }; + auto w = attn_odd::weights(qrow, each, scale, lane); + const auto vis = [&](int m, int) { return pf + m < npairs; }; // a pair sees every key of its group + for (int t = t0; t * HB < k1; t += 4) // after the last tile: no rescale turns an infinity into NaN + attn_odd::add_back(o, [&](int k) { return at(vc, vn, k); }, [&](int k) { return at(kc, kn, k); }, + t * HB, min(k1, t * HB + HB), lane, vis, l0, l1, w); + } if constexpr (W == 1) { // one warp a block: each slice's partial (o, max, sum) of its 16 pairs; the merge combines the 4 in order #pragma unroll @@ -254,7 +276,7 @@ __device__ __forceinline__ void tree_group( } if (qd == 0) pm[part * G + head] = half ? m1 : m0, pl[part * G + head] = half ? l1 : l0; } - return; + return false; } else { __syncthreads(); float* os = reinterpret_cast(smem); // [4 warps][16 heads][D], then max and sum [4][16] @@ -288,19 +310,37 @@ __device__ __forceinline__ void tree_group( acc = __fmaf_rn(a, os[(w * 16 + hd) * D + e], acc); l = __fmaf_rn(a, ls[w * 16 + hd], l); } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a weight of 0; again, kept + acc = 0.0f; + for (int w = 0; w < 4; ++w) + if (ms[w * 16 + hd] != -INFINITY) + acc = attn_odd::fold(os[(w * 16 + hd) * D + e], texp(__fsub_rn(ms[w * 16 + hd], mm)), acc); + } po[(part * G + head) * D + e] = acc; if (e == 0) pm[part * G + head] = mm, pl[part * G + head] = l; } + return false; } } +// A block whose partial is NaN or infinite, run again as ODD (out of the common path). +template +__device__ __noinline__ void tree_odd( + const __nv_bfloat16* q, const __nv_bfloat16* kn, const __nv_bfloat16* vn, const __nv_bfloat16* origin, + const long long* offs, const int* row_stream, const int* streams, const int* paths, const int* depths, float* po, + float* pm, float* pl, int h, int hk, float scale, int gi, int ng, int64_t rq, int pf, int npairs, int64_t row, + int slice) { + tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, + gi, ng, rq, pf, npairs, row, slice); +} + // A stream's full committed groups once for its rows, 16 (row, head) pairs a block: each partial a row block's bits. template __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, float* __restrict__ pm, float* __restrict__ pl, int h, int hk, - float scale, int ng, int n) { + float scale, int ng, int n, const __nv_bfloat16* __restrict__ kn) { const int G = h / hk; const int per = W == 4 ? 1 : 4; // slices an item takes a block for (int item = blockIdx.x; item < n * ng * per; item += gridDim.x) { // (row, group[, slice]): a fixed grid @@ -311,8 +351,11 @@ __device__ __forceinline__ void tree_shared( const int pairs = streams[4 * st + 1] * G, p = streams[4 * st + 2]; const int tile = static_cast(r - r0); // a stream's j-th row hosts its j-th pair tile (tiles <= rows) if (tile * 16 >= pairs || (gi + 1) * SPAN > p) continue; - tree_group(q, q, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, ng, - r0, tile * 16, pairs, r, item % per); + // vn: never read, a full committed group has no path keys (the replay weighs them, from kn) + if (tree_group(q, kn, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, + scale, gi, ng, r0, tile * 16, pairs, r, item % per)) + tree_odd(q, kn, q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + ng, r0, tile * 16, pairs, r, item % per); __syncthreads(); } } @@ -322,8 +365,10 @@ __device__ __forceinline__ void tree_shared( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ origin, \ const long long* __restrict__ offs, const int* __restrict__ row_stream, const int* __restrict__ streams, \ const int* __restrict__ paths, const int* __restrict__ depths, float* __restrict__ po, \ - float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n) { \ - tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n); \ + float* __restrict__ pm, float* __restrict__ pl, int h, int hk, float scale, int ng, int n, \ + const __nv_bfloat16* __restrict__ kn) { \ + tree_shared(q, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, ng, n, \ + kn); \ } TF_TREE_SHARED(tf_tree_shared, 128, 32, 4, 1) @@ -344,8 +389,10 @@ __device__ __forceinline__ void tree_part( const int p = streams[4 * row_stream[r] + 2], total = p + depths[r]; const int per = W == 4 ? 1 : 4, slot = blockIdx.z / per, slots = gridDim.z / per; for (int gi = slot + p / SPAN; gi < ng && gi * SPAN < total; gi += slots) { // full committed: shared - tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, - ng, r, 0, h / hk, r, blockIdx.z % per); + if (tree_group(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, + scale, gi, ng, r, 0, h / hk, r, blockIdx.z % per)) + tree_odd(q, kn, vn, origin, offs, row_stream, streams, paths, depths, po, pm, pl, h, hk, scale, gi, + ng, r, 0, h / hk, r, blockIdx.z % per); __syncthreads(); // the next group's tiles reuse the partials' shared memory } } @@ -417,6 +464,26 @@ extern "C" __global__ void __launch_bounds__(256) tf_tree_merge( m = next; } } + if (acc != acc) { // NaN: a replayed infinity (weighed in float64) may have met a weight of 0 here; again, kept + float mm = -INFINITY; + acc = 0.0f; + for (int s0 = 0; s0 < slots; s0 += per) { + float cm = -INFINITY, co = 0.0f, cl = 0.0f; + for (int w = 0; w < per; ++w) cm = fmaxf(cm, pm[(base + s0 + w) * G + hd]); + for (int w = 0; w < per; ++w) { + const int64_t q = (base + s0 + w) * G + hd; + if (pm[q] == -INFINITY) continue; + const float aw = texp(__fsub_rn(pm[q], cm)); + co = attn_odd::fold(po[q * d + e], aw, co); + cl = __fmaf_rn(aw, pl[q], cl); + } + if (!(cl > 0.0f)) continue; + const float next = fmaxf(mm, cm); + acc = attn_odd::fold(acc, mm == -INFINITY ? 0.0f : texp(__fsub_rn(mm, next)), + attn_odd::keep(co, texp(__fsub_rn(cm, next)))); + mm = next; + } + } out[(r * h + kh * G + hd) * d + e] = __float2bfloat16_rn(__fdiv_rn(acc, l)); } } diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 3fe36e2e50..6083adcc2c 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -46,19 +46,19 @@ pub const cuda = [_]CudaRecord{ .{ .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 = "tree attention tf_tree_shared", .symbol = "tf_tree_shared", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1", .symbol = "tf_tree_shared_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_w1s3", .symbol = "tf_tree_shared_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part", .symbol = "tf_tree_part", .image = "tree_attention", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1", .symbol = "tf_tree_part_w1", .image = "tree_attention", .threads = 32, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_w1s3", .symbol = "tf_tree_part_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 49152, .opt_in = true }, .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256", .symbol = "tf_tree_shared_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1", .symbol = "tf_tree_shared_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_shared_256_w1s3", .symbol = "tf_tree_shared_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256", .symbol = "tf_tree_part_256", .image = "tree_attention", .threads = 128, .dynamic = 73728, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1", .symbol = "tf_tree_part_256_w1", .image = "tree_attention", .threads = 32, .dynamic = 40960, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_part_256_w1s3", .symbol = "tf_tree_part_256_w1s3", .image = "tree_attention", .threads = 32, .dynamic = 57344, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tree attention tf_tree_merge", .symbol = "tf_tree_merge", .image = "tree_attention", .threads = 256 }, .usage = .{ .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "block attention tf_block_attention (up to 8 query heads a key head)", .symbol = "tf_block_attention", .image = "block_attention", .threads = 256, .dynamic = 16384, .opt_in = true }, .usage = .{ .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 } } }, }; diff --git a/zig/src/cuda/tree_attention.zig b/zig/src/cuda/tree_attention.zig index 69b2ad598f..d4fc29afb2 100644 --- a/zig/src/cuda/tree_attention.zig +++ b/zig/src/cuda/tree_attention.zig @@ -109,6 +109,7 @@ pub const TreeAttention = struct { c.add(scale); c.add(@as(i32, @intCast(ng))); c.add(@as(i32, @intCast(n))); + c.add(kn); // the replay of a NaN or infinite partial weighs a row's path keys too const per: usize = if (one) slices else 1; const block: u32 = if (one) 32 else 128; // three stages where the one-warp grid fits one wave of the SMs (a block each), else two (two blocks an SM) diff --git a/zig/tests/cuda/contract_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig index 3dd57f47e6..466526de7a 100644 --- a/zig/tests/cuda/contract_tree_attn.zig +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -153,7 +153,8 @@ pub const Tree = struct { return true; } - /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key. + /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key; + /// +inf at float64 weights of 0 and below float32's least. fn plant(self: *Tree) void { const d = self.v.d; const kv = self.v.hk * d; @@ -171,6 +172,27 @@ pub const Tree = struct { self.cache.items[vo + (p0 - 3) * kv + 2 * d + 11] = ninf; } self.cache.items[ko + (p0 - 1) * kv + 3 * d + 0] = nan; + self.under(1, 21); + } + + /// Key head kh, row 0: committed key p0 - 1 all 8, its other keys all -8 with +inf at e. Head kh's first query is + /// all 2 (weights below float32's least: inf), its second all 8 (float64 weights of 0, 0 x inf: NaN). + fn under(self: *Tree, kh: usize, e: usize) void { + const d = self.v.d; + const kv = self.v.hk * d; + const G = self.v.h / self.v.hk; + const p0 = self.committed.items[0]; + const ko: usize = @intCast(self.offs.items[0]); + const vo: usize = @intCast(self.offs.items[1]); + const inf = c.toBf16(std.math.inf(f32)); + @memset(self.q.items[(kh * G) * d ..][0..d], c.toBf16(2)); + @memset(self.q.items[(kh * G + 1) * d ..][0..d], c.toBf16(8)); + for (0..p0) |j| { + @memset(self.cache.items[ko + j * kv + kh * d ..][0..d], c.toBf16(if (j == p0 - 1) 8 else -8)); + if (j != p0 - 1) self.cache.items[vo + j * kv + kh * d + e] = inf; + } + @memset(self.kn.items[kh * d ..][0..d], c.toBf16(-8)); // row 0's path: itself + self.vn.items[kh * d + e] = inf; } /// Gate 4: every key alike (each p is 1), committed values mirrored per column, path values 0. From b502747c7fe1a28faf006c1b92e4e341eda524bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 18 Oct 2026 09:25:26 +0200 Subject: [PATCH 14/14] cuda/tree_attention, block_attention: the replay's comments on one line each --- zig/kernels/cuda/block_attention.cu | 3 +-- zig/kernels/cuda/tree_attention.cu | 6 ++---- zig/tests/cuda/contract_block_attn.zig | 3 +-- zig/tests/cuda/contract_tree_attn.zig | 6 ++---- 4 files changed, 6 insertions(+), 12 deletions(-) diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu index 8af5fa4852..b8675f80d0 100644 --- a/zig/kernels/cuda/block_attention.cu +++ b/zig/kernels/cuda/block_attention.cu @@ -41,8 +41,7 @@ __device__ __forceinline__ uint32_t pack(float lo, float hi) { __device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); } -// One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V. -// KEEP: the replay's infinities already in o stay (attn_odd::keep), whatever the rescale. +// One warp's fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax, P x V (KEEP: o's inf stays). struct Fold { float o[NB][4]; float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f; diff --git a/zig/kernels/cuda/tree_attention.cu b/zig/kernels/cuda/tree_attention.cu index 2673273642..5ab4bdf8f8 100644 --- a/zig/kernels/cuda/tree_attention.cu +++ b/zig/kernels/cuda/tree_attention.cu @@ -58,9 +58,7 @@ extern "C" __global__ void tf_tree_paths(const int* __restrict__ parents, int* _ for (int slot = depth - 1; slot >= 0; --slot, cur = parents[cur]) paths[node * MAXD + slot] = cur; } -// Group gi of key head kh for 16 (row, head) pairs from pair pf (pair m: row rq + m / G, head m % G), one mma tile. -// ODD: every NaN and infinite value leaves the mma, and the pairs add them back at float64's weights. -// !ODD: true (nothing written) when some warp's partial is NaN or infinite: the block runs again as ODD. +// Group gi of key head kh for 16 (row, head) pairs from pf (m: row rq + m / G, head m % G); !ODD true: run tree_odd. template __device__ __forceinline__ bool tree_group( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kn, const __nv_bfloat16* __restrict__ vn, @@ -323,7 +321,7 @@ __device__ __forceinline__ bool tree_group( } } -// A block whose partial is NaN or infinite, run again as ODD (out of the common path). +// A block with a NaN or infinite partial, again as ODD: those values out of the mma, added back at float64's weights. template __device__ __noinline__ void tree_odd( const __nv_bfloat16* q, const __nv_bfloat16* kn, const __nv_bfloat16* vn, const __nv_bfloat16* origin, diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig index 1197339202..bc343b9938 100644 --- a/zig/tests/cuda/contract_block_attn.zig +++ b/zig/tests/cuda/contract_block_attn.zig @@ -148,8 +148,7 @@ pub const Block = struct { if (self.ns > 1) self.ownMax(1, 7, 23); } - /// Stream j, key head kh: block key 0 all 8, every other key all -8 with +inf at e, the block's row 0 querying all 2 - /// (head 4 kh: weights below float32's least, inf) and all 8 (head 4 kh + 1: float64 weights of 0, NaN). + /// Stream j, key head kh: block key 0 all 8, other keys -8 with +inf at e; row 0 queries 2 (inf) and 8 (NaN). fn ownMax(self: *Block, j: usize, kh: usize, e: usize) void { const L = self.len(); const rows = self.ns * L; diff --git a/zig/tests/cuda/contract_tree_attn.zig b/zig/tests/cuda/contract_tree_attn.zig index 466526de7a..a81f6493df 100644 --- a/zig/tests/cuda/contract_tree_attn.zig +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -153,8 +153,7 @@ pub const Tree = struct { return true; } - /// NaN in a query, +inf and -inf in values (committed and on a path), both in one column, NaN in a committed key; - /// +inf at float64 weights of 0 and below float32's least. + /// NaN in a query, +inf and -inf in values (committed, on a path, in one column), a NaN key, +inf at tiny weights. fn plant(self: *Tree) void { const d = self.v.d; const kv = self.v.hk * d; @@ -175,8 +174,7 @@ pub const Tree = struct { self.under(1, 21); } - /// Key head kh, row 0: committed key p0 - 1 all 8, its other keys all -8 with +inf at e. Head kh's first query is - /// all 2 (weights below float32's least: inf), its second all 8 (float64 weights of 0, 0 x inf: NaN). + /// Key head kh, row 0: key p0 - 1 all 8, its other keys -8 with +inf at e; queries 2 (inf) and 8 (weight 0: NaN). fn under(self: *Tree, kh: usize, e: usize) void { const d = self.v.d; const kv = self.v.hk * d;