diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..b36accc65b 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 @@ -269,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/kernels/cuda/attn_odd.cuh b/zig/kernels/cuda/attn_odd.cuh new file mode 100644 index 0000000000..a7c4354f75 --- /dev/null +++ b/zig/kernels/cuda/attn_odd.cuh @@ -0,0 +1,124 @@ +// 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 +#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 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); } + +// 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], 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], w0 ? hi : __fmul_rn(0.0f, hi)); + if (v1) o[i][3] = __fadd_rn(o[i][3], w1 ? hi : __fmul_rn(0.0f, hi)); + } + } + } +} + +} // namespace attn_odd diff --git a/zig/kernels/cuda/block_attention.cu b/zig/kernels/cuda/block_attention.cu new file mode 100644 index 0000000000..b8675f80d0 --- /dev/null +++ b/zig/kernels/cuda/block_attention.cu @@ -0,0 +1,236 @@ +// 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 "attn_odd.cuh" +#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 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; + + __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] = 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 + 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) + const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; }; + 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; + 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) 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 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, 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); + 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(); + attn_odd::clean(vs, blockDim.x); + __syncthreads(); + 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); + } + 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..5ab4bdf8f8 --- /dev/null +++ b/zig/kernels/cuda/tree_attention.cu @@ -0,0 +1,487 @@ +// Tree attention (float64 references: core/attention_ref.zig): committed keys, then the row's path, groups in order. + +#include +#include +#include +#include + +#include "attn_odd.cuh" +#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 = 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; + +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 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, + 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, 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]; + 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 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) + 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 * 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; + 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; + // 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; + 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(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(); + }; + // 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; + if constexpr (ODD) attn_odd::clean_rows(vs, lane, 32); // the chunks this lane copied + __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 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) { + 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(); + } + 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 + 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 * 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 * G + head] = half ? m1 : m0, pl[part * G + head] = half ? l1 : l0; + } + return false; + } 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; + 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); + } + 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 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, + 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, 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 + 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; + // 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(); + } +} + +#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, \ + 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) +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 +__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]; + 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 + 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 + } +} + +#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, \ + ng); \ + } + +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). +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, + __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]; + 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 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; + } + } + 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/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/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..6083adcc2c 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 = 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 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. 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..d4fc29afb2 --- /dev/null +++ b/zig/src/cuda/tree_attention.zig @@ -0,0 +1,130 @@ +//! 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 = 512; +pub const max_depth = 128; +const part_smem = [2]u32{ 4 * 2 * 32 * 16 * 16, 4 * 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 = 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 }; + +pub const Scratch = struct { o: u64, m: u64, l: u64, so: u64, sm: u64, sl: u64 }; + +pub const TreeAttention = struct { + mod: Module, + paths_fn: Function, + part_fns: [2]Function, // head dim 128, 256 + shared_fns: [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 (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 { + 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); + 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") }; + } + + 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; + } + + /// 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, 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 * (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 }; + } + + /// 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; + // 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); + 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, 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))); + 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) + 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[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))); + 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/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/contract_attn.zig b/zig/tests/cuda/contract_attn.zig new file mode 100644 index 0000000000..afa284d55c --- /dev/null +++ b/zig/tests/cuda/contract_attn.zig @@ -0,0 +1,144 @@ +//! 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); + errdefer a.free(keys); + for (keys, 0..) |*k, i| k.* = i; + 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 { + 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)); +} + +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, .{}); +} diff --git a/zig/tests/cuda/contract_block_attn.zig b/zig/tests/cuda/contract_block_attn.zig new file mode 100644 index 0000000000..bc343b9938 --- /dev/null +++ b/zig/tests/cuda/contract_block_attn.zig @@ -0,0 +1,275 @@ +//! 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; + // +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); + if (self.ns > 1) self.ownMax(1, 7, 23); + } + + /// 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; + 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. + 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. + 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..a81f6493df --- /dev/null +++ b/zig/tests/cuda/contract_tree_attn.zig @@ -0,0 +1,353 @@ +//! 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, 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; + 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; + self.under(1, 21); + } + + /// 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; + 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. + 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; + } +}; 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..fedfdb134f --- /dev/null +++ b/zig/tests/cuda/tree_tests.zig @@ -0,0 +1,217 @@ +//! 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, 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, h, 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, 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; + defer g.sms = 0; + 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, 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 { + 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, four-warp or one-warp blocks", .{worst}); +}