From 1dd63e39a9180054df02b59a5d86b2d0b9b22c04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:41:51 +0200 Subject: [PATCH 1/6] cuda: act_rows, gdn_prep and qk_rows for every family (SwiGLU, output gates, DFlash2's conv, Gated DeltaNet inputs, QK norm and RoPE), checked against float64 --- zig/build/cuda.zig | 3 + zig/kernels/cuda/act_rows.cu | 142 ++++++++++++++++++ zig/kernels/cuda/gdn_prep.cu | 99 +++++++++++++ zig/kernels/cuda/qk_rows.cu | 194 +++++++++++++++++++++++++ zig/src/core/act_ref.zig | 11 ++ zig/src/core/gdn_prep_ref.zig | 45 ++++++ zig/src/core/qk_ref.zig | 30 ++++ zig/src/core/root.zig | 4 + zig/src/cuda/act_rows.zig | 95 ++++++++++++ zig/src/cuda/gdn_prep.zig | 72 +++++++++ zig/src/cuda/kernels.zig | 3 + zig/src/cuda/qk_rows.zig | 107 ++++++++++++++ zig/src/cuda/root.zig | 6 + zig/tests/cuda/act_rows_tests.zig | 132 +++++++++++++++++ zig/tests/cuda/gdn_prep_tests.zig | 156 ++++++++++++++++++++ zig/tests/cuda/main.zig | 9 ++ zig/tests/cuda/norm_rows_tests.zig | 24 +-- zig/tests/cuda/qk_rows_tests.zig | 226 +++++++++++++++++++++++++++++ 18 files changed, 1346 insertions(+), 12 deletions(-) create mode 100644 zig/kernels/cuda/act_rows.cu create mode 100644 zig/kernels/cuda/gdn_prep.cu create mode 100644 zig/kernels/cuda/qk_rows.cu create mode 100644 zig/src/core/gdn_prep_ref.zig create mode 100644 zig/src/core/qk_ref.zig create mode 100644 zig/src/cuda/act_rows.zig create mode 100644 zig/src/cuda/gdn_prep.zig create mode 100644 zig/src/cuda/qk_rows.zig create mode 100644 zig/tests/cuda/act_rows_tests.zig create mode 100644 zig/tests/cuda/gdn_prep_tests.zig create mode 100644 zig/tests/cuda/qk_rows_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..35a77ebc58 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -45,6 +45,9 @@ const kernels = [_]Kernel{ .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle + .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) + .{ .name = "gdn_prep", .flags = &.{"-O3"} }, // ours: Gated DeltaNet conv, norms, g and beta (cuda/gdn_prep.zig) + .{ .name = "qk_rows", .flags = &.{"-O3"} }, // ours: QK norm, RoPE and the cache write (cuda/qk_rows.zig) .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu new file mode 100644 index 0000000000..515cc6bf67 --- /dev/null +++ b/zig/kernels/cuda/act_rows.cu @@ -0,0 +1,142 @@ +// Activation rows for every CUDA family (float64 references: core/act_ref.zig): SwiGLU, output gates, group sums. + +// A thread 8 values with 16-byte loads; a row's bits depend on that row alone. + +#include +#include + +namespace tf_act { + +constexpr int NT = 256; + +__device__ __forceinline__ void unpack8(const uint4& u, float* f) { + const __nv_bfloat162* h = reinterpret_cast(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 t = __bfloat1622float2(h[i]); + f[2 * i] = t.x; + f[2 * i + 1] = t.y; + } +} + +__device__ __forceinline__ uint4 pack8(const float* f) { + uint4 u; + __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&u); +#pragma unroll + for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(f[2 * i], f[2 * i + 1]); + return u; +} + +__device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } + +// A thread's 8 stored values in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). +__device__ __forceinline__ float group_sum8(const uint4& packed, bool live) { + float f[8]; + unpack8(packed, f); + float s = 0.f; + if (live) +#pragma unroll + for (int i = 0; i < 8; ++i) s = __fadd_rn(s, f[i]); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 4)); + s = __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 2)); + return __fadd_rn(s, __shfl_xor_sync(0xFFFFFFFFu, s, 1)); +} + +} // namespace tf_act + +using namespace tf_act; + +// out[r, i] = bf16(silu(gate[r * g_ld + i]) * up[r * u_ld + i]) for i < n (n % 64 == 0); xs (or null) out's sums. +extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows(const __nv_bfloat16* __restrict__ gate, long long g_ld, + const __nv_bfloat16* __restrict__ up, long long u_ld, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, + int n) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const bool live = 8 * c < n; + float g[8], u[8], o[8]; + if (live) { + unpack8(*reinterpret_cast(gate + r * g_ld + 8 * c), g); + unpack8(*reinterpret_cast(up + r * u_ld + 8 * c), u); + } +#pragma unroll + for (int i = 0; i < 8; ++i) o[i] = live ? __fmul_rn(__fmul_rn(g[i], sigmoid(g[i])), u[i]) : 0.f; + const uint4 op = pack8(o); + if (live) *reinterpret_cast(out + r * n + 8 * c) = op; + if (xs != nullptr) { + const float s = group_sum8(op, live); + if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s; + } +} + +// out[r, h * d + e] = bf16(o[r, h * d + e] * sigmoid(g[r * g_ld + h * g_hs + e])) (d % 64 == 0); xs (or null) its sums. +extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat16* __restrict__ o, + const __nv_bfloat16* __restrict__ g, long long g_ld, + long long g_hs, __nv_bfloat16* __restrict__ out, + float* __restrict__ xs, int heads, int d) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int n = heads * d; + const bool live = 8 * c < n; + const int h = 8 * c / d, e = 8 * c % d; + float ov[8], gv[8], y[8]; + if (live) { + unpack8(*reinterpret_cast(o + r * n + 8 * c), ov); + unpack8(*reinterpret_cast(g + r * g_ld + h * g_hs + e), gv); + } +#pragma unroll + for (int i = 0; i < 8; ++i) y[i] = live ? __fmul_rn(ov[i], sigmoid(gv[i])) : 0.f; + const uint4 yp = pack8(y); + if (live) *reinterpret_cast(out + r * n + 8 * c) = yp; + if (xs != nullptr) { + const float s = group_sum8(yp, live); + if (live && (c & 7) == 0) xs[r * (n / 64) + c / 8] = s; + } +} + +// xs[r, j] = the sum of x[r * ld + 64 j ..][0 .. 64] (k % 64 == 0), in group_sum8's order. +extern "C" __global__ void __launch_bounds__(NT) tf_group_sums_rows(const __nv_bfloat16* __restrict__ x, long long ld, + float* __restrict__ xs, int k) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + const bool live = 8 * c < k; + const uint4 v = live ? *reinterpret_cast(x + r * ld + 8 * c) : make_uint4(0, 0, 0, 0); + const float s = group_sum8(v, live); + if (live && (c & 7) == 0) xs[r * (k / 64) + c / 8] = s; +} + +// DFlash2's two-tap conv: out = bf16(x k0 + x[r - 1] k1) [+ res], no r - 1 at a seg block's start, k = bf16(b + dyn). +extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ dyn, + const __nv_bfloat16* __restrict__ base, + const __nv_bfloat16* __restrict__ res, + __nv_bfloat16* __restrict__ out, int d, int gs, int branch, + int seg) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + if (8 * c >= d) return; + const int g = d / gs, grp = 8 * c / gs; + float xv[8], pv[8], b0[8], b1[8], y[8]; + unpack8(*reinterpret_cast(x + r * d + 8 * c), xv); + if (r % seg > 0) unpack8(*reinterpret_cast(x + (r - 1) * d + 8 * c), pv); + else +#pragma unroll + for (int i = 0; i < 8; ++i) pv[i] = 0.f; + unpack8(*reinterpret_cast(base + (int64_t)(branch * 2) * d + 8 * c), b0); + unpack8(*reinterpret_cast(base + (int64_t)(branch * 2 + 1) * d + 8 * c), b1); + const float d0 = __bfloat162float(dyn[((r * 2 + branch) * 2) * g + grp]); + const float d1 = __bfloat162float(dyn[((r * 2 + branch) * 2 + 1) * g + grp]); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float k0 = __bfloat162float(__float2bfloat16_rn(__fadd_rn(b0[i], d0))); + const float k1 = __bfloat162float(__float2bfloat16_rn(__fadd_rn(b1[i], d1))); + y[i] = __fmaf_rn(pv[i], k1, __fmul_rn(xv[i], k0)); + } + if (res != nullptr) { + float rv[8]; + unpack8(*reinterpret_cast(res + r * d + 8 * c), rv); +#pragma unroll + for (int i = 0; i < 8; ++i) y[i] = __fadd_rn(rv[i], __bfloat162float(__float2bfloat16_rn(y[i]))); + } + *reinterpret_cast(out + r * d + 8 * c) = pack8(y); +} diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu new file mode 100644 index 0000000000..04f330f4e7 --- /dev/null +++ b/zig/kernels/cuda/gdn_prep.cu @@ -0,0 +1,99 @@ +// Gated DeltaNet inputs for every CUDA family (float64 reference: core/gdn_prep_ref.zig): conv, norms, g, beta. + +// A warp a head (q heads, k heads, then v heads), lane l holding channels J l .. J l + J - 1; a row's bits are its own. + +#include +#include +#include + +namespace tf_gdn_prep { + +constexpr int MAXT = 8; // conv taps at most + +struct Io { + const __nv_bfloat16* qkv; // [rows, qkv_ld]: q (kh dk), k (kh dk), v (vh dk) channels + long long qkv_ld; + const __nv_bfloat16* state; // [streams * keep, c]: each stream's kept conv rows + const __nv_bfloat16* conv; // [c, keep + 1] taps a channel + const int* win; // [rows, keep + 1]: < keep a kept row, else qkv row win - keep + const int* sid; // [rows] stream (null: stream 0) + const __nv_bfloat16* a; // [rows, a_ld] the v heads' dt inputs + long long a_ld; + const __nv_bfloat16* b; // [rows, b_ld] the v heads' beta inputs + long long b_ld; + const float* a_log; // [vh] + const float* dt_bias; // [vh] + __nv_bfloat16* q; // [rows, kh, dk] + __nv_bfloat16* k; // [rows, kh, dk] + __nv_bfloat16* v; // [rows, vh, dk] + float* g; // [rows, vh] + float* beta; // [rows, vh] +}; + +__device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } + +template +__device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int keep, float eps) { + constexpr int DK = 32 * J; + const int lane = threadIdx.x & 31; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y * (blockDim.x / 32) + threadIdx.x / 32; + if (hd >= 2 * kh + vh) return; + const int ch = hd * DK + J * lane; + const int64_t srow = (io.sid != nullptr ? static_cast(io.sid[r]) : 0) * keep; + int src[MAXT]; +#pragma unroll + for (int t = 0; t < MAXT; ++t) src[t] = t <= keep ? io.win[r * (keep + 1) + t] : 0; + float acc[J]; +#pragma unroll + for (int i = 0; i < J; ++i) acc[i] = 0.f; +#pragma unroll + for (int t = 0; t < MAXT; ++t) { + if (t > keep) break; + const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; +#pragma unroll + for (int i = 0; i < J; ++i) + acc[i] = __fmaf_rn(__bfloat162float(x[ch + i]), __bfloat162float(io.conv[(ch + i) * (keep + 1) + t]), acc[i]); + } + float cv[J]; +#pragma unroll + for (int i = 0; i < J; ++i) cv[i] = __bfloat162float(__float2bfloat16_rn(__fmul_rn(acc[i], sigmoid(acc[i])))); + if (hd < 2 * kh) { // q: l2norm / sqrt(dk), k: l2norm (rms with eps, scaled) + float part = 0.f; +#pragma unroll + for (int i = 0; i < J; ++i) part = __fmaf_rn(cv[i], cv[i], part); + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + const float inv = rsqrtf(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps)); + const bool is_q = hd < kh; + const float scale = is_q ? 1.f / DK : rsqrtf(static_cast(DK)); + __nv_bfloat16* out = (is_q ? io.q : io.k) + (r * kh + (is_q ? hd : hd - kh)) * DK + J * lane; +#pragma unroll + for (int i = 0; i < J; ++i) out[i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(cv[i], inv), scale)); + return; + } + const int hv = hd - 2 * kh; +#pragma unroll + for (int i = 0; i < J; ++i) io.v[(r * vh + hv) * DK + J * lane + i] = __float2bfloat16_rn(cv[i]); + if (lane == 0) { + const float x = __fadd_rn(__bfloat162float(io.a[r * io.a_ld + hv]), io.dt_bias[hv]); + const float sp = x > 20.f ? x : log1pf(expf(x)); + io.g[r * vh + hv] = expf(-__fmul_rn(expf(io.a_log[hv]), sp)); + io.beta[r * vh + hv] = sigmoid(__bfloat162float(io.b[r * io.b_ld + hv])); + } +} + +} // namespace tf_gdn_prep + +using namespace tf_gdn_prep; + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_64(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<2>(io, c, kh, vh, keep, eps); +} + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_128(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<4>(io, c, kh, vh, keep, eps); +} + +extern "C" __global__ void __launch_bounds__(128) tf_gdn_prep_256(const Io io, int c, int kh, int vh, int keep, float eps) { + prep<8>(io, c, kh, vh, keep, eps); +} diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu new file mode 100644 index 0000000000..146e14017c --- /dev/null +++ b/zig/kernels/cuda/qk_rows.cu @@ -0,0 +1,194 @@ +// Attention inputs for every CUDA family (float64 references: core/qk_ref.zig): QK norm, RoPE and the cache write. + +// A block a (row, head), a thread an element; a head's norm sums in one fixed order, never depending on the row count. + +#include +#include + +namespace tf_qk { + +constexpr int MOST = 256; // head_dim at most + +// Where a row's heads come from: q, k and v (null: none) at row stride *_ld and head stride *_hs (elements). +struct QkSrc { + const __nv_bfloat16* q; + long long q_ld, q_hs; + const __nv_bfloat16* k; + long long k_ld, k_hs; + const __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +// The q and k norm weights: fp32 or bf16 (null q: no norm). +struct QkW { + const void* q; + const void* k; + int bf16; +}; + +__device__ __forceinline__ float wv(const QkW& w, bool is_k, int i) { + const void* p = is_k ? w.k : w.q; + return w.bf16 ? __bfloat162float(static_cast(p)[i]) : static_cast(p)[i]; +} + +// cos and sin of p's angle on rotated dim d: from tables [positions, rot], or computed from inv [rot / 2]. +__device__ __forceinline__ void rope_at(const float* cosv, const float* sinv, const float* inv, int p, int rot, int d, + float& c, float& sn) { + if (inv != nullptr) { + const float a = __fmul_rn(static_cast(p), inv[d < rot / 2 ? d : d - rot / 2]); + c = cosf(a); + sn = sinf(a); + } else { + c = cosv[static_cast(p) * rot + d]; + sn = sinv[static_cast(p) * rot + d]; + } +} + +// Where a row's q, k and v (null: none) go, at row stride *_ld and head stride *_hs (elements). +struct QkDst { + __nv_bfloat16* q; + long long q_ld, q_hs; + __nv_bfloat16* k; + long long k_ld, k_hs; + __nv_bfloat16* v; + long long v_ld, v_hs; +}; + +__device__ __forceinline__ const __nv_bfloat16* head_src(const QkSrc& s, int64_t r, int hd, int h, int hk) { + if (hd < h) return s.q + r * s.q_ld + hd * s.q_hs; + if (hd < h + hk) return s.k + r * s.k_ld + (hd - h) * s.k_hs; + return s.v + r * s.v_ld + (hd - h - hk) * s.v_hs; +} + +} // namespace tf_qk + +using namespace tf_qk; + +// QkSrc's heads -> q, k, v (bf16); q, k normed by nw, rotate-half RoPE on `rot` dims; k, v also into kc/vc. +extern "C" __global__ void __launch_bounds__(MOST) tf_qk_rows(const QkSrc in, const QkW nw, const float* cosv, + const float* sinv, const float* inv, int rot, const int* pos, + const int* slot, const QkDst dst, __nv_bfloat16* kc, __nv_bfloat16* vc, + long long len, int h, int hk, float eps) { + __shared__ float warps[MOST / 32]; + const int d = threadIdx.x; + const int dim = blockDim.x; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y; + const __nv_bfloat16* src = head_src(in, r, hd, h, hk); + const float x = __bfloat162float(src[d]); + const int p = pos[r]; + const int64_t pl = p < len ? p : p % len; // a 64-bit modulo only past the cache's end + const int64_t at = slot != nullptr ? static_cast(slot[r]) * len + pl : pl; + if (hd >= h + hk) { // v: as is + const int k = hd - h - hk; + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + d] = src[d]; + if (vc != nullptr) vc[(at * hk + k) * dim + d] = src[d]; + return; + } + float y = x; + float yp = 0.f; + const int half = rot / 2; + const int partner = d < half ? d + half : d - half; + if (nw.q != nullptr) { + float part = __fmul_rn(x, x); + for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); + if ((d & 31) == 0) warps[d >> 5] = part; + __syncthreads(); + float s = 0.f; + for (int w = 0; w < dim / 32; ++w) s = __fadd_rn(s, warps[w]); + const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps)); + y = __fmul_rn(__fmul_rn(x, rstd), wv(nw, hd >= h, d)); + if (d < rot) yp = __fmul_rn(__fmul_rn(__bfloat162float(src[partner]), rstd), wv(nw, hd >= h, partner)); + } else if (d < rot) { + yp = __bfloat162float(src[partner]); + } + if ((cosv != nullptr || inv != nullptr) && d < rot) { + float c, sn; + rope_at(cosv, sinv, inv, p, rot, d, c, sn); + y = __fadd_rn(__fmul_rn(y, c), __fmul_rn(d < half ? -yp : yp, sn)); + } + const __nv_bfloat16 out = __float2bfloat16_rn(y); + if (hd < h) { + dst.q[r * dst.q_ld + hd * dst.q_hs + d] = out; + } else { + const int k = hd - h; + dst.k[r * dst.k_ld + k * dst.k_hs + d] = out; + if (kc != nullptr) kc[(at * hk + k) * dim + d] = out; + } +} + +// tf_qk_rows at dim 128, no or full rotation: a warp a head, lane l dims l + 32 j (same fold order, partner j ^ 2). +extern "C" __global__ void __launch_bounds__(128) tf_qk_rows_w(const QkSrc in, const QkW nw, + const float* cosv, const float* sinv, const float* inv, int rot, + const int* pos, const int* slot, const QkDst dst, + __nv_bfloat16* kc, __nv_bfloat16* vc, long long len, + int h, int hk, float eps) { + constexpr int D = 128, J = D / 32; + const int lane = threadIdx.x & 31; + const int64_t r = blockIdx.x; + const int hd = blockIdx.y * 4 + (threadIdx.x >> 5); + if (hd >= h + (in.v != nullptr ? 2 * hk : hk)) return; + const __nv_bfloat16* src = head_src(in, r, hd, h, hk); + const int p = pos[r]; + __nv_bfloat16 raw[J]; +#pragma unroll + for (int j = 0; j < J; ++j) raw[j] = src[lane + 32 * j]; + // every load up front (the rotation's row waits on pos, then on memory) + float c[J], sn[J], w[J]; + const bool qk = hd < h + hk; +#pragma unroll + for (int j = 0; j < J; ++j) { + c[j] = sn[j] = 0.f; + if (qk && rot != 0 && (cosv != nullptr || inv != nullptr)) rope_at(cosv, sinv, inv, p, D, lane + 32 * j, c[j], sn[j]); + w[j] = qk && nw.q != nullptr ? wv(nw, hd >= h, lane + 32 * j) : 0.f; + } + const int64_t pl = p < len ? p : p % len; + const int64_t at = slot != nullptr ? static_cast(slot[r]) * len + pl : pl; + if (hd >= h + hk) { // v: as is + const int k = hd - h - hk; +#pragma unroll + for (int j = 0; j < J; ++j) { + if (dst.v != nullptr) dst.v[r * dst.v_ld + k * dst.v_hs + lane + 32 * j] = raw[j]; + if (vc != nullptr) vc[(at * hk + k) * D + lane + 32 * j] = raw[j]; + } + return; + } + float x[J], y[J], yp[J]; +#pragma unroll + for (int j = 0; j < J; ++j) x[j] = __bfloat162float(raw[j]), y[j] = x[j], yp[j] = 0.f; + if (nw.q != nullptr) { + float part[J]; +#pragma unroll + for (int j = 0; j < J; ++j) part[j] = __fmul_rn(x[j], x[j]); + for (int m = 16; m > 0; m >>= 1) +#pragma unroll + for (int j = 0; j < J; ++j) part[j] = __fadd_rn(part[j], __shfl_xor_sync(0xFFFFFFFFu, part[j], m)); + float s = 0.f; +#pragma unroll + for (int j = 0; j < J; ++j) s = __fadd_rn(s, part[j]); + const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps)); +#pragma unroll + for (int j = 0; j < J; ++j) { + y[j] = __fmul_rn(__fmul_rn(x[j], rstd), w[j]); + if (rot != 0) yp[j] = __fmul_rn(__fmul_rn(x[j ^ 2], rstd), w[j ^ 2]); + } + } else if (rot != 0) { +#pragma unroll + for (int j = 0; j < J; ++j) yp[j] = x[j ^ 2]; + } + if (rot != 0 && (cosv != nullptr || inv != nullptr)) { +#pragma unroll + for (int j = 0; j < J; ++j) y[j] = __fadd_rn(__fmul_rn(y[j], c[j]), __fmul_rn(j < 2 ? -yp[j] : yp[j], sn[j])); + } +#pragma unroll + for (int j = 0; j < J; ++j) { + const __nv_bfloat16 out = __float2bfloat16_rn(y[j]); + if (hd < h) { + dst.q[r * dst.q_ld + hd * dst.q_hs + lane + 32 * j] = out; + } else { + const int k = hd - h; + dst.k[r * dst.k_ld + k * dst.k_hs + lane + 32 * j] = out; + if (kc != nullptr) kc[(at * hk + k) * D + lane + 32 * j] = out; + } + } +} diff --git a/zig/src/core/act_ref.zig b/zig/src/core/act_ref.zig index ccab2c254e..f56f282175 100644 --- a/zig/src/core/act_ref.zig +++ b/zig/src/core/act_ref.zig @@ -5,6 +5,16 @@ pub fn sigmoid(x: f64) f64 { return 1 / (1 + @exp(-x)); } +/// silu(gate) * up. +pub fn swiglu(gate: f64, up: f64) f64 { + return gate * sigmoid(gate) * up; +} + +/// DFlash2's dynamic conv at one value: x k0 + prev k1 with the taps as rounded (k_t = bf16(base_t + dyn_t)). +pub fn dconv(x: f64, prev: f64, k0: f64, k1: f64) f64 { + return x * k0 + prev * k1; +} + /// A 4-bit affine table's value: q * scale + bias. pub fn affine4(q: u4, scale: f64, bias: f64) f64 { return @as(f64, @floatFromInt(q)) * scale + bias; @@ -18,5 +28,6 @@ pub fn tolerance(want: f64) f64 { test "activations against hand values" { try std.testing.expectEqual(@as(f64, 0.5), sigmoid(0)); try std.testing.expectEqual(@as(f64, 7 * 0.5 - 1), affine4(7, 0.5, -1)); + try std.testing.expectApproxEqRel(@as(f64, 2 * sigmoid(2) * 3), swiglu(2, 3), 1e-15); try std.testing.expect(tolerance(1) > tolerance(0.5)); } diff --git a/zig/src/core/gdn_prep_ref.zig b/zig/src/core/gdn_prep_ref.zig new file mode 100644 index 0000000000..59a812d6df --- /dev/null +++ b/zig/src/core/gdn_prep_ref.zig @@ -0,0 +1,45 @@ +//! Gated DeltaNet inputs in float64, the reference every backend's GDN prep kernels are checked against. +const std = @import("std"); + +/// A head's conv over `taps` rows (x [taps][dk], w [dk][taps]), then SiLU -> c; err: a channel's fp32 bound. +pub fn conv(x: []const []const f64, w: []const f64, taps: usize, c: []f64, err: []f64) void { + for (c, err, 0..) |*o, *e, i| { + var acc: f64 = 0; + var mag: f64 = 0; + for (0..taps) |t| { + acc += x[t][i] * w[i * taps + t]; + mag += @abs(x[t][i] * w[i * taps + t]); + } + o.* = acc / (1 + @exp(-acc)); + e.* = 1.1 * @as(f64, @floatFromInt(taps)) * std.math.pow(f64, 2, -24) * mag; // silu' is at most 1.1 + } +} + +/// q (is_q) or k from a head's c: c / sqrt(mean(c^2) + eps) * (1 / dk for q, 1 / sqrt(dk) for k). +pub fn norm(c: []const f64, is_q: bool, eps: f64, out: []f64) f64 { + var ss: f64 = 0; + for (c) |v| ss += v * v; + const dk: f64 = @floatFromInt(c.len); + const f = 1 / @sqrt(ss / dk + eps) * (if (is_q) 1 / dk else 1 / @sqrt(dk)); + for (out, c) |*o, v| o.* = v * f; + return f; +} + +/// g = exp(-exp(a_log) * softplus(a + dt_bias)). +pub fn gate(a: f64, dt_bias: f64, a_log: f64) f64 { + const x = a + dt_bias; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + return @exp(-@exp(a_log) * sp); +} + +test "conv, norms and the gate against hand values" { + var c: [2]f64 = undefined; + var e: [2]f64 = undefined; + conv(&.{ &.{ 1, 2 }, &.{ 3, 4 } }, &.{ 1, 0, 0, 1 }, 2, &c, &e); // channel 0 tap 0, channel 1 tap 1 + try std.testing.expectApproxEqRel(@as(f64, 1.0 / (1.0 + @exp(-1.0))), c[0], 1e-15); + try std.testing.expectApproxEqRel(@as(f64, 4.0 / (1.0 + @exp(-4.0))), c[1], 1e-15); + var o: [2]f64 = undefined; + _ = norm(&.{ 3, 4 }, false, 0, &o); + try std.testing.expectApproxEqRel(@as(f64, 0.6), o[0], 1e-15); + try std.testing.expectApproxEqRel(@as(f64, @exp(-@exp(0.5) * 25.0)), gate(25, 0, 0.5), 1e-15); +} diff --git a/zig/src/core/qk_ref.zig b/zig/src/core/qk_ref.zig new file mode 100644 index 0000000000..b1763d9f34 --- /dev/null +++ b/zig/src/core/qk_ref.zig @@ -0,0 +1,30 @@ +//! QK norm and RoPE of one head in float64, the reference for every backend's attention inputs (cuda/qk_rows.zig). +const std = @import("std"); + +/// One head: RMS-normed by w (none: as is), then rotate-half RoPE on its first `rot` dims (no cos/sin: none). +pub fn head(x: []const f64, w: ?[]const f64, eps: f64, cos: ?[]const f64, sin: ?[]const f64, rot: usize, y: []f64) void { + const n = x.len; + var ss: f64 = 0; + for (x) |v| ss += v * v; + const inv = if (w != null) 1 / @sqrt(ss / @as(f64, @floatFromInt(n)) + eps) else 1; + var z: [256]f64 = undefined; + for (x, 0..) |v, i| z[i] = v * inv * (if (w) |ww| ww[i] else 1); + const half = rot / 2; + for (y, 0..) |*o, i| { + if (cos == null or i >= rot) { + o.* = z[i]; + continue; + } + const p = if (i < half) i + half else i - half; + const rotated = if (i < half) -z[p] else z[p]; + o.* = z[i] * cos.?[i] + rotated * sin.?[i]; + } +} + +test "rotate-half pairs (d, d + rot / 2) and leave the rest" { + var y: [4]f64 = undefined; + head(&.{ 1, 2, 3, 4 }, null, 0, &.{ 0, 0 }, &.{ 1, 1 }, 2, &y); // angle 90 degrees on the first pair + try std.testing.expectEqualSlices(f64, &.{ -2, 1, 3, 4 }, &y); + head(&.{ 2, 2, 2, 2 }, &.{ 1, 1, 1, 3 }, 0, null, null, 0, &y); + try std.testing.expectEqualSlices(f64, &.{ 1, 1, 1, 3 }, &y); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 122ffe9e4c..01934d154d 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -28,6 +28,8 @@ pub const modelopt = @import("modelopt.zig"); pub const picture = @import("picture/picture.zig"); /// Qwen-VL image input every Qwen3.5-VL checkpoint shares: settings, geometry, M-RoPE positions, the tower's tables. pub const qwen_vl = @import("qwen_vl/qwen_vl.zig"); +pub const gdn_prep_ref = @import("gdn_prep_ref.zig"); +pub const qk_ref = @import("qk_ref.zig"); test { _ = safetensors; @@ -41,6 +43,8 @@ test { _ = norms_ref; _ = act_ref; _ = contract; + _ = gdn_prep_ref; + _ = qk_ref; _ = exl3_format; _ = exlr8; _ = eaq_container; diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig new file mode 100644 index 0000000000..832844ce01 --- /dev/null +++ b/zig/src/cuda/act_rows.zig @@ -0,0 +1,95 @@ +//! act_rows.cu's activation rows for every CUDA family: SwiGLU, sigmoid output gates and 64-value group sums. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +const threads = 256; + +fn aligned(ptrs: []const u64, strides: []const usize) bool { + for (ptrs) |p| if (p % 16 != 0) return false; + for (strides) |v| if (v % 8 != 0) return false; + return true; +} + +fn grid(rows: usize, n: usize) launch_.Dim3 { + return .{ .x = @intCast(rows), .y = @intCast((n / 8 + threads - 1) / threads), .z = 1 }; +} + +pub const ActRows = struct { + mod: Module, + swiglu_fn: Function, + gate_fn: Function, + sums_fn: Function, + dconv_fn: Function, + + pub fn load(d: *const Driver) !ActRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.act_rows); + errdefer mod.unload(); + return .{ .mod = mod, .swiglu_fn = try mod.function("tf_swiglu_rows"), .gate_fn = try mod.function("tf_gate_rows"), .sums_fn = try mod.function("tf_group_sums_rows"), .dconv_fn = try mod.function("tf_dconv_rows") }; + } + + pub fn unload(g: *ActRows) void { + g.mod.unload(); + } + + /// out [rows, n] = bf16(silu(gate) * up), rows at strides g_ld, u_ld; xs (or null): out's 64-value sums. + pub fn swiglu(g: *const ActRows, s: Stream, gate: u64, g_ld: usize, up: u64, u_ld: usize, out: u64, xs: ?u64, rows: usize, n: usize) !void { + if (n % 64 != 0 or !aligned(&.{ gate, up, out }, &.{ g_ld, u_ld })) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(gate); + a.add(@as(i64, @intCast(g_ld))); + a.add(up); + a.add(@as(i64, @intCast(u_ld))); + a.add(out); + a.add(xs orelse 0); + a.add(@as(i32, @intCast(n))); + try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n), .block = .{ .x = threads } }, s, &a); + } + + /// out = bf16(o * sigmoid(gate)), head h's gate at r * g_ld + h * g_hs ([q | gate]: q + d, 2 d); xs: sums. + pub fn outputGate(g: *const ActRows, s: Stream, o: u64, gate_: u64, g_ld: usize, g_hs: usize, out: u64, xs: ?u64, rows: usize, heads: usize, d: usize) !void { + if (d % 64 != 0 or !aligned(&.{ o, gate_, out }, &.{ g_ld, g_hs })) return error.RowShape; + if (rows == 0 or heads == 0) return; + var a: launch_.Args = .{}; + a.add(o); + a.add(gate_); + a.add(@as(i64, @intCast(g_ld))); + a.add(@as(i64, @intCast(g_hs))); + a.add(out); + a.add(xs orelse 0); + a.add(@as(i32, @intCast(heads))); + a.add(@as(i32, @intCast(d))); + try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d), .block = .{ .x = threads } }, s, &a); + } + + /// xs [rows, k / 64]: x's rows (stride ld) summed over each 64 values. k % 64 == 0. + pub fn groupSums(g: *const ActRows, s: Stream, x: u64, ld: usize, xs: u64, rows: usize, k: usize) !void { + if (k % 64 != 0 or !aligned(&.{x}, &.{ld})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(ld))); + a.add(xs); + a.add(@as(i32, @intCast(k))); + try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k), .block = .{ .x = threads } }, s, &a); + } + + /// DFlash2's two-tap conv over [rows, d] bf16: dyn [rows, 2, 2, d / gs], base [2, 2, d]; res (or null) added. + pub fn dconv(g: *const ActRows, s: Stream, x: u64, dyn: u64, base: u64, res: ?u64, out: u64, rows: usize, d: usize, gs: usize, branch: usize, seg: usize) !void { + if (d % 8 != 0 or gs % 8 != 0 or d % gs != 0 or branch > 1 or seg == 0 or !aligned(&.{ x, base, res orelse 0, out }, &.{})) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ x, dyn, base, res orelse 0, out }) |v| a.add(v); + for ([_]usize{ d, gs, branch, seg }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d), .block = .{ .x = threads } }, s, &a); + } +}; diff --git a/zig/src/cuda/gdn_prep.zig b/zig/src/cuda/gdn_prep.zig new file mode 100644 index 0000000000..d11f869096 --- /dev/null +++ b/zig/src/cuda/gdn_prep.zig @@ -0,0 +1,72 @@ +//! gdn_prep.cu's Gated DeltaNet inputs for every CUDA family: each row's conv window, SiLU, q/k norms, g and beta. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// The kernel's buffers (addresses; strides in elements), as gdn_prep.cu's Io. +pub const Io = extern struct { + qkv: u64, // [rows, qkv_ld] bf16: q (kh dk), k (kh dk), v (vh dk) channels + qkv_ld: i64, + state: u64, // [streams * keep, c] bf16: each stream's kept conv rows + conv: u64, // [c, keep + 1] bf16 taps a channel + win: u64, // [rows, keep + 1] i32: < keep a kept row of the row's stream, else qkv row (win - keep) + sid: u64 = 0, // [rows] i32 stream (0: every row stream 0) + a: u64, // [rows, a_ld] bf16 + a_ld: i64, + b: u64, // [rows, b_ld] bf16 + b_ld: i64, + a_log: u64, // [vh] fp32 + dt_bias: u64, // [vh] fp32 + q: u64, // [rows, kh, dk] bf16 + k: u64, + v: u64, // [rows, vh, dk] bf16 + g: u64, // [rows, vh] fp32 + beta: u64, +}; + +/// Channels a row (c = 2 kh dk + vh dk), heads and kept conv rows (taps - 1). +pub const Shape = struct { kh: usize, vh: usize, dk: usize, keep: usize, eps: f32 = 1e-6 }; + +pub const GdnPrep = struct { + mod: Module, + fns: [3]Function, // dk 64, 128, 256 + + pub fn load(d: *const Driver) !GdnPrep { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.gdn_prep); + errdefer mod.unload(); + return .{ .mod = mod, .fns = .{ try mod.function("tf_gdn_prep_64"), try mod.function("tf_gdn_prep_128"), try mod.function("tf_gdn_prep_256") } }; + } + + pub fn unload(g: *GdnPrep) void { + g.mod.unload(); + } + + /// Per row's conv window c = silu(conv): q = l2norm(c) / sqrt(dk), k = l2norm(c), v = c, g and beta. + pub fn run(g: *const GdnPrep, s: Stream, io: Io, rows: usize, sh: Shape) !void { + const which: usize = switch (sh.dk) { + 64 => 0, + 128 => 1, + 256 => 2, + else => return error.GdnShape, + }; + if (sh.keep + 1 > 8 or sh.keep == 0) return error.GdnShape; + if (rows == 0) return; + const heads = 2 * sh.kh + sh.vh; + var a: launch_.Args = .{}; + a.add(io); + a.add(@as(i32, @intCast(heads * sh.dk))); // c: a kept row's channels + a.add(@as(i32, @intCast(sh.kh))); + a.add(@as(i32, @intCast(sh.vh))); + a.add(@as(i32, @intCast(sh.keep))); + a.add(sh.eps); + try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + 3) / 4), .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..3326909329 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -56,6 +56,9 @@ pub const ck_lane: []const u8 = if (available) &Blob("fatbin_ck_lane").bytes els pub const ck_experts: []const u8 = if (available) &Blob("fatbin_ck_experts").bytes else &.{}; pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; +pub const act_rows: []const u8 = if (available) &Blob("fatbin_act_rows").bytes else &.{}; +pub const gdn_prep: []const u8 = if (available) &Blob("fatbin_gdn_prep").bytes else &.{}; +pub const qk_rows: []const u8 = if (available) &Blob("fatbin_qk_rows").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/cuda/qk_rows.zig b/zig/src/cuda/qk_rows.zig new file mode 100644 index 0000000000..d1ccb1a4ed --- /dev/null +++ b/zig/src/cuda/qk_rows.zig @@ -0,0 +1,107 @@ +//! qk_rows.cu's attention inputs for every CUDA family: QK norm, RoPE (rotate-half, partial) and the cache write. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// A packed row's heads: h query, hk key and hk value heads of `dim` (a multiple of 32, at most 256). +pub const Heads = struct { h: usize, hk: usize, dim: usize }; + +/// Where a row's heads come from: q, k and v (0: none) at row stride *_ld and head stride *_hs (elements). +pub const Src = extern struct { q: u64, q_ld: i64, q_hs: i64, k: u64, k_ld: i64, k_hs: i64, v: u64 = 0, v_ld: i64 = 0, v_hs: i64 = 0 }; + +/// q/k norms [2, dim] (or q's with norm_k), fp32 or bf16; RoPE on `rot` dims: cos/sin tables or inv [rot / 2]. +pub const Shape = struct { norms: ?u64 = null, norm_k: ?u64 = null, norms_bf16: bool = false, cos: ?u64 = null, sin: ?u64 = null, inv: ?u64 = null, rot: usize = 0, eps: f32 = 1e-6 }; + +/// qk_rows.cu's QkW. +const W = extern struct { q: u64, k: u64, bf16: i32 }; + +/// Where the rows go: q [n, h, dim], k and v [n, hk, dim]; k and v also into kc/vc at (slot * len + pos % len). +pub const Out = struct { + q: u64, + k: u64, + v: u64, + kc: ?u64 = null, + vc: ?u64 = null, + len: usize = 1, + /// Row and head strides in elements (0: packed [n, heads, dim]); head-major [heads, n, dim]: ld dim, hs n dim. + q_ld: usize = 0, + q_hs: usize = 0, + kv_ld: usize = 0, + kv_hs: usize = 0, +}; + +/// qk_rows.cu's QkDst. +const Dst = extern struct { q: u64, q_ld: i64, q_hs: i64, k: u64, k_ld: i64, k_hs: i64, v: u64, v_ld: i64, v_hs: i64 }; + +pub const QkRows = struct { + mod: Module, + fn_: Function, + /// tf_qk_rows_w: dim 128 with no or full rotation, a warp a head. + warp_fn: Function, + + pub fn load(d: *const Driver) !QkRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.qk_rows); + errdefer mod.unload(); + return .{ .mod = mod, .fn_ = try mod.function("tf_qk_rows"), .warp_fn = try mod.function("tf_qk_rows_w") }; + } + + pub fn unload(g: *QkRows) void { + g.mod.unload(); + } + + /// n packed rows qkv [n, (h + 2 hk) dim] bf16 at positions pos (int32 [n]) in cache slots slot (int32 [n] or none). + pub fn rows(g: *const QkRows, s: Stream, qkv: u64, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + const ld: i64 = @intCast((hs.h + 2 * hs.hk) * hs.dim); + const dim: i64 = @intCast(hs.dim); + const src: Src = .{ .q = qkv, .q_ld = ld, .q_hs = dim, .k = qkv + hs.h * hs.dim * 2, .k_ld = ld, .k_hs = dim, .v = qkv + (hs.h + hs.hk) * hs.dim * 2, .v_ld = ld, .v_hs = dim }; + return g.rowsFrom(s, src, hs, shape, pos, slot, out, n); + } + + /// As `rows`, the heads read from `src` (e.g. [q | gate] rows: q_hs = 2 dim); no v source: no v out or vc. + pub fn rowsFrom(g: *const QkRows, s: Stream, src: Src, hs: Heads, shape: Shape, pos: u64, slot: ?u64, out: Out, n: usize) !void { + if (hs.dim % 32 != 0 or hs.dim > 256 or shape.rot % 2 != 0 or shape.rot > hs.dim) return error.HeadShape; + if ((shape.cos == null) != (shape.sin == null) or (shape.cos != null and shape.inv != null)) return error.HeadShape; + const roped = shape.cos != null or shape.inv != null; + if (roped and shape.rot == 0) return error.HeadShape; + if (src.v == 0 and out.vc != null) return error.HeadShape; + if (n == 0) return; + var a: launch_.Args = .{}; + a.add(src); + const wsize: usize = if (shape.norms_bf16) 2 else 4; + const nq = shape.norms orelse 0; + a.add(W{ .q = nq, .k = shape.norm_k orelse (if (nq != 0) nq + hs.dim * wsize else 0), .bf16 = @intFromBool(shape.norms_bf16) }); + a.add(shape.cos orelse 0); + a.add(shape.sin orelse 0); + a.add(shape.inv orelse 0); + a.add(@as(i32, @intCast(if (roped) shape.rot else 0))); + a.add(pos); + a.add(slot orelse 0); + const dim: i64 = @intCast(hs.dim); + const qld: i64 = if (out.q_ld != 0) @intCast(out.q_ld) else @as(i64, @intCast(hs.h)) * dim; + const qhs: i64 = if (out.q_hs != 0) @intCast(out.q_hs) else dim; + const kld: i64 = if (out.kv_ld != 0) @intCast(out.kv_ld) else @as(i64, @intCast(hs.hk)) * dim; + const khs: i64 = if (out.kv_hs != 0) @intCast(out.kv_hs) else dim; + a.add(Dst{ .q = out.q, .q_ld = qld, .q_hs = qhs, .k = out.k, .k_ld = kld, .k_hs = khs, .v = out.v, .v_ld = kld, .v_hs = khs }); + a.add(out.kc orelse 0); + a.add(out.vc orelse 0); + a.add(@as(i64, @intCast(out.len))); + a.add(@as(i32, @intCast(hs.h))); + a.add(@as(i32, @intCast(hs.hk))); + a.add(shape.eps); + const heads = hs.h + (if (src.v != 0) 2 * hs.hk else hs.hk); + if (hs.dim == 128 and (!roped or shape.rot == 128)) { + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast((heads + 3) / 4), .z = 1 }; + return launch_.launch(g.warp_fn, .{ .grid = grid, .block = .{ .x = 128 } }, s, &a); + } + const grid: launch_.Dim3 = .{ .x = @intCast(n), .y = @intCast(heads), .z = 1 }; + try launch_.launch(g.fn_, .{ .grid = grid, .block = .{ .x = @intCast(hs.dim) } }, s, &a); + } +}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6c4e0991b6..710c2e76a7 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -34,6 +34,9 @@ pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); pub const vit = @import("vit.zig"); +pub const act_rows = @import("act_rows.zig"); +pub const gdn_prep = @import("gdn_prep.zig"); +pub const qk_rows = @import("qk_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const rank_scatter = @import("rank_scatter.zig"); @@ -49,6 +52,9 @@ test { _ = @import("norms.zig"); _ = @import("embed_rows.zig"); _ = @import("vit.zig"); + _ = @import("act_rows.zig"); + _ = @import("gdn_prep.zig"); + _ = @import("qk_rows.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = @import("rank_scatter.zig"); diff --git a/zig/tests/cuda/act_rows_tests.zig b/zig/tests/cuda/act_rows_tests.zig new file mode 100644 index 0000000000..43ddaf4b7e --- /dev/null +++ b/zig/tests/cuda/act_rows_tests.zig @@ -0,0 +1,132 @@ +//! act_rows.cu's SwiGLU, output gate, group sums and DFlash2's dynamic conv against float64, bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rows = @import("norm_rows_tests.zig"); + +const aref = core.act_ref; +const counts = rows.counts; +const most = rows.most; +const Rig = rows.Rig; +const bf16 = rows.bf16; +const bf = rows.bf; + +fn swiglu(r: *Rig, g: *const cuda.act_rows.ActRows, n: usize) !void { + const gu = r.bfs(most * 2 * n, 3); // [gate | up] rows + const d = try r.dev(std.mem.sliceAsBytes(gu)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * n * 2); + xss[c] = try r.out(m * n / 64 * 4); + try g.swiglu(r.s, d, 2 * n, d + 2 * n, 2 * n, os[c], xss[c], m, n); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "swiglu {d}", .{n}); + const o = try r.back(u16, os[counts.len - 1], most * n); + const xs = try r.back(f32, xss[counts.len - 1], most * n / 64); + for (0..most) |row| for (0..n) |i| { + const w = aref.swiglu(bf(gu[row * 2 * n + i]), bf(gu[row * 2 * n + n + i])); + try r.within(label, bf(o[row * n + i]), w, aref.tolerance(w)); + }; + try r.sums(label, o, xs, most, n); + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * n)), std.mem.sliceAsBytes(o[0 .. m * n])); +} + +fn outputGate(r: *Rig, g: *const cuda.act_rows.ActRows, heads: usize, d: usize) !void { + const n = heads * d; + const o = r.bfs(most * n, 1); + const qg = r.bfs(most * 2 * n, 3); // [q_h | gate_h] a head + const do_ = try r.dev(std.mem.sliceAsBytes(o)); + const dq = try r.dev(std.mem.sliceAsBytes(qg)); + var os: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * n * 2); + xss[c] = try r.out(m * n / 64 * 4); + try g.outputGate(r.s, do_, dq + 2 * d, 2 * n, 2 * d, os[c], xss[c], m, heads, d); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "output gate {d} heads of {d}", .{ heads, d }); + const got = try r.back(u16, os[counts.len - 1], most * n); + const xs = try r.back(f32, xss[counts.len - 1], most * n / 64); + for (0..most) |row| for (0..heads) |h| for (0..d) |e| { + const w = bf(o[row * n + h * d + e]) * aref.sigmoid(bf(qg[row * 2 * n + h * 2 * d + d + e])); + try r.within(label, bf(got[row * n + h * d + e]), w, aref.tolerance(w)); + }; + try r.sums(label, got, xs, most, n); + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * n)), std.mem.sliceAsBytes(got[0 .. m * n])); +} + +fn groupSums(r: *Rig, g: *const cuda.act_rows.ActRows, k: usize) !void { + const ld = k + 64; + const x = r.bfs(most * ld, 2); + const d = try r.dev(std.mem.sliceAsBytes(x)); + const xsd = try r.out(most * k / 64 * 4); + try g.groupSums(r.s, d, ld, xsd, most, k); + const xs = try r.back(f32, xsd, most * k / 64); + const dense = try r.a.alloc(u16, most * k); + for (0..most) |row| @memcpy(dense[row * k ..][0..k], x[row * ld ..][0..k]); + try r.sums("group sums", dense, xs, most, k); +} + +fn dconv(r: *Rig, g: *const cuda.act_rows.ActRows, d: usize, gs: usize, seg: usize, branch: usize, residual: bool) !void { + const x = r.bfs(most * d, 1); + const dyn = r.bfs(most * 4 * (d / gs), 0.3); + const base = r.bfs(4 * d, 0.7); + const res = r.bfs(most * d, 1); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const ddyn = try r.dev(std.mem.sliceAsBytes(dyn)); + const dbase = try r.dev(std.mem.sliceAsBytes(base)); + const dres = try r.dev(std.mem.sliceAsBytes(res)); + var os: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + os[c] = try r.out(m * d * 2); + try g.dconv(r.s, dx, ddyn, dbase, if (residual) dres else null, os[c], m, d, gs, branch, seg); + } + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "dconv d {d} gs {d} seg {d} branch {d}{s}", .{ d, gs, seg, branch, if (residual) " +res" else "" }); + const got = try r.back(u16, os[counts.len - 1], most * d); + const half = std.math.pow(f64, 2, -8); + for (0..most) |row| for (0..d) |i| { + const grp = i / gs; + const k0 = bf(bf16(@as(f32, @floatCast(bf(base[(branch * 2) * d + i]))) + @as(f32, @floatCast(bf(dyn[((row * 2 + branch) * 2) * (d / gs) + grp]))))); + const k1 = bf(bf16(@as(f32, @floatCast(bf(base[(branch * 2 + 1) * d + i]))) + @as(f32, @floatCast(bf(dyn[((row * 2 + branch) * 2 + 1) * (d / gs) + grp]))))); + const xv = bf(x[row * d + i]); + const pv = if (row % seg > 0) bf(x[(row - 1) * d + i]) else 0; + const y = aref.dconv(xv, pv, k0, k1); + const yt = (@abs(xv * k0) + @abs(pv * k1)) * std.math.pow(f64, 2, -22); + if (residual) { + const rv = bf(res[row * d + i]); + try r.within(label, bf(got[row * d + i]), rv + y, @abs(rv + y) * half + @abs(y) * half + yt + 1e-30); + } else try r.within(label, bf(got[row * d + i]), y, @abs(y) * half + yt + 1e-30); + }; + for (counts[0 .. counts.len - 1], 0..) |m, c| + try check.sameBytes(label, std.mem.sliceAsBytes(try r.back(u16, os[c], m * d)), std.mem.sliceAsBytes(got[0 .. m * d])); +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(27); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var g = try cuda.act_rows.ActRows.load(gpu.d); + defer g.unload(); + try swiglu(&r, &g, 17408); // Qwen3.8-27B's MLP + try swiglu(&r, &g, 192); + try outputGate(&r, &g, 24, 256); // the 27B's attention output gate + try outputGate(&r, &g, 3, 64); + try groupSums(&r, &g, 5120); + try dconv(&r, &g, 5120, 16, 16, 0, false); // DFlash2's block conv (groups of 16) + try dconv(&r, &g, 5120, 16, 16, 1, true); + try dconv(&r, &g, 256, 64, 7, 0, true); + check.pass("act rows: swiglu, output gate, dconv within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); +} diff --git a/zig/tests/cuda/gdn_prep_tests.zig b/zig/tests/cuda/gdn_prep_tests.zig new file mode 100644 index 0000000000..1ad46ec614 --- /dev/null +++ b/zig/tests/cuda/gdn_prep_tests.zig @@ -0,0 +1,156 @@ +//! gdn_prep.cu against core/gdn_prep_ref.zig: random conv windows over two streams, bits at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const ref = core.gdn_prep_ref; +const counts = [_]usize{ 1, 3, 17, 40 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +fn bfs(a: std.mem.Allocator, rng: std.Random, n: usize, scale: f32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*x| x.* = bf16(rng.floatNorm(f32) * scale); + return v; +} + +const Bufs = struct { + gpu: check.Gpu, + list: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn up(b: *Bufs, bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(b.gpu.d, bytes); + try b.list.append(b.gpu.gpa, d); + return d.ptr; + } + + fn out(b: *Bufs, s: cuda.Stream, n: usize) !u64 { + const d = try cuda.DeviceBuffer.alloc(b.gpu.d, @max(n, 16)); + try b.list.append(b.gpu.gpa, d); + try d.fill8(0xff, s.handle); + return d.ptr; + } + + fn back(b: *Bufs, a: std.mem.Allocator, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try a.alloc(T, n); + const d: cuda.DeviceBuffer = .{ .d = b.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try d.download(0, std.mem.sliceAsBytes(o)); + return o; + } +}; + +fn case(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.gdn_prep.GdnPrep, rng: std.Random, sh: cuda.gdn_prep.Shape, worst: *f64) !void { + var b: Bufs = .{ .gpu = gpu }; + defer { + for (b.list.items) |*d| d.free(); + b.list.deinit(gpu.gpa); + } + const c = (2 * sh.kh + sh.vh) * sh.dk; + const taps = sh.keep + 1; + const ld = c + 64; // stacked rows: something after the channels + const ab_ld = sh.vh + 3; + const qkv = try bfs(a, rng, most * ld, 1); + const state = try bfs(a, rng, 2 * sh.keep * c, 1); + const conv = try bfs(a, rng, c * taps, 0.4); + const av = try bfs(a, rng, most * ab_ld, 2); + const bv = try bfs(a, rng, most * ab_ld, 2); + const a_log = try a.alloc(f32, sh.vh); + const dtb = try a.alloc(f32, sh.vh); + for (a_log, dtb) |*x, *y| { + x.* = rng.float(f32) * 2 - 1; + y.* = rng.floatNorm(f32); + } + const sid = try a.alloc(i32, most); + const win = try a.alloc(i32, most * taps); + for (sid, 0..) |*x, i| x.* = @intCast(i % 2); + for (win, 0..) |*w, i| { // kept rows or any qkv row (a tree's window is any rows) + w.* = if (rng.uintLessThan(u32, 3) == 0) @intCast(rng.uintLessThan(usize, sh.keep)) else @intCast(sh.keep + rng.uintLessThan(usize, most)); + if (i % taps == taps - 1) w.* = @intCast(sh.keep + i / taps); // the row itself last + } + const base: cuda.gdn_prep.Io = .{ .qkv = try b.up(std.mem.sliceAsBytes(qkv)), .qkv_ld = @intCast(ld), .state = try b.up(std.mem.sliceAsBytes(state)), .conv = try b.up(std.mem.sliceAsBytes(conv)), .win = try b.up(std.mem.sliceAsBytes(win)), .sid = try b.up(std.mem.sliceAsBytes(sid)), .a = try b.up(std.mem.sliceAsBytes(av)), .a_ld = @intCast(ab_ld), .b = try b.up(std.mem.sliceAsBytes(bv)), .b_ld = @intCast(ab_ld), .a_log = try b.up(std.mem.sliceAsBytes(a_log)), .dt_bias = try b.up(std.mem.sliceAsBytes(dtb)), .q = 0, .k = 0, .v = 0, .g = 0, .beta = 0 }; + var ios: [counts.len]cuda.gdn_prep.Io = undefined; + for (counts, &ios) |m, *io| { + io.* = base; + io.q = try b.out(s, m * sh.kh * sh.dk * 2); + io.k = try b.out(s, m * sh.kh * sh.dk * 2); + io.v = try b.out(s, m * sh.vh * sh.dk * 2); + io.g = try b.out(s, m * sh.vh * 4); + io.beta = try b.out(s, m * sh.vh * 4); + try g.run(s, io.*, m, sh); + } + try s.synchronize(); + const io = ios[counts.len - 1]; + const q = try b.back(a, u16, io.q, most * sh.kh * sh.dk); + const k = try b.back(a, u16, io.k, most * sh.kh * sh.dk); + const v = try b.back(a, u16, io.v, most * sh.vh * sh.dk); + const gg = try b.back(a, f32, io.g, most * sh.vh); + const be = try b.back(a, f32, io.beta, most * sh.vh); + var name: [64]u8 = undefined; + const label = try std.fmt.bufPrint(&name, "gdn prep kh {d} vh {d} dk {d} taps {d}", .{ sh.kh, sh.vh, sh.dk, taps }); + const rows = try a.alloc([]f64, taps); + for (rows) |*x| x.* = try a.alloc(f64, sh.dk); + const xs = try a.alloc([]const f64, taps); + const w = try a.alloc(f64, sh.dk * taps); + const cv = try a.alloc(f64, sh.dk); + const err = try a.alloc(f64, sh.dk); + const want = try a.alloc(f64, sh.dk); + var most_q: f64 = 0; + const rel = 3 * std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -18); // c rounded, the norm through it, the output + for (0..most) |r| for (0..2 * sh.kh + sh.vh) |hd| { + for (rows, xs, 0..) |x, *p, t| { + const src: usize = @intCast(win[r * taps + t]); + for (x, 0..) |*o, i| o.* = if (src < sh.keep) bf(state[(@as(usize, @intCast(sid[r])) * sh.keep + src) * c + hd * sh.dk + i]) else bf(qkv[(src - sh.keep) * ld + hd * sh.dk + i]); + p.* = x; + } + for (0..sh.dk) |i| for (0..taps) |t| { + w[i * taps + t] = bf(conv[(hd * sh.dk + i) * taps + t]); + }; + ref.conv(xs, w, taps, cv, err); + if (hd < 2 * sh.kh) { + const is_q = hd < sh.kh; + const f = ref.norm(cv, is_q, sh.eps, want); + const got = if (is_q) q[(r * sh.kh + hd) * sh.dk ..] else k[(r * sh.kh + hd - sh.kh) * sh.dk ..]; + for (want, err, 0..) |x, e, i| most_q = @max(most_q, @abs(bf(got[i]) - x) / (@abs(x) * rel + f * e + 1e-30)); + } else { + const hv = hd - 2 * sh.kh; + for (cv, err, 0..) |x, e, i| most_q = @max(most_q, @abs(bf(v[(r * sh.vh + hv) * sh.dk + i]) - x) / (@abs(x) * (std.math.pow(f64, 2, -8) + std.math.pow(f64, 2, -20)) + e + 1e-30)); + const x = bf(av[r * ab_ld + hv]) + dtb[hv]; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + const wg = ref.gate(bf(av[r * ab_ld + hv]), dtb[hv], a_log[hv]); + most_q = @max(most_q, @abs(gg[r * sh.vh + hv] - wg) / (@abs(wg) * std.math.pow(f64, 2, -20) * (4 + 4 * @exp(@as(f64, a_log[hv])) * sp) + 1e-30)); + const wb = 1 / (1 + @exp(-bf(bv[r * ab_ld + hv]))); + most_q = @max(most_q, @abs(be[r * sh.vh + hv] - wb) / (wb * std.math.pow(f64, 2, -19) + 1e-30)); + } + }; + worst.* = @max(worst.*, most_q); + try check.expect(most_q <= 1, "{s}: within float64 tolerance (worst {d:.3})", .{ label, most_q }); + for (counts[0 .. counts.len - 1], ios[0 .. counts.len - 1]) |m, o| { + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, u16, o.q, m * sh.kh * sh.dk)), std.mem.sliceAsBytes(q[0 .. m * sh.kh * sh.dk])); + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, u16, o.v, m * sh.vh * sh.dk)), std.mem.sliceAsBytes(v[0 .. m * sh.vh * sh.dk])); + try check.sameBytes(label, std.mem.sliceAsBytes(try b.back(a, f32, o.g, m * sh.vh)), std.mem.sliceAsBytes(gg[0 .. m * sh.vh])); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try cuda.gdn_prep.GdnPrep.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(48); + var worst: f64 = 0; + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 16, .vh = 48, .dk = 128, .keep = 3 }, &worst); // Qwen3.8-27B + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 2, .vh = 3, .dk = 64, .keep = 3 }, &worst); + try case(gpu, arena.allocator(), s, &g, prng.random(), .{ .kh = 1, .vh = 2, .dk = 256, .keep = 1 }, &worst); + check.pass("gdn prep: conv windows, q/k norms, v, g and beta within float64 tolerance (worst {d:.3}), bits at every row count", .{worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..15e70e0162 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,9 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const act_rows_tests = @import("act_rows_tests.zig"); +const gdn_prep_tests = @import("gdn_prep_tests.zig"); +const qk_rows_tests = @import("qk_rows_tests.zig"); const usage = \\usage: tf-cuda-test @@ -59,6 +62,9 @@ const usage = \\ triton a Triton cubin's bits against the Python oracle's fixture \\ norms RMS norm and the residual step against float64 references (synthetic) \\ norm-rows bf16-residual and gated norms, token rows against float64 (synthetic) + \\ act-rows SwiGLU, output gates, group sums, DFlash2's dynamic conv against float64 (synthetic) + \\ gdn-prep Gated DeltaNet conv, norms, g and beta against float64 (synthetic) + \\ qk-rows QK norm, RoPE and the cache write against float64 references (synthetic) \\ sample sample.cu's keyed draws against the Metal rule's host references (synthetic rows) \\ glue Nemotron's glue kernels against their host references at its shapes (synthetic rows) \\ lfm2-kernels LFM2's glue kernels against their host references at d1's shapes (synthetic rows) @@ -164,6 +170,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "act-rows")) return act_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "gdn-prep")) return gdn_prep_tests.run(gpu); + if (std.mem.eql(u8, cmd, "qk-rows")) return qk_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "lfm2-kernels")) return lfm2_kernel_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig index 8f1d295ec1..6610270ace 100644 --- a/zig/tests/cuda/norm_rows_tests.zig +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -6,20 +6,20 @@ const check = @import("check.zig"); const nref = core.norms_ref; const aref = core.act_ref; -const counts = [_]usize{ 1, 3, 17, 64 }; -const most = counts[counts.len - 1]; +pub const counts = [_]usize{ 1, 3, 17, 64 }; +pub const most = counts[counts.len - 1]; const eps: f32 = 1e-6; -fn bf16(x: f32) u16 { +pub fn bf16(x: f32) u16 { const u: u32 = @bitCast(x); return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); } -fn bf(h: u16) f64 { +pub fn bf(h: u16) f64 { return @as(f32, @bitCast(@as(u32, h) << 16)); } -const Rig = struct { +pub const Rig = struct { gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, @@ -27,20 +27,20 @@ const Rig = struct { bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, worst: f64 = 0, - fn dev(r: *Rig, bytes: []const u8) !u64 { + pub fn dev(r: *Rig, bytes: []const u8) !u64 { const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); try r.bufs.append(r.gpu.gpa, b); return b.ptr; } - fn out(r: *Rig, bytes: usize) !u64 { + pub fn out(r: *Rig, bytes: usize) !u64 { const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); try r.bufs.append(r.gpu.gpa, b); try b.fill8(0xff, r.s.handle); // a value the kernel misses reads NaN return b.ptr; } - fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + pub fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { const o = try r.a.alloc(T, n); try r.s.synchronize(); const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; @@ -49,20 +49,20 @@ const Rig = struct { } /// bf16 values: mostly within [-scale, scale), one in 64 twenty times larger. - fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + pub fn bfs(r: *Rig, n: usize, scale: f32) []u16 { const v = r.a.alloc(u16, n) catch unreachable; for (v, 0..) |*x, i| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); return v; } /// Weights near 1. - fn weights(r: *Rig, n: usize) []f32 { + pub fn weights(r: *Rig, n: usize) []f32 { const v = r.a.alloc(f32, n) catch unreachable; for (v) |*x| x.* = 0.5 + r.rng.float(f32); return v; } - fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + pub fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { const q = @abs(got - want) / tol; if (!(q <= 1)) { try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); @@ -71,7 +71,7 @@ const Rig = struct { } /// Group sums `xs` of the stored bf16 rows `y` ([rows, k]). - fn sums(r: *Rig, what: []const u8, y: []const u16, xs: []const f32, rows: usize, k: usize) !void { + pub fn sums(r: *Rig, what: []const u8, y: []const u16, xs: []const f32, rows: usize, k: usize) !void { var row: [64]f64 = undefined; for (0..rows) |i| for (0..k / 64) |j| { for (&row, 0..) |*v, e| v.* = bf(y[i * k + 64 * j + e]); diff --git a/zig/tests/cuda/qk_rows_tests.zig b/zig/tests/cuda/qk_rows_tests.zig new file mode 100644 index 0000000000..cdb5105ef3 --- /dev/null +++ b/zig/tests/cuda/qk_rows_tests.zig @@ -0,0 +1,226 @@ +//! qk_rows.cu against core/qk_ref.zig: q, k within float64 tolerance, v and the caches exact, bits at any row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const counts = [_]usize{ 1, 3, 17, 300 }; +const most = counts[counts.len - 1]; +const positions = 4096; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Case = struct { h: usize, hk: usize, dim: usize, rot: usize, norm: bool }; + +fn run1(gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, g: *const cuda.qk_rows.QkRows, rng: std.Random, c: Case, worst: *f64) !void { + const heads = c.h + 2 * c.hk; + const width = heads * c.dim; + const xs = try a.alloc(u16, most * width); + for (xs, 0..) |*h, i| h.* = bf16(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const nw = try a.alloc(f32, 2 * c.dim); + for (nw) |*w| w.* = 0.5 + rng.float(f32); + const tc = try a.alloc(f32, positions * @max(c.rot, 1)); + const ts = try a.alloc(f32, positions * @max(c.rot, 1)); + for (0..positions) |p| for (0..c.rot) |i| { + const j = i % (c.rot / 2); + const ang = @as(f64, @floatFromInt(p)) / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(c.rot))); + tc[p * c.rot + i] = @floatCast(@cos(ang)); + ts[p * c.rot + i] = @floatCast(@sin(ang)); + }; + const len = 600; + const pos = try a.alloc(i32, most); + const slot = try a.alloc(i32, most); + for (pos, slot, 0..) |*p, *sl, i| { + p.* = @intCast((i * 7919 + 13) % positions); + sl.* = @intCast(i % 2); + } + // distinct (slot, pos % len) a row, so no row's cache entry overwrites another's + for (0..most) |i| pos[i] = @intCast((i / 2) % len + len * (i % 5)); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*b| b.free(); + bufs.deinit(gpu.gpa); + } + const up = struct { + fn f(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(gp.d, bytes); + try l.append(gp.gpa, b); + return b.ptr; + } + fn z(gp: check.Gpu, l: *std.ArrayList(cuda.DeviceBuffer), st: cuda.Stream, n: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(gp.d, n); + try l.append(gp.gpa, b); + try b.fill8(0xff, st.handle); + return b.ptr; + } + }; + const dx = try up.f(gpu, &bufs, std.mem.sliceAsBytes(xs)); + const dnw = try up.f(gpu, &bufs, std.mem.sliceAsBytes(nw)); + const dc = try up.f(gpu, &bufs, std.mem.sliceAsBytes(tc)); + const ds = try up.f(gpu, &bufs, std.mem.sliceAsBytes(ts)); + const dp = try up.f(gpu, &bufs, std.mem.sliceAsBytes(pos)); + const dsl = try up.f(gpu, &bufs, std.mem.sliceAsBytes(slot)); + const shape: cuda.qk_rows.Shape = .{ .norms = if (c.norm) dnw else null, .cos = if (c.rot > 0) dc else null, .sin = if (c.rot > 0) ds else null, .rot = c.rot }; + const cache_bytes = 2 * len * c.hk * c.dim * 2; + const outs = try a.alloc(cuda.qk_rows.Out, counts.len); + for (counts, outs) |m, *o| { + o.* = .{ .q = try up.z(gpu, &bufs, s, m * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, m * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, m * c.hk * c.dim * 2), .kc = try up.z(gpu, &bufs, s, cache_bytes), .vc = try up.z(gpu, &bufs, s, cache_bytes), .len = len }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o.*, m); + } + try s.synchronize(); + const o = outs[counts.len - 1]; + const back = struct { + fn f(gp: check.Gpu, al: std.mem.Allocator, ptr: u64, n: usize) ![]u16 { + const raw = try al.alloc(u16, n); + const b: cuda.DeviceBuffer = .{ .d = gp.d, .ptr = ptr, .len = n * 2 }; + try b.download(0, std.mem.sliceAsBytes(raw)); + return raw; + } + }; + const gq = try back.f(gpu, a, o.q, most * c.h * c.dim); + const gk = try back.f(gpu, a, o.k, most * c.hk * c.dim); + const gv = try back.f(gpu, a, o.v, most * c.hk * c.dim); + const gkc = try back.f(gpu, a, o.kc.?, cache_bytes / 2); + const gvc = try back.f(gpu, a, o.vc.?, cache_bytes / 2); + const xv = try a.alloc(f64, c.dim); + const wq = try a.alloc(f64, c.dim); + const wk = try a.alloc(f64, c.dim); + for (wq, wk, 0..) |*p, *q, i| { + p.* = nw[i]; + q.* = nw[c.dim + i]; + } + const cv = try a.alloc(f64, @max(c.rot, 1)); + const sv = try a.alloc(f64, @max(c.rot, 1)); + const y = try a.alloc(f64, c.dim); + var q_most: f64 = 0; + for (0..most) |r| { + const p: usize = @intCast(pos[r]); + for (0..c.rot) |i| { + cv[i] = tc[p * c.rot + i]; + sv[i] = ts[p * c.rot + i]; + } + const at = @as(usize, @intCast(slot[r])) * len + p % len; + for (0..heads) |hd| { + for (xv, 0..) |*v, i| v.* = bf(xs[r * width + hd * c.dim + i]); + if (hd >= c.h + c.hk) { // v: exact copies, in the output and the cache + const k = hd - c.h - c.hk; + for (0..c.dim) |i| { + try check.expect(gv[(r * c.hk + k) * c.dim + i] == xs[r * width + hd * c.dim + i], "qk rows: v copied", .{}); + try check.expect(gvc[(at * c.hk + k) * c.dim + i] == xs[r * width + hd * c.dim + i], "qk rows: v cached", .{}); + } + continue; + } + const is_q = hd < c.h; + core.qk_ref.head(xv, if (c.norm) (if (is_q) wq else wk) else null, 1e-6, if (c.rot > 0) cv else null, if (c.rot > 0) sv else null, c.rot, y); + for (y, 0..) |want, i| { + const got = if (is_q) bf(gq[(r * c.h + hd) * c.dim + i]) else bf(gk[(r * c.hk + hd - c.h) * c.dim + i]); + const tol = (@abs(want) + 2 * @abs(xv[i]) + 1e-3) * std.math.pow(f64, 2, -8); + q_most = @max(q_most, @abs(got - want) / tol); + if (!is_q) try check.expect(gkc[(at * c.hk + hd - c.h) * c.dim + i] == gk[(r * c.hk + hd - c.h) * c.dim + i], "qk rows: k cached as output", .{}); + } + } + } + worst.* = @max(worst.*, q_most); + // the same heads from [q | gate] rows and a padded k buffer, no v: the same q, k and k cache bits + const qg = try a.alloc(u16, most * c.h * 2 * c.dim); + const kb = try a.alloc(u16, most * (c.hk * c.dim + 64)); + @memset(qg, 0x7fc0); // the gate halves: NaN, never read + @memset(kb, 0x7fc0); + for (0..most) |r| { + for (0..c.h) |hd| @memcpy(qg[(r * c.h + hd) * 2 * c.dim ..][0..c.dim], xs[r * width + hd * c.dim ..][0..c.dim]); + @memcpy(kb[r * (c.hk * c.dim + 64) ..][0 .. c.hk * c.dim], xs[r * width + c.h * c.dim ..][0 .. c.hk * c.dim]); + } + const src: cuda.qk_rows.Src = .{ .q = try up.f(gpu, &bufs, std.mem.sliceAsBytes(qg)), .q_ld = @intCast(c.h * 2 * c.dim), .q_hs = @intCast(2 * c.dim), .k = try up.f(gpu, &bufs, std.mem.sliceAsBytes(kb)), .k_ld = @intCast(c.hk * c.dim + 64), .k_hs = @intCast(c.dim) }; + const o2: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = 0, .kc = try up.z(gpu, &bufs, s, cache_bytes), .len = len }; + try g.rowsFrom(s, src, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o2, most); + try s.synchronize(); + try check.sameBytes("qk rows from strided heads: q", std.mem.sliceAsBytes(try back.f(gpu, a, o2.q, most * c.h * c.dim)), std.mem.sliceAsBytes(gq)); + try check.sameBytes("qk rows from strided heads: k", std.mem.sliceAsBytes(try back.f(gpu, a, o2.k, most * c.hk * c.dim)), std.mem.sliceAsBytes(gk)); + try check.sameBytes("qk rows from strided heads: k cache", std.mem.sliceAsBytes(try back.f(gpu, a, o2.kc.?, cache_bytes / 2)), std.mem.sliceAsBytes(gkc)); + // head-major outputs [heads, rows, dim] (strides): the packed call's bits, rearranged + { + const o4: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .len = len, .q_ld = c.dim, .q_hs = most * c.dim, .kv_ld = c.dim, .kv_hs = most * c.dim }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, shape, dp, dsl, o4, most); + try s.synchronize(); + const hq = try back.f(gpu, a, o4.q, most * c.h * c.dim); + const hk_ = try back.f(gpu, a, o4.k, most * c.hk * c.dim); + const hv = try back.f(gpu, a, o4.v, most * c.hk * c.dim); + for (0..most) |r| { + for (0..c.h) |hd| try check.sameBytes("qk rows head-major: q", std.mem.sliceAsBytes(hq[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gq[(r * c.h + hd) * c.dim ..][0..c.dim])); + for (0..c.hk) |hd| { + try check.sameBytes("qk rows head-major: k", std.mem.sliceAsBytes(hk_[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gk[(r * c.hk + hd) * c.dim ..][0..c.dim])); + try check.sameBytes("qk rows head-major: v", std.mem.sliceAsBytes(hv[(hd * most + r) * c.dim ..][0..c.dim]), std.mem.sliceAsBytes(gv[(r * c.hk + hd) * c.dim ..][0..c.dim])); + } + } + } + // bf16 norm weights at their own addresses and RoPE from inv (angle = pos * inv in fp32): float64 tolerance + if (c.norm) { + const wqb = try a.alloc(u16, c.dim); + const wkb = try a.alloc(u16, c.dim); + for (wqb, wkb, 0..) |*p, *q, i| { + p.* = bf16(nw[i]); + q.* = bf16(nw[c.dim + i]); + } + const half = @max(c.rot / 2, 1); + const inv = try a.alloc(f32, half); + for (inv, 0..) |*v, j| v.* = @floatCast(1 / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(@max(c.rot, 1))))); + const sh3: cuda.qk_rows.Shape = .{ .norms = try up.f(gpu, &bufs, std.mem.sliceAsBytes(wqb)), .norm_k = try up.f(gpu, &bufs, std.mem.sliceAsBytes(wkb)), .norms_bf16 = true, .inv = if (c.rot > 0) try up.f(gpu, &bufs, std.mem.sliceAsBytes(inv)) else null, .rot = c.rot }; + const o3: cuda.qk_rows.Out = .{ .q = try up.z(gpu, &bufs, s, most * c.h * c.dim * 2), .k = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .v = try up.z(gpu, &bufs, s, most * c.hk * c.dim * 2), .len = len }; + try g.rows(s, dx, .{ .h = c.h, .hk = c.hk, .dim = c.dim }, sh3, dp, dsl, o3, most); + try s.synchronize(); + const q3 = try back.f(gpu, a, o3.q, most * c.h * c.dim); + const k3 = try back.f(gpu, a, o3.k, most * c.hk * c.dim); + for (wq, wk, 0..) |*p, *q, i| { + p.* = bf(wqb[i]); + q.* = bf(wkb[i]); + } + for (0..most) |r| { + const p: usize = @intCast(pos[r]); + for (0..c.rot) |i| { + const ang: f64 = @as(f32, @floatFromInt(p)) * inv[i % half]; + cv[i] = @cos(ang); + sv[i] = @sin(ang); + } + for (0..c.h + c.hk) |hd| { + for (xv, 0..) |*v, i| v.* = bf(xs[r * width + hd * c.dim + i]); + const is_q = hd < c.h; + core.qk_ref.head(xv, if (is_q) wq else wk, 1e-6, if (c.rot > 0) cv else null, if (c.rot > 0) sv else null, c.rot, y); + for (y, 0..) |want, i| { + const got = if (is_q) bf(q3[(r * c.h + hd) * c.dim + i]) else bf(k3[(r * c.hk + hd - c.h) * c.dim + i]); + q_most = @max(q_most, @abs(got - want) / ((@abs(want) + 2 * @abs(xv[i]) + 1e-3) * std.math.pow(f64, 2, -8))); + } + } + } + worst.* = @max(worst.*, q_most); + } + try check.expect(q_most <= 1, "qk rows {d}/{d}x{d}, rot {d}{s}: within float64 tolerance (worst {d:.3})", .{ c.h, c.hk, c.dim, c.rot, if (c.norm) ", normed" else "", q_most }); + for (counts[0 .. counts.len - 1], outs[0 .. counts.len - 1]) |m, oo| { + try check.sameBytes("qk rows: q at a smaller row count", std.mem.sliceAsBytes(try back.f(gpu, a, oo.q, m * c.h * c.dim)), std.mem.sliceAsBytes(gq[0 .. m * c.h * c.dim])); + try check.sameBytes("qk rows: k at a smaller row count", std.mem.sliceAsBytes(try back.f(gpu, a, oo.k, m * c.hk * c.dim)), std.mem.sliceAsBytes(gk[0 .. m * c.hk * c.dim])); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try cuda.qk_rows.QkRows.load(gpu.d); + defer g.unload(); + var prng = std.Random.DefaultPrng.init(31); + var worst: f64 = 0; + const a = arena.allocator(); + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .dim = 128, .rot = 128, .norm = true }, &worst); // Kolibri, sliding + try run1(gpu, a, s, &g, prng.random(), .{ .h = 48, .hk = 4, .dim = 128, .rot = 0, .norm = true }, &worst); // Kolibri, full + try run1(gpu, a, s, &g, prng.random(), .{ .h = 8, .hk = 2, .dim = 64, .rot = 32, .norm = false }, &worst); // partial rotary + try run1(gpu, a, s, &g, prng.random(), .{ .h = 16, .hk = 2, .dim = 128, .rot = 128, .norm = false }, &worst); // a warp a head, no norm + try run1(gpu, a, s, &g, prng.random(), .{ .h = 24, .hk = 4, .dim = 256, .rot = 64, .norm = true }, &worst); // Qwen3.8-27B + check.pass("qk rows: q and k within float64 tolerance (worst {d:.3}), v and the caches exact, bits at any row count", .{worst}); +} From bffc7047c8769d4f39ba5ae3ed55a99cf2f0349e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:27:32 +0200 Subject: [PATCH 2/6] cuda: qk_rows and gdn_prep take 1/sqrt by IEEE division and square root, correctly rounded on every architecture --- zig/kernels/cuda/gdn_prep.cu | 4 ++-- zig/kernels/cuda/qk_rows.cu | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index 04f330f4e7..c1af9cb37e 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -63,9 +63,9 @@ __device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int ke #pragma unroll for (int i = 0; i < J; ++i) part = __fmaf_rn(cv[i], cv[i], part); for (int m = 16; m > 0; m >>= 1) part = __fadd_rn(part, __shfl_xor_sync(0xFFFFFFFFu, part, m)); - const float inv = rsqrtf(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps)); + const float inv = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(part, static_cast(DK)), eps))); const bool is_q = hd < kh; - const float scale = is_q ? 1.f / DK : rsqrtf(static_cast(DK)); + const float scale = is_q ? 1.f / DK : __fdiv_rn(1.f, __fsqrt_rn(static_cast(DK))); __nv_bfloat16* out = (is_q ? io.q : io.k) + (r * kh + (is_q ? hd : hd - kh)) * DK + J * lane; #pragma unroll for (int i = 0; i < J; ++i) out[i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(cv[i], inv), scale)); diff --git a/zig/kernels/cuda/qk_rows.cu b/zig/kernels/cuda/qk_rows.cu index 146e14017c..395b1ca677 100644 --- a/zig/kernels/cuda/qk_rows.cu +++ b/zig/kernels/cuda/qk_rows.cu @@ -96,7 +96,7 @@ extern "C" __global__ void __launch_bounds__(MOST) tf_qk_rows(const QkSrc in, co __syncthreads(); float s = 0.f; for (int w = 0; w < dim / 32; ++w) s = __fadd_rn(s, warps[w]); - const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps)); + const float rstd = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(s, static_cast(dim)), eps))); y = __fmul_rn(__fmul_rn(x, rstd), wv(nw, hd >= h, d)); if (d < rot) yp = __fmul_rn(__fmul_rn(__bfloat162float(src[partner]), rstd), wv(nw, hd >= h, partner)); } else if (d < rot) { @@ -166,7 +166,7 @@ extern "C" __global__ void __launch_bounds__(128) tf_qk_rows_w(const QkSrc in, c float s = 0.f; #pragma unroll for (int j = 0; j < J; ++j) s = __fadd_rn(s, part[j]); - const float rstd = rsqrtf(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps)); + const float rstd = __fdiv_rn(1.f, __fsqrt_rn(__fadd_rn(__fdiv_rn(s, static_cast(D)), eps))); #pragma unroll for (int j = 0; j < J; ++j) { y[j] = __fmul_rn(__fmul_rn(x[j], rstd), w[j]); From a97b9c0ff97f2cef42cd97165bc8c0cc1936d900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:55:00 +0200 Subject: [PATCH 3/6] cuda: decode-sized row calls take small blocks (launch.rowThreads, Triton's tuned sizes): swiglu 128 and gate 64 threads, gdn_prep a warp a block pair; the same bits at any row count --- zig/kernels/cuda/act_rows.cu | 8 ++++---- zig/src/cuda/act_rows.zig | 14 ++++++++------ zig/src/cuda/gdn_prep.zig | 3 ++- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index 515cc6bf67..9fc24a05b6 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/zig/kernels/cuda/act_rows.cu @@ -52,7 +52,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_swiglu_rows(const __nv_bfloa __nv_bfloat16* __restrict__ out, float* __restrict__ xs, int n) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values const bool live = 8 * c < n; float g[8], u[8], o[8]; if (live) { @@ -75,7 +75,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat1 long long g_hs, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, int heads, int d) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; // 8 values + const int c = blockIdx.y * blockDim.x + threadIdx.x; // 8 values const int n = heads * d; const bool live = 8 * c < n; const int h = 8 * c / d, e = 8 * c % d; @@ -98,7 +98,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_gate_rows(const __nv_bfloat1 extern "C" __global__ void __launch_bounds__(NT) tf_group_sums_rows(const __nv_bfloat16* __restrict__ x, long long ld, float* __restrict__ xs, int k) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; const bool live = 8 * c < k; const uint4 v = live ? *reinterpret_cast(x + r * ld + 8 * c) : make_uint4(0, 0, 0, 0); const float s = group_sum8(v, live); @@ -113,7 +113,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_dconv_rows(const __nv_bfloat __nv_bfloat16* __restrict__ out, int d, int gs, int branch, int seg) { const int64_t r = blockIdx.x; - const int c = blockIdx.y * NT + threadIdx.x; + const int c = blockIdx.y * blockDim.x + threadIdx.x; if (8 * c >= d) return; const int g = d / gs, grp = 8 * c / gs; float xv[8], pv[8], b0[8], b1[8], y[8]; diff --git a/zig/src/cuda/act_rows.zig b/zig/src/cuda/act_rows.zig index 832844ce01..1d343f85f3 100644 --- a/zig/src/cuda/act_rows.zig +++ b/zig/src/cuda/act_rows.zig @@ -18,8 +18,8 @@ fn aligned(ptrs: []const u64, strides: []const usize) bool { return true; } -fn grid(rows: usize, n: usize) launch_.Dim3 { - return .{ .x = @intCast(rows), .y = @intCast((n / 8 + threads - 1) / threads), .z = 1 }; +fn grid(rows: usize, n: usize, t: u32) launch_.Dim3 { + return .{ .x = @intCast(rows), .y = @intCast((n / 8 + t - 1) / t), .z = 1 }; } pub const ActRows = struct { @@ -52,7 +52,8 @@ pub const ActRows = struct { a.add(out); a.add(xs orelse 0); a.add(@as(i32, @intCast(n))); - try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n), .block = .{ .x = threads } }, s, &a); + const t = launch_.rowThreads(rows, threads, 128); + try launch_.launch(g.swiglu_fn, .{ .grid = grid(rows, n, t), .block = .{ .x = t } }, s, &a); } /// out = bf16(o * sigmoid(gate)), head h's gate at r * g_ld + h * g_hs ([q | gate]: q + d, 2 d); xs: sums. @@ -68,7 +69,8 @@ pub const ActRows = struct { a.add(xs orelse 0); a.add(@as(i32, @intCast(heads))); a.add(@as(i32, @intCast(d))); - try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d), .block = .{ .x = threads } }, s, &a); + const t = launch_.rowThreads(rows, threads, 64); + try launch_.launch(g.gate_fn, .{ .grid = grid(rows, heads * d, t), .block = .{ .x = t } }, s, &a); } /// xs [rows, k / 64]: x's rows (stride ld) summed over each 64 values. k % 64 == 0. @@ -80,7 +82,7 @@ pub const ActRows = struct { a.add(@as(i64, @intCast(ld))); a.add(xs); a.add(@as(i32, @intCast(k))); - try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k), .block = .{ .x = threads } }, s, &a); + try launch_.launch(g.sums_fn, .{ .grid = grid(rows, k, threads), .block = .{ .x = threads } }, s, &a); } /// DFlash2's two-tap conv over [rows, d] bf16: dyn [rows, 2, 2, d / gs], base [2, 2, d]; res (or null) added. @@ -90,6 +92,6 @@ pub const ActRows = struct { var a: launch_.Args = .{}; for ([_]u64{ x, dyn, base, res orelse 0, out }) |v| a.add(v); for ([_]usize{ d, gs, branch, seg }) |v| a.add(@as(i32, @intCast(v))); - try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d), .block = .{ .x = threads } }, s, &a); + try launch_.launch(g.dconv_fn, .{ .grid = grid(rows, d, threads), .block = .{ .x = threads } }, s, &a); } }; diff --git a/zig/src/cuda/gdn_prep.zig b/zig/src/cuda/gdn_prep.zig index d11f869096..641a01cd8a 100644 --- a/zig/src/cuda/gdn_prep.zig +++ b/zig/src/cuda/gdn_prep.zig @@ -67,6 +67,7 @@ pub const GdnPrep = struct { a.add(@as(i32, @intCast(sh.vh))); a.add(@as(i32, @intCast(sh.keep))); a.add(sh.eps); - try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + 3) / 4), .z = 1 }, .block = .{ .x = 128 } }, s, &a); + const t = launch_.rowThreads(rows, 128, 64); // a warp a head either way + try launch_.launch(g.fns[which], .{ .grid = .{ .x = @intCast(rows), .y = @intCast((heads + t / 32 - 1) / (t / 32)), .z = 1 }, .block = .{ .x = t } }, s, &a); } }; From 7bf17c648eb8dba94f559184f738bf56201af480 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:07:48 +0200 Subject: [PATCH 4/6] cuda: act_rows and gdn_prep say why their sigmoid divides exactly --- zig/kernels/cuda/act_rows.cu | 1 + zig/kernels/cuda/gdn_prep.cu | 1 + 2 files changed, 2 insertions(+) diff --git a/zig/kernels/cuda/act_rows.cu b/zig/kernels/cuda/act_rows.cu index 9fc24a05b6..a10b0ed2c9 100644 --- a/zig/kernels/cuda/act_rows.cu +++ b/zig/kernels/cuda/act_rows.cu @@ -27,6 +27,7 @@ __device__ __forceinline__ uint4 pack8(const float* f) { return u; } +// Exactly rounded division on purpose: the same bits on every GPU, where div.full and ex2.approx are not. __device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } // A thread's 8 stored values in order, then the group's 8 lanes in a fixed tree (every lane of the warp calls). diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index c1af9cb37e..05b3f140f6 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -30,6 +30,7 @@ struct Io { float* beta; // [rows, vh] }; +// Exactly rounded division on purpose: the same bits on every GPU, where div.full and ex2.approx are not. __device__ __forceinline__ float sigmoid(float x) { return __fdiv_rn(1.f, __fadd_rn(1.f, expf(-x))); } template From 011c851c9a7671f5bcf14a92b8bc61fa0369b928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:47:35 +0200 Subject: [PATCH 5/6] cuda/gdn_prep: every conv tap's loads issued before any is used, the same bits --- zig/kernels/cuda/gdn_prep.cu | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/zig/kernels/cuda/gdn_prep.cu b/zig/kernels/cuda/gdn_prep.cu index 05b3f140f6..7915bea2f7 100644 --- a/zig/kernels/cuda/gdn_prep.cu +++ b/zig/kernels/cuda/gdn_prep.cu @@ -45,16 +45,22 @@ __device__ __forceinline__ void prep(const Io& io, int c, int kh, int vh, int ke int src[MAXT]; #pragma unroll for (int t = 0; t < MAXT; ++t) src[t] = t <= keep ? io.win[r * (keep + 1) + t] : 0; + __nv_bfloat16 xv[MAXT][J], kv[MAXT][J]; +#pragma unroll + for (int t = 0; t < MAXT; ++t) { // every tap's loads issued before any is used + if (t > keep) break; + const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; +#pragma unroll + for (int i = 0; i < J; ++i) xv[t][i] = x[ch + i], kv[t][i] = io.conv[(ch + i) * (keep + 1) + t]; + } float acc[J]; #pragma unroll for (int i = 0; i < J; ++i) acc[i] = 0.f; #pragma unroll for (int t = 0; t < MAXT; ++t) { if (t > keep) break; - const __nv_bfloat16* x = src[t] < keep ? io.state + (srow + src[t]) * c : io.qkv + (src[t] - keep) * io.qkv_ld; #pragma unroll - for (int i = 0; i < J; ++i) - acc[i] = __fmaf_rn(__bfloat162float(x[ch + i]), __bfloat162float(io.conv[(ch + i) * (keep + 1) + t]), acc[i]); + for (int i = 0; i < J; ++i) acc[i] = __fmaf_rn(__bfloat162float(xv[t][i]), __bfloat162float(kv[t][i]), acc[i]); } float cv[J]; #pragma unroll From 31d3baf48020b992e2e0c3a445475adad94535e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 17 Oct 2026 07:40:22 +0200 Subject: [PATCH 6/6] cuda/act_rows, gdn_prep, qk_rows: registered with the kernel contract --- zig/src/core/contract/registry.zig | 15 ++ zig/tests/cuda/contract_act.zig | 352 +++++++++++++++++++++++++++++ zig/tests/cuda/contract_gdn.zig | 305 +++++++++++++++++++++++++ zig/tests/cuda/contract_qk.zig | 271 ++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 23 +- 5 files changed, 965 insertions(+), 1 deletion(-) create mode 100644 zig/tests/cuda/contract_act.zig create mode 100644 zig/tests/cuda/contract_gdn.zig create mode 100644 zig/tests/cuda/contract_qk.zig diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..565030a2c8 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,21 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows", .symbol = "tf_swiglu_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_swiglu_rows, 16 rows or fewer", .symbol = "tf_swiglu_rows", .image = "act_rows", .threads = 128 }, .usage = .{ .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 35, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_gate_rows", .symbol = "tf_gate_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_gate_rows, 16 rows or fewer", .symbol = "tf_gate_rows", .image = "act_rows", .threads = 64 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_group_sums_rows", .symbol = "tf_group_sums_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "act_rows tf_dconv_rows", .symbol = "tf_dconv_rows", .image = "act_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 }, .{ .regs = 37, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_64", .symbol = "tf_gdn_prep_64", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_64, 16 rows or fewer", .symbol = "tf_gdn_prep_64", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_128", .symbol = "tf_gdn_prep_128", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_128, 16 rows or fewer", .symbol = "tf_gdn_prep_128", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 }, .{ .regs = 77, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_256", .symbol = "tf_gdn_prep_256", .image = "gdn_prep", .threads = 128 }, .usage = .{ .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "gdn_prep tf_gdn_prep_256, 16 rows or fewer", .symbol = "tf_gdn_prep_256", .image = "gdn_prep", .threads = 64 }, .usage = .{ .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 146, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 }, .{ .regs = 145, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows at dim 64", .symbol = "tf_qk_rows", .image = "qk_rows", .threads = 64 }, .usage = .{ .{ .regs = 30, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows at dim 256", .symbol = "tf_qk_rows", .image = "qk_rows", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 36, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 }, .{ .regs = 29, .static = 32, .bound = 256 } } }, + .{ .launch = .{ .name = "qk_rows tf_qk_rows_w", .symbol = "tf_qk_rows_w", .image = "qk_rows", .threads = 128 }, .usage = .{ .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. diff --git a/zig/tests/cuda/contract_act.zig b/zig/tests/cuda/contract_act.zig new file mode 100644 index 0000000000..1f4804f713 --- /dev/null +++ b/zig/tests/cuda/contract_act.zig @@ -0,0 +1,352 @@ +//! act_rows.cu (#797) under the contract on CUDA: SwiGLU, the output gate, 64-value sums and DFlash2's two-tap conv. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const aref = core.act_ref; + +pub const Which = enum { swiglu, swiglu_sums, gate, gate_sums, group_sums, dconv, dconv_res }; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// A thread's 8 values in order, then 8 lanes' butterfly. +fn groupDepth(_: c.Shape) usize { + return 11; +} + +/// x k0 rounded, then the fma that adds the previous row's tap. +fn convDepth(_: c.Shape) usize { + return 2; +} + +const gs = 16; +const tiles_rows: c.Tile = .{ .size = 16 }; // launch.rowThreads: 16 rows or fewer take the small blocks +const tiles_n: c.Tile = .{ .size = 2048, .step = 64, .min = 64 }; // 256 threads of 8 values +const swiglu_served = &[_]c.Shape{ .{ .rows = 1, .n = 17408, .k = 1 }, .{ .rows = 64, .n = 17408, .k = 1 } }; // the 27B's MLP +const gate_served = &[_]c.Shape{ .{ .rows = 1, .n = 24, .k = 256 }, .{ .rows = 64, .n = 24, .k = 256 } }; // its output gate +const conv_served = &[_]c.Shape{ .{ .rows = 16, .n = 5120, .k = 16 }, .{ .rows = 64, .n = 5120, .k = 16 } }; // DFlash2's block + +/// silu and sigmoid: expf (2 ulp, 4 units of 2^-24), the add, the division and the products. +const act_ops = 8; +const res_rule = "own check, core/act_ref's bound: the conv is rounded to bf16 before the residual add, an error of the conv's size"; + +/// swiglu, gate: rows x outputs (gate: heads x dim); sums: rows x 64-value groups; dconv: rows x channels x segment. +pub fn spec(which: Which) c.Spec { + return switch (which) { + .swiglu => .{ .name = "cuda act_rows tf_swiglu_rows", .accum = .none, .depth = noSums, .ops = act_ops, .out = .bf16, .tiles = .{ tiles_rows, tiles_n, .{} }, .served = swiglu_served }, + .swiglu_sums => .{ .name = "cuda act_rows tf_swiglu_rows group sums (of its stored bf16 outputs)", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ tiles_rows, tiles_n, .{} }, .served = swiglu_served }, + .gate => .{ .name = "cuda act_rows tf_gate_rows", .accum = .none, .depth = noSums, .ops = act_ops, .out = .bf16, .tiles = .{ tiles_rows, .{}, .{ .size = 256, .step = 64, .min = 64 } }, .served = gate_served }, + .gate_sums => .{ .name = "cuda act_rows tf_gate_rows group sums (of its stored bf16 outputs)", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ tiles_rows, .{}, .{ .size = 256, .step = 64, .min = 64 } }, .served = gate_served }, + .group_sums => .{ .name = "cuda act_rows tf_group_sums_rows", .accum = .f32, .depth = groupDepth, .out = .f32, .tiles = .{ .{}, .{}, tiles_n }, .served = &.{ .{ .rows = 1, .n = 1, .k = 5120 }, .{ .rows = 64, .n = 1, .k = 17408 } } }, + .dconv => .{ .name = "cuda act_rows tf_dconv_rows", .accum = .f32, .depth = convDepth, .out = .bf16, .tiles = .{ tiles_rows, .{ .size = 2048, .step = gs, .min = gs }, .{} }, .served = conv_served }, + .dconv_res => .{ .name = "cuda act_rows tf_dconv_rows with residual", .accum = .f32, .depth = convDepth, .out = .bf16, .nan = .{ .stated = res_rule }, .tiles = .{ tiles_rows, .{ .size = 2048, .step = gs, .min = gs }, .{} }, .served = conv_served }, + }; +} + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +fn word(v: f32) u16 { + return c.toBf16(v); +} + +/// Gate 1's places: NaN, +inf, -inf, then both infinities in row 3 (rows that exist), in `x` [rows, w]. +fn plant(x: []u16, rows_: usize, w: usize, first: bool) void { + const at = [_][2]usize{ .{ 0, w / 3 }, .{ 1, w / 2 }, .{ 2, w - 1 }, .{ 3, 0 }, .{ 3, w - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put, 0..) |p, v, i| if (p[0] < rows_ and (i < 2) == first) { + x[p[0] * w + p[1]] = word(v); + }; +} + +/// Gate 4's 64-value groups: 16 integers, then their negatives (four nudged by one), twice: big partials in each order. +fn mirrored(rng: std.Random, x: []u16, scale: f32, top: u32) void { + var i: usize = 0; + while (i + 64 <= x.len) : (i += 64) for (0..2) |h| for (0..16) |j| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, top)); + x[i + 32 * h + j] = word(v * scale); + x[i + 32 * h + 16 + j] = word((-v + @as(f32, if (j < 4) 1 else 0)) * scale); + }; +} + +pub const Act = struct { + rig: rig_.Rig, + which: Which, + k: *const cuda.act_rows.ActRows, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + y: []u16 = &.{}, + base: []u16 = &.{}, + dyn: []u16 = &.{}, + res: []u16 = &.{}, + in: [5]u64 = @splat(0), + out: u64 = 0, + xs: u64 = 0, + fill: g.Fill = .random, + ran: bool = false, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Act) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Act) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = if (self.which == .dconv_res) &vt_res else &vt }; + } + + const vt = g.hooks(Act); + const vt_res = blk: { + var v = g.hooks(Act); + v.check = checkRes; + break :blk v; + }; + + /// Outputs a row's call writes: values, or (sums) one a 64-value group. + fn width(self: *const Act) usize { + return switch (self.which) { + .gate, .gate_sums => self.shape.n * self.shape.k, + .group_sums => self.shape.k, + else => self.shape.n, + }; + } + + fn sums(self: *const Act) bool { + return self.which == .swiglu_sums or self.which == .gate_sums or self.which == .group_sums; + } + + fn bfs(self: *Act, n: usize, rng: std.Random, scale: f32) ![]u16 { + const v = try self.arena.allocator().alloc(u16, n); + for (v, 0..) |*x, i| x.* = word((rng.float(f32) * 2 - 1) * scale * (if (i % 64 == 7) @as(f32, 20) else 1)); + return v; + } + + pub fn prepare(self: *Act, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + self.fill = fill; + self.ran = false; + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const w = self.width(); + if (w % 64 != 0 and self.which != .dconv and self.which != .dconv_res) return false; + switch (self.which) { + .swiglu, .swiglu_sums, .gate, .gate_sums => { + // x: the gates ([gate | up] rows, or [q | gate] heads), y: up or o + const gated = self.which == .gate or self.which == .gate_sums; + self.x = try self.bfs(r * 2 * w, rng, 3); + self.y = try self.bfs(r * w, rng, 1); + if (fill == .nonfinite) { + plant(self.y, r, w, false); + for (0..2) |i| if (i < r) { + self.x[i * 2 * w + (if (gated) shape.k else 0) + (if (i == 0) w / 3 else w / 2) % shape.k] = word(if (i == 0) std.math.nan(f32) else std.math.inf(f32)); + }; + } + if (fill == .stress) { + @memset(self.x, word(17)); // silu(17) and sigmoid(17) are 17 and 1 in fp32 + mirrored(rng, self.y, if (gated) 17 else 1, 15); // outputs 17 u: odd, up to 255 + } + const both = try self.arena.allocator().alloc(u16, r * 2 * w); + @memcpy(both, self.x); + // [gate | up]: up after a row's gates; [q | gate] heads: the gates are the second half of each head + if (!gated) for (0..r) |i| @memcpy(both[i * 2 * w + w ..][0..w], self.y[i * w ..][0..w]); + self.in[0] = try self.rig.input(std.mem.sliceAsBytes(both)); + if (gated) self.in[1] = try self.rig.input(std.mem.sliceAsBytes(self.y)); + if (self.sums()) { + self.out = try self.rig.scratch(r * w * 2); + self.xs = try self.rig.output(r * w / 64 * 4); + } else { + self.out = try self.rig.output(r * w * 2); + self.xs = try self.rig.scratch(r * w / 64 * 4); + } + }, + .group_sums => { + self.x = try self.bfs(r * (w + 64), rng, 2); + if (fill == .nonfinite) plant(self.x, r, w + 64, true); + if (fill == .nonfinite) plant(self.x, r, w + 64, false); + if (fill == .stress) for (0..r) |i| mirrored(rng, self.x[i * (w + 64) ..][0..w], 1, 255); + self.in[0] = try self.rig.input(std.mem.sliceAsBytes(self.x)); + self.xs = try self.rig.output(r * w / 64 * 4); + }, + .dconv, .dconv_res => { + if (w % gs != 0) return false; + self.x = try self.bfs(r * w, rng, 1); + self.dyn = try self.bfs(r * 4 * (w / gs), rng, 0.3); + self.base = try self.bfs(4 * w, rng, 0.7); + self.res = try self.bfs(r * w, rng, 1); + if (fill == .nonfinite) { + plant(self.x, r, w, true); + plant(self.x, r, w, false); + } + if (fill == .stress) { + // rows alternate 255 and 254, taps 255 and -256: x k0 + p k1 = 1, which fp16 and bf16 lose + @memset(self.base, 0); + for (0..r) |i| { + @memset(self.x[i * w ..][0..w], word(if (i % 2 == 0) 255 else 254)); + @memset(self.res[i * w ..][0..w], word(0.25)); + for (0..2) |b| for (0..2) |t| { + const v: f32 = if ((i % 2 == 0) == (t == 0)) 255 else -256; + @memset(self.dyn[((i * 2 + b) * 2 + t) * (w / gs) ..][0 .. w / gs], word(v)); + }; + } + } + for ([_][]u16{ self.x, self.dyn, self.base, self.res }, 0..) |v, i| self.in[i] = try self.rig.input(std.mem.sliceAsBytes(v)); + self.out = try self.rig.output(r * w * 2); + }, + } + return true; + } + + fn branch(self: *const Act) usize { + return if (self.which == .dconv_res) 1 else 0; + } + + pub fn run(self: *Act) !void { + const s = self.rig.s; + const r = self.shape.rows; + const w = self.width(); + switch (self.which) { + .swiglu, .swiglu_sums => try self.k.swiglu(s, self.in[0], 2 * w, self.in[0] + 2 * w, 2 * w, self.out, self.xs, r, w), + .gate, .gate_sums => try self.k.outputGate(s, self.in[1], self.in[0] + 2 * self.shape.k, 2 * w, 2 * self.shape.k, self.out, self.xs, r, self.shape.n, self.shape.k), + .group_sums => try self.k.groupSums(s, self.in[0], w + 64, self.xs, r, w), + .dconv, .dconv_res => try self.k.dconv(s, self.in[0], self.in[1], self.in[2], if (self.which == .dconv_res) self.in[3] else null, self.out, r, w, gs, self.branch(), self.shape.k), + } + try s.synchronize(); + self.ran = true; + } + + pub fn count(self: *Act) usize { + const n = self.shape.rows * self.width(); + return if (self.sums()) n / 64 else n; + } + + /// The values a group sums: x, or the bf16 outputs stored (before a run, gate 4's: silu(17) and sigmoid(17) exact). + fn summed(self: *Act) ![]const u16 { + if (self.which == .group_sums) return self.x; + const n = self.shape.rows * self.width(); + const v = try self.arena.allocator().alloc(u16, n); + if (self.ran) { + @memcpy(std.mem.sliceAsBytes(v), try self.rig.bytes(self.out, n * 2)); + return v; + } + if (self.fill != .stress) return error.NoRunYet; + const k = if (self.which == .swiglu_sums) aref.swiglu(17, 1) else aref.sigmoid(17); + for (v, self.y) |*o, y| o.* = word(@floatCast(k * bf(y))); + return v; + } + + fn group(self: *Act, vals: []const u16, o: usize, t: *[64]f64) void { + const w = self.width(); + const ld = if (self.which == .group_sums) w + 64 else w; + const row = o / (w / 64); + const j = o % (w / 64); + for (t, 0..) |*v, e| v.* = bf(vals[row * ld + 64 * j + e]); + } + + /// The conv's two terms at output o, and the residual (0 without). + fn conv(self: *const Act, o: usize, t: *[2]f64) f64 { + const w = self.shape.n; + const row = o / w; + const i = o % w; + const b = self.branch(); + const grp = i / gs; + const tap = struct { + fn f(base: f64, dyn: f64) f64 { + return bf(word(@as(f32, @floatCast(base)) + @as(f32, @floatCast(dyn)))); // k = bf16(base + dyn) in fp32 + } + }.f; + const k0 = tap(bf(self.base[(b * 2) * w + i]), bf(self.dyn[((row * 2 + b) * 2) * (w / gs) + grp])); + const k1 = tap(bf(self.base[(b * 2 + 1) * w + i]), bf(self.dyn[((row * 2 + b) * 2 + 1) * (w / gs) + grp])); + const prev = if (row % self.shape.k > 0) bf(self.x[(row - 1) * w + i]) else 0; + t.* = .{ bf(self.x[row * w + i]) * k0, prev * k1 }; + return if (self.which == .dconv_res) bf(self.res[row * w + i]) else 0; + } + + pub fn reference(self: *Act, want: []f64, mag: []f64) !void { + const w = self.width(); + if (self.sums()) { + const vals = try self.summed(); + var t: [64]f64 = undefined; + for (want, mag, 0..) |*o, *m, i| { + self.group(vals, i, &t); + o.* = 0; + m.* = 0; + for (t) |v| { + o.* += v; + m.* += @abs(v); + } + } + return; + } + for (want, mag, 0..) |*o, *m, i| { + const row = i / w; + const e = i % w; + switch (self.which) { + .swiglu => o.* = aref.swiglu(bf(self.x[row * 2 * w + e]), bf(self.y[i])), + .gate => o.* = bf(self.y[i]) * aref.sigmoid(bf(self.x[row * 2 * w + (e / self.shape.k) * 2 * self.shape.k + self.shape.k + e % self.shape.k])), + .dconv, .dconv_res => { + var t: [2]f64 = undefined; + const res = self.conv(i, &t); + o.* = res + aref.dconv(1, 1, t[0], t[1]); + m.* = @abs(t[0]) + @abs(t[1]); + continue; + }, + else => unreachable, + } + m.* = @abs(o.*); + } + } + + pub fn simulate(self: *Act, acc: c.Accum, order: g.Order, out: []f64) !void { + if (self.sums()) { + const vals = try self.summed(); + var t: [64]f64 = undefined; + for (out, 0..) |*o, i| { + self.group(vals, i, &t); + o.* = rows.sumIn(acc, order, &t); + } + return; + } + if (self.which != .dconv and self.which != .dconv_res) return error.NoSums; + var t: [2]f64 = undefined; + for (out, 0..) |*o, i| { + const res = self.conv(i, &t); + o.* = res + rows.sumIn(acc, order, &t); + } + } + + /// core/act_ref's dconv bound with the residual: half a bf16 ulp of the output and of the conv, 2^-22 of its terms. + fn checkRes(ctx: *anyopaque, at: g.Gate) anyerror!?[]const u8 { + const self: *Act = @ptrCast(@alignCast(ctx)); + const out = try self.rig.outputBytes(); + const half = std.math.pow(f64, 2, -8); + var planted: usize = 0; + const n = self.count(); + var t: [2]f64 = undefined; + for (0..n) |i| { + const res = self.conv(i, &t); + const y = t[0] + t[1]; + const want = res + y; + const got = c.read(.bf16, out, i); + if (c.classOf(want) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(want)) return try std.fmt.bufPrint(&self.msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(got), c.classOf(want) }); + if (at == .nonfinite or c.classOf(want) != .finite) continue; + const tol = @abs(want) * half + @abs(y) * half + (@abs(t[0]) + @abs(t[1])) * std.math.pow(f64, 2, -22) + 1e-30; + if (!(@abs(got - want) <= tol)) return try std.fmt.bufPrint(&self.msg, "output {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of its bound", .{ i, got, want, @abs(got - want) / tol }); + } + if (at == .nonfinite and (planted == 0 or planted == n)) return try std.fmt.bufPrint(&self.msg, "the planted values make {d} of {d} outputs non-finite: no proof", .{ planted, n }); + return null; + } +}; diff --git a/zig/tests/cuda/contract_gdn.zig b/zig/tests/cuda/contract_gdn.zig new file mode 100644 index 0000000000..7c9653463e --- /dev/null +++ b/zig/tests/cuda/contract_gdn.zig @@ -0,0 +1,305 @@ +//! gdn_prep.cu (#797) under the contract on CUDA: a row's conv window, SiLU, the q/k norms, v, g and beta. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const ref = core.gdn_prep_ref; +const gp = cuda.gdn_prep; + +/// Which outputs an entry reads (the others are working memory): q and k, v, or g and beta. +pub const Part = enum { qk, v, gb }; + +/// The conv's taps (one fma chain). +fn taps(s: c.Shape) usize { + return s.k; +} + +fn noSums(_: c.Shape) usize { + return 0; +} + +const qk_why = "q and k are normed rows: a bf16 conv or square sum moves them by about an output ulp (fp16 overflows)"; +const v_why = "v is the conv's bf16 rounding: a bf16 sum of its few taps is within an output ulp or two (fp16 overflows)"; +const gb_rule = "own check, the PR's float64 bound: g = exp(-y) carries y times y's relative error, beyond a fixed op count"; + +/// Shape: rows x key heads (n; 3 v heads a key head) x conv taps (k). +pub fn spec(dk: usize, part: Part) c.Spec { + const served: []const c.Shape = switch (dk) { + 128 => &.{ .{ .rows = 1, .n = 16, .k = 4 }, .{ .rows = 64, .n = 16, .k = 4 } }, // Qwen3.8-27B: 16 key heads, 4 taps + 64 => &.{ .{ .rows = 1, .n = 2, .k = 4 }, .{ .rows = 40, .n = 2, .k = 4 } }, + else => &.{ .{ .rows = 1, .n = 1, .k = 2 }, .{ .rows = 40, .n = 1, .k = 2 } }, + }; + const tiles: [3]c.Tile = .{ .{ .size = 16 }, .{ .size = 4 }, .{ .size = 8, .min = 2 } }; // small blocks to 16 rows; 4 heads a block + const name: []const u8 = switch (part) { + inline else => |p| switch (dk) { + 64 => "cuda gdn_prep tf_gdn_prep_64 " ++ @tagName(p), + 128 => "cuda gdn_prep tf_gdn_prep_128 " ++ @tagName(p), + else => "cuda gdn_prep tf_gdn_prep_256 " ++ @tagName(p), + }, + }; + return switch (part) { + // silu 8, the norm's sum and rsqrt 12, the scale and products 4; c rounded to bf16, and the norm through it + .qk => .{ .name = name, .accum = .f32, .depth = taps, .ops = 24, .mids = 2, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .bf16 = 0, .why = qk_why } }, + .v => .{ .name = name, .accum = .f32, .depth = taps, .ops = 8, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .bf16 = 0, .why = v_why } }, + .gb => .{ .name = name, .accum = .none, .depth = noSums, .out = .f32, .nan = .{ .stated = gb_rule }, .tiles = tiles, .served = served }, + }; +} + +const vh_per = 3; + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +pub const Gdn = struct { + rig: rig_.Rig, + dk: usize, + part: Part, + k: *const gp.GdnPrep, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + sh: gp.Shape = .{ .kh = 0, .vh = 0, .dk = 0, .keep = 0 }, + qkv: []u16 = &.{}, + state: []u16 = &.{}, + conv: []u16 = &.{}, + av: []u16 = &.{}, + bv: []u16 = &.{}, + a_log: []f32 = &.{}, + dtb: []f32 = &.{}, + sid: []i32 = &.{}, + win: []i32 = &.{}, + io: gp.Io = undefined, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Gdn) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Gdn) g.Kernel { + return .{ .spec = spec(self.dk, self.part), .ctx = self, .vt = if (self.part == .gb) &vt_gb else &vt }; + } + + const vt = g.hooks(Gdn); + const vt_gb = blk: { + var v = g.hooks(Gdn); + v.check = checkGb; + break :blk v; + }; + + fn chans(self: *const Gdn) usize { + return (2 * self.sh.kh + self.sh.vh) * self.dk; + } + + fn ld(self: *const Gdn) usize { + return self.chans() + 64; + } + + pub fn prepare(self: *Gdn, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + if (shape.rows == 0 or shape.n == 0 or shape.k < 2 or shape.k > 8) return false; + self.shape = shape; + self.sh = .{ .kh = shape.n, .vh = vh_per * shape.n, .dk = self.dk, .keep = shape.k - 1 }; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const cc = self.chans(); + const keep = self.sh.keep; + const vh = self.sh.vh; + self.qkv = try a.alloc(u16, r * self.ld()); + self.state = try a.alloc(u16, 2 * keep * cc); + self.conv = try a.alloc(u16, cc * shape.k); + self.av = try a.alloc(u16, r * (vh + 3)); + self.bv = try a.alloc(u16, r * (vh + 3)); + for ([_][]u16{ self.qkv, self.state, self.conv, self.av, self.bv }, [_]f32{ 1, 1, 0.4, 2, 2 }) |v, sc| for (v) |*x| { + x.* = c.toBf16(rng.floatNorm(f32) * sc); + }; + self.a_log = try a.alloc(f32, vh); + self.dtb = try a.alloc(f32, vh); + for (self.a_log, self.dtb) |*x, *y| { + x.* = rng.float(f32) * 2 - 1; + y.* = rng.floatNorm(f32); + } + self.sid = try a.alloc(i32, r); + self.win = try a.alloc(i32, r * shape.k); + for (self.sid, 0..) |*x, i| x.* = @intCast(i % 2); + for (self.win, 0..) |*w, i| { // kept rows or any qkv row (a tree's window is any rows), the row itself last + w.* = if (rng.uintLessThan(u32, 3) == 0) @intCast(rng.uintLessThan(usize, keep)) else @intCast(keep + rng.uintLessThan(usize, r)); + if (i % shape.k == shape.k - 1) w.* = @intCast(keep + i / shape.k); + } + if (fill == .nonfinite) { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + const at = [_][2]usize{ .{ 0, cc / 3 }, .{ 1, cc / 2 }, .{ 2, cc - 1 }, .{ 3, 0 }, .{ 3, cc - 1 } }; + for (at, put) |p, v| if (p[0] < r) { + self.qkv[p[0] * self.ld() + p[1]] = c.toBf16(v); + }; + for ([_][]u16{ self.av, self.bv }, 0..) |v, i| if (i < r) { + v[i * (vh + 3) + i % vh] = c.toBf16(put[i]); + }; + } + if (fill == .stress) { + // every input 256, taps 255 but the last one 250 to 254: big integer sums, fp16 overflows past 65504 + @memset(self.qkv, c.toBf16(256)); + @memset(self.state, c.toBf16(256)); + for (0..cc) |i| for (0..shape.k) |t| { + self.conv[i * shape.k + t] = c.toBf16(if (t + 1 < shape.k) 255 else @floatFromInt(250 + i % 5)); + }; + } + self.io = .{ .qkv = try self.rig.input(std.mem.sliceAsBytes(self.qkv)), .qkv_ld = @intCast(self.ld()), .state = try self.rig.input(std.mem.sliceAsBytes(self.state)), .conv = try self.rig.input(std.mem.sliceAsBytes(self.conv)), .win = try self.rig.input(std.mem.sliceAsBytes(self.win)), .sid = try self.rig.input(std.mem.sliceAsBytes(self.sid)), .a = try self.rig.input(std.mem.sliceAsBytes(self.av)), .a_ld = @intCast(vh + 3), .b = try self.rig.input(std.mem.sliceAsBytes(self.bv)), .b_ld = @intCast(vh + 3), .a_log = try self.rig.input(std.mem.sliceAsBytes(self.a_log)), .dt_bias = try self.rig.input(std.mem.sliceAsBytes(self.dtb)), .q = 0, .k = 0, .v = 0, .g = 0, .beta = 0 }; + const qk_bytes = r * self.sh.kh * self.dk * 2; + const kinds = [5]bool{ self.part == .qk, self.part == .qk, self.part == .v, self.part == .gb, self.part == .gb }; + const sizes = [5]usize{ qk_bytes, qk_bytes, r * vh * self.dk * 2, r * vh * 4, r * vh * 4 }; + var ptrs: [5]u64 = undefined; + for (kinds, sizes, &ptrs) |out, n, *p| p.* = if (out) try self.rig.output(n) else try self.rig.scratch(n); + self.io.q = ptrs[0]; + self.io.k = ptrs[1]; + self.io.v = ptrs[2]; + self.io.g = ptrs[3]; + self.io.beta = ptrs[4]; + return true; + } + + pub fn run(self: *Gdn) !void { + try self.k.run(self.rig.s, self.io, self.shape.rows, self.sh); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Gdn) usize { + const r = self.shape.rows; + return switch (self.part) { + .qk => 2 * r * self.sh.kh * self.dk, + .v => r * self.sh.vh * self.dk, + .gb => 2 * r * self.sh.vh, + }; + } + + /// Output o's row, head (q heads, k heads, then v heads) and channel. + fn locate(self: *const Gdn, o: usize) [3]usize { + const r = self.shape.rows; + const kh = self.sh.kh; + const dk = self.dk; + if (self.part == .v) return .{ o / (self.sh.vh * dk), 2 * kh + o / dk % self.sh.vh, o % dk }; + const k_half = o >= r * kh * dk; + const i = if (k_half) o - r * kh * dk else o; + return .{ i / (kh * dk), (if (k_half) kh else 0) + i / dk % kh, i % dk }; + } + + /// Head hd of row r: each channel's tap terms into t [dk][taps]. + fn terms(self: *const Gdn, r: usize, hd: usize, t: []f64) void { + const n = self.shape.k; + const keep = self.sh.keep; + for (0..n) |tp| { + const src: usize = @intCast(self.win[r * n + tp]); + for (0..self.dk) |i| { + const ch = hd * self.dk + i; + const x = if (src < keep) bf(self.state[(@as(usize, @intCast(self.sid[r])) * keep + src) * self.chans() + ch]) else bf(self.qkv[(src - keep) * self.ld() + ch]); + t[i * n + tp] = x * bf(self.conv[ch * n + tp]); + } + } + } + + /// Head hd of row r's outputs and magnitudes (dk each), its conv sums in `acc` and `order` (f32 none: float64). + fn head(self: *Gdn, r: usize, hd: usize, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const a = self.arena.allocator(); + const n = self.shape.k; + const t = try a.alloc(f64, self.dk * n); + const cv = try a.alloc(f64, self.dk); + self.terms(r, hd, t); + for (cv, mag, 0..) |*o, *m, i| { + const row = t[i * n ..][0..n]; + var s: f64 = 0; + m.* = 0; + for (row) |v| { + s += v; + m.* += 1.1 * @abs(v); // silu' is at most 1.1 + } + if (acc) |ac| s = rows.sumIn(ac, order, row); + o.* = s / (1 + @exp(-s)); + } + if (hd >= 2 * self.sh.kh) { + @memcpy(want, cv); + return; + } + const is_q = hd < self.sh.kh; + const dkf: f64 = @floatFromInt(self.dk); + var f = ref.norm(cv, is_q, self.sh.eps, want); + if (acc) |ac| { + const sq = try a.alloc(f64, self.dk); + for (sq, cv) |*q, v| q.* = v * v; + f = 1 / @sqrt(rows.sumIn(ac, order, sq) / dkf + self.sh.eps) * (if (is_q) 1 / dkf else 1 / @sqrt(dkf)); + for (want, cv) |*o, v| o.* = v * f; + } + for (mag) |*m| m.* *= @abs(f); + } + + fn outputs(self: *Gdn, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const a = self.arena.allocator(); + const hw = try a.alloc(f64, self.dk); + const hm = try a.alloc(f64, self.dk); + var o: usize = 0; + while (o < want.len) : (o += self.dk) { + const at = self.locate(o); + try self.head(at[0], at[1], acc, order, hw, hm); + @memcpy(want[o..][0..self.dk], hw); + @memcpy(mag[o..][0..self.dk], hm); + } + } + + pub fn reference(self: *Gdn, want: []f64, mag: []f64) !void { + if (self.part == .gb) { + for (want, mag, 0..) |*w, *m, i| { + w.* = self.gbWant(i); + m.* = @abs(w.*); + } + return; + } + return self.outputs(null, .sequential, want, mag); + } + + pub fn simulate(self: *Gdn, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.arena.allocator().alloc(f64, out.len); + return self.outputs(acc, order, out, mag); + } + + /// g (first rows x vh) then beta, in float64. + fn gbWant(self: *const Gdn, i: usize) f64 { + const vh = self.sh.vh; + const half = self.shape.rows * vh; + const j = i % half; + const r = j / vh; + const hv = j % vh; + if (i < half) return ref.gate(bf(self.av[r * (vh + 3) + hv]), self.dtb[hv], self.a_log[hv]); + return 1 / (1 + @exp(-bf(self.bv[r * (vh + 3) + hv]))); + } + + /// gdn_prep_tests' bounds: g within 2^-20 (4 + 4 exp(a_log) softplus) of itself, beta within 2^-19. + fn checkGb(ctx: *anyopaque, at: g.Gate) anyerror!?[]const u8 { + const self: *Gdn = @ptrCast(@alignCast(ctx)); + const out = try self.rig.outputBytes(); + const vh = self.sh.vh; + const half = self.shape.rows * vh; + var planted: usize = 0; + const n = self.count(); + for (0..n) |i| { + const want = self.gbWant(i); + const got = c.read(.f32, out, i); + if (c.classOf(want) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(want)) return try std.fmt.bufPrint(&self.msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(got), c.classOf(want) }); + if (at == .nonfinite or c.classOf(want) != .finite) continue; + const j = i % half; + const x = bf(self.av[j / vh * (vh + 3) + j % vh]) + self.dtb[j % vh]; + const sp = if (x > 20) x else std.math.log1p(@exp(x)); + const tol = if (i < half) @abs(want) * std.math.pow(f64, 2, -20) * (4 + 4 * @exp(@as(f64, self.a_log[j % vh])) * sp) + 1e-30 else want * std.math.pow(f64, 2, -19) + 1e-30; + if (!(@abs(got - want) <= tol)) return try std.fmt.bufPrint(&self.msg, "output {d} is {e:.6} where float64 gives {e:.6}: {d:.2} of its bound", .{ i, got, want, @abs(got - want) / tol }); + } + if (at == .nonfinite and (planted == 0 or planted == n)) return try std.fmt.bufPrint(&self.msg, "the planted values make {d} of {d} outputs non-finite: no proof", .{ planted, n }); + return null; + } +}; diff --git a/zig/tests/cuda/contract_qk.zig b/zig/tests/cuda/contract_qk.zig new file mode 100644 index 0000000000..137d24a2f4 --- /dev/null +++ b/zig/tests/cuda/contract_qk.zig @@ -0,0 +1,271 @@ +//! qk_rows.cu (#797) under the contract on CUDA: QK norm, rotate-half RoPE and v, against core/qk_ref.zig. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows = contract.rows; +const qk = cuda.qk_rows; + +pub const Norm = enum { none, f32w, bf16w }; +pub const Rope = enum { none, table, inv }; + +/// One registered form: head dim, query heads a key head, rotated dims, norm weights and where cos and sin come from. +pub const Form = struct { name: []const u8, dim: usize, group: usize, rot: usize, norm: Norm, rope: Rope, h: usize }; + +/// Kolibri (48/4 x 128, sliding and full), a warp a head without the norm, the 27B (24/4 x 256, rotary 64), dim 64. +pub const forms = [_]Form{ + .{ .name = "cuda qk_rows tf_qk_rows_w (norm, rotation from a table)", .dim = 128, .group = 12, .rot = 128, .norm = .f32w, .rope = .table, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (norm, no rotation)", .dim = 128, .group = 12, .rot = 0, .norm = .f32w, .rope = .none, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (bf16 norm, rotation from inv_freq)", .dim = 128, .group = 12, .rot = 128, .norm = .bf16w, .rope = .inv, .h = 48 }, + .{ .name = "cuda qk_rows tf_qk_rows_w (no norm, rotation from a table)", .dim = 128, .group = 8, .rot = 128, .norm = .none, .rope = .table, .h = 16 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 256, norm, rotary 64 from a table)", .dim = 256, .group = 6, .rot = 64, .norm = .f32w, .rope = .table, .h = 24 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 256, bf16 norm, rotary 64 from inv_freq)", .dim = 256, .group = 6, .rot = 64, .norm = .bf16w, .rope = .inv, .h = 24 }, + .{ .name = "cuda qk_rows tf_qk_rows (dim 64, no norm, rotary 32 from a table)", .dim = 64, .group = 4, .rot = 32, .norm = .none, .rope = .table, .h = 8 }, +}; + +/// A rotated value: x rstd, times w, times cos (cosf and sinf: 2 ulp, 4 units more), then the add; unrotated fewer. +fn depthTable(_: c.Shape) usize { + return 4; +} + +fn depthInv(_: c.Shape) usize { + return 8; +} + +fn depthBare(_: c.Shape) usize { + return 2; +} + +const norm_why = "a normed row follows its sum of squares as ss^-1/2: a bf16 sum moves it by about an output ulp (fp16 overflows)"; + +/// Shape: rows x query heads (n) x head dim (k). +pub fn spec(f: Form) c.Spec { + const served = switch (f.dim) { + 128 => if (f.h == 48) &[_]c.Shape{ .{ .rows = 1, .n = 48, .k = 128 }, .{ .rows = 300, .n = 48, .k = 128 } } else &[_]c.Shape{ .{ .rows = 1, .n = 16, .k = 128 }, .{ .rows = 300, .n = 16, .k = 128 } }, + 256 => &[_]c.Shape{ .{ .rows = 1, .n = 24, .k = 256 }, .{ .rows = 300, .n = 24, .k = 256 } }, + else => &[_]c.Shape{ .{ .rows = 1, .n = 8, .k = 64 }, .{ .rows = 300, .n = 8, .k = 64 } }, + }; + const depth: *const fn (c.Shape) usize = if (f.norm == .none) &depthBare else if (f.rope == .inv) &depthInv else &depthTable; + // the norm's sum of squares (14 deep, half of it in rsqrt), the division, eps, the root and the reciprocal + const ops: f64 = if (f.norm == .none) 0 else 12; + const sums: c.Sums = if (f.norm == .none) .{} else .{ .bf16 = 0, .why = norm_why }; + return .{ .name = f.name, .accum = .f32, .depth = depth, .ops = ops, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{}, .{} }, .served = served, .sums = sums }; +} + +const positions = 4096; +const len = 600; + +fn bf(h: u16) f64 { + return c.fromBf16(h); +} + +pub const Qk = struct { + rig: rig_.Rig, + form: Form, + k: *const qk.QkRows, + arena: std.heap.ArenaAllocator, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + w: []f64 = &.{}, + cos: []f32 = &.{}, + sin: []f32 = &.{}, + inv: []f32 = &.{}, + pos: []i32 = &.{}, + dx: u64 = 0, + sh: qk.Shape = .{}, + dpos: u64 = 0, + dslot: u64 = 0, + out: qk.Out = undefined, + + pub fn deinit(self: *Qk) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Qk) g.Kernel { + return .{ .spec = spec(self.form), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Qk); + + fn hk(self: *const Qk) usize { + return self.shape.n / self.form.group; + } + + fn heads(self: *const Qk) usize { + return self.shape.n + 2 * self.hk(); + } + + pub fn prepare(self: *Qk, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + const f = self.form; + if (shape.rows == 0 or shape.k != f.dim or shape.n % f.group != 0 or shape.n == 0) return false; + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const r = shape.rows; + const width = self.heads() * f.dim; + self.x = try a.alloc(u16, r * width); + for (self.x, 0..) |*h, i| h.* = c.toBf16(rng.floatNorm(f32) * (if (i % 37 == 0) @as(f32, 8) else 1)); + const rot = @max(f.rot, 2); + self.cos = try a.alloc(f32, positions * rot); + self.sin = try a.alloc(f32, positions * rot); + for (0..positions) |p| for (0..rot) |i| { + const ang = @as(f64, @floatFromInt(p)) / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * (i % (rot / 2)))) / @as(f64, @floatFromInt(rot))); + self.cos[p * rot + i] = @floatCast(@cos(ang)); + self.sin[p * rot + i] = @floatCast(@sin(ang)); + }; + self.inv = try a.alloc(f32, rot / 2); + for (self.inv, 0..) |*v, j| v.* = @floatCast(1 / std.math.pow(f64, 10000, @as(f64, @floatFromInt(2 * j)) / @as(f64, @floatFromInt(rot)))); + const wf = try a.alloc(f32, 2 * f.dim); + for (wf) |*v| v.* = 0.5 + rng.float(f32); + if (fill == .nonfinite) { + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + const at = [_][2]usize{ .{ 0, width / 3 }, .{ 1, width / 2 }, .{ 2, width - 1 }, .{ 3, 0 }, .{ 3, f.dim - 1 } }; + for (at, put) |p, v| if (p[0] < r) { + self.x[p[0] * width + p[1]] = c.toBf16(v); + }; + } + if (fill == .stress) { + if (f.norm == .none) { + // x c - x' s with c, s = 65535, 65534, x = 255: products near 2^24, exact in fp32, a difference of 255 + @memset(self.x, c.toBf16(255)); + @memset(self.cos, 65535); + @memset(self.sin, 65534); + } else { + // integers 32 to 64 of either sign: squares sum past fp16's range, exactly in fp32; cos 1, sin 0 + for (self.x) |*h| h.* = c.toBf16(@as(f32, @floatFromInt(32 + rng.uintLessThan(u32, 33))) * (if (rng.boolean()) @as(f32, 1) else -1)); + @memset(self.cos, 1); + @memset(self.sin, 0); + @memset(self.inv, 0); + } + } + self.w = try a.alloc(f64, 2 * f.dim); + const wb = try a.alloc(u16, 2 * f.dim); + for (self.w, wf, wb) |*v, x, *b| { + b.* = c.toBf16(x); + v.* = if (f.norm == .bf16w) bf(b.*) else x; + } + self.pos = try a.alloc(i32, r); + const slot = try a.alloc(i32, r); + for (self.pos, slot, 0..) |*p, *sl, i| { // distinct (slot, pos % len) a row: no row's cache entry overwrites another's + p.* = @intCast((i / 2) % len + len * (i % 5)); + sl.* = @intCast(i % 2); + } + self.dx = try self.rig.input(std.mem.sliceAsBytes(self.x)); + const nw: ?u64 = switch (f.norm) { + .none => null, + .f32w => try self.rig.input(std.mem.sliceAsBytes(wf)), + .bf16w => try self.rig.input(std.mem.sliceAsBytes(wb)), + }; + const tc = try self.rig.input(std.mem.sliceAsBytes(self.cos)); + const ts = try self.rig.input(std.mem.sliceAsBytes(self.sin)); + const ti = try self.rig.input(std.mem.sliceAsBytes(self.inv)); + self.sh = .{ .norms = nw, .norms_bf16 = f.norm == .bf16w, .cos = if (f.rope == .table) tc else null, .sin = if (f.rope == .table) ts else null, .inv = if (f.rope == .inv) ti else null, .rot = f.rot }; + self.dpos = try self.rig.input(std.mem.sliceAsBytes(self.pos)); + self.dslot = try self.rig.input(std.mem.sliceAsBytes(slot)); + const kv = r * self.hk() * f.dim * 2; + const cache = 2 * len * self.hk() * f.dim * 2; + self.out = .{ .q = try self.rig.output(r * shape.n * f.dim * 2), .k = try self.rig.output(kv), .v = try self.rig.output(kv), .kc = try self.rig.scratch(cache), .vc = try self.rig.scratch(cache), .len = len }; + return true; + } + + pub fn run(self: *Qk) !void { + try self.k.rows(self.rig.s, self.dx, .{ .h = self.shape.n, .hk = self.hk(), .dim = self.form.dim }, self.sh, self.dpos, self.dslot, self.out, self.shape.rows); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Qk) usize { + return self.shape.rows * self.heads() * self.form.dim; + } + + /// Output o's row and head (q heads, k heads, v heads as the outputs lie: all q, then all k, then all v). + fn locate(self: *const Qk, o: usize) [2]usize { + const r = self.shape.rows; + const h = self.shape.n; + const d = self.form.dim; + const i = o / d; + if (i < r * h) return .{ i / h, i % h }; + const j = i - r * h; + const hk_ = self.hk(); + if (j < r * hk_) return .{ j / hk_, h + j % hk_ }; + return .{ (j - r * hk_) / hk_, h + hk_ + (j - r * hk_) % hk_ }; + } + + /// One head's outputs and term magnitudes; its sums in `acc` and `order` (null: float64, core/qk_ref.head). + fn head(self: *Qk, r: usize, hd: usize, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const f = self.form; + const d = f.dim; + const a = self.arena.allocator(); + const x = try a.alloc(f64, d); + for (x, 0..) |*v, i| v.* = bf(self.x[(r * self.heads() + hd) * d + i]); + if (hd >= self.shape.n + self.hk()) { + @memcpy(want, x); + for (mag, x) |*m, v| m.* = @abs(v); + return; + } + const w: ?[]const f64 = if (f.norm == .none) null else self.w[(if (hd >= self.shape.n) d else 0)..][0..d]; + const rot = @max(f.rot, 2); + const p: usize = @intCast(self.pos[r]); + const cv = try a.alloc(f64, rot); + const sv = try a.alloc(f64, rot); + for (0..rot) |i| { + if (f.rope == .inv) { + const ang: f64 = @as(f32, @floatFromInt(p)) * self.inv[i % (rot / 2)]; // the angle in fp32, as served + cv[i] = @cos(ang); + sv[i] = @sin(ang); + } else { + cv[i] = self.cos[p * rot + i]; + sv[i] = self.sin[p * rot + i]; + } + } + const roped = f.rope != .none; + var ss: f64 = 0; + const sq = try a.alloc(f64, d); + for (sq, x) |*q, v| { + q.* = v * v; + ss += q.*; + } + if (acc) |ac| ss = rows.sumIn(ac, order, sq); + const inv = if (w != null) 1 / @sqrt(ss / @as(f64, @floatFromInt(d)) + 1e-6) else 1; + const z = try a.alloc(f64, d); + for (z, x, 0..) |*o, v, i| o.* = v * inv * (if (w) |ww| ww[i] else 1); + if (acc == null) core.qk_ref.head(x, w, 1e-6, if (roped) cv else null, if (roped) sv else null, f.rot, want); + const half = f.rot / 2; + for (want, mag, 0..) |*o, *m, i| { + if (!roped or i >= f.rot) { + if (acc != null) o.* = z[i]; + m.* = @abs(z[i]); + continue; + } + const pi = if (i < half) i + half else i - half; + const t = [2]f64{ z[i] * cv[i], (if (i < half) -z[pi] else z[pi]) * sv[i] }; + if (acc) |ac| o.* = rows.sumIn(ac, order, &t); + m.* = @abs(t[0]) + @abs(t[1]); + } + } + + fn outputs(self: *Qk, acc: ?c.Accum, order: g.Order, want: []f64, mag: []f64) !void { + const d = self.form.dim; + var o: usize = 0; + while (o < want.len) : (o += d) { + const at = self.locate(o); + try self.head(at[0], at[1], acc, order, want[o..][0..d], mag[o..][0..d]); + } + } + + pub fn reference(self: *Qk, want: []f64, mag: []f64) !void { + return self.outputs(null, .sequential, want, mag); + } + + pub fn simulate(self: *Qk, acc: c.Accum, order: g.Order, out: []f64) !void { + const mag = try self.arena.allocator().alloc(f64, out.len); + return self.outputs(acc, order, out, mag); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e87a2e90e5..bfb343528c 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -8,6 +8,9 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const rank_k = @import("contract_rank.zig"); +const act_k = @import("contract_act.zig"); +const gdn_k = @import("contract_gdn.zig"); +const qk_k = @import("contract_qk.zig"); const contract = core.contract; const g = contract.gates; @@ -21,7 +24,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "act_rows", "gdn_prep", "qk_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } @@ -89,6 +92,24 @@ pub fn run(gpu: check.Gpu) !void { var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var act_rows = try cuda.act_rows.ActRows.load(gpu.d); + defer act_rows.unload(); + var acts: [std.enums.values(act_k.Which).len]act_k.Act = undefined; + for (std.enums.values(act_k.Which), &acts) |w, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .k = &act_rows, .arena = .init(a) }; + defer for (&acts) |*x| x.deinit(); + for (&acts) |*x| try kernels.append(a, x.kernel()); + var gdn_prep = try cuda.gdn_prep.GdnPrep.load(gpu.d); + defer gdn_prep.unload(); + var gdns: [9]gdn_k.Gdn = undefined; + for (&gdns, 0..) |*x, i| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .dk = @as(usize, 64) << @intCast(i / 3), .part = @fromBackingInt(@intCast(i % 3)), .k = &gdn_prep, .arena = .init(a) }; + defer for (&gdns) |*x| x.deinit(); + for (&gdns) |*x| try kernels.append(a, x.kernel()); + var qk_rows = try cuda.qk_rows.QkRows.load(gpu.d); + defer qk_rows.unload(); + var qks: [qk_k.forms.len]qk_k.Qk = undefined; + for (&qks, qk_k.forms) |*x, f| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .form = f, .k = &qk_rows, .arena = .init(a) }; + defer for (&qks) |*x| x.deinit(); + for (&qks) |*x| try kernels.append(a, x.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t);