cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64

zig/build/cuda.zig

@@ -47,6 +47,11 @@ const kernels = [_]Kernel{
     .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle
     .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3
     .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1
+    .{ .name = "tree_attention", .flags = &.{"-O3"} }, // ours: committed keys and paths (cuda/tree_attention.zig)
+    .{ .name = "block_attention", .flags = &.{"-O3"} }, // ours: a drafter's block over its context (cuda/block_attention.zig)
+    .{ .name = "window_attention", .flags = &.{"-O3"} }, // ours: sliding-window attention (cuda/window_attention.zig)
+    .{ .name = "moe_glue", .flags = &.{"-O3"} }, // ours: sigmoid top-k and the weighted combine (cuda/moe_glue.zig)
+    .{ .name = "dense_rows", .flags = &.{"-O3"} }, // ours: bf16 rows times weight rows into fp32 (cuda/dense_rows.zig)
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts
     .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu
     .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op

zig/kernels/cuda/attn_odd.cuh

@@ -0,0 +1,63 @@
+// NaN and infinite values in a tensor-core attention tile's V: a key a row cannot see adds nothing to that row's sum.
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+namespace attn_odd {
+
+// Each bf16 half of x that is NaN or an infinity (all-ones exponent) as +0, the finite halves as they are.
+__device__ __forceinline__ uint32_t finite2(uint32_t x) {
+    return x & (((x & 0x7f80u) == 0x7f80u ? 0u : 0xffffu) | ((x & 0x7f800000u) == 0x7f800000u ? 0u : 0xffff0000u));
+}
+
+// Some NaN or infinite value in a warp's mma rows g (r0: a real row) and g + 8 (r1) of its fold.
+template <int NB>
+__device__ __forceinline__ bool nonfinite_out(const float (&o)[NB][4], bool r0, bool r1) {
+    float a = 0.0f, b = 0.0f; // x - x is 0 for a finite x, NaN for NaN and the infinities
+#pragma unroll
+    for (int i = 0; i < NB; ++i) {
+        a = __fadd_rn(a, __fadd_rn(__fsub_rn(o[i][0], o[i][0]), __fsub_rn(o[i][1], o[i][1])));
+        b = __fadd_rn(b, __fadd_rn(__fsub_rn(o[i][2], o[i][2]), __fsub_rn(o[i][3], o[i][3])));
+    }
+    return (r0 && a != a) || (r1 && b != b);
+}
+
+// A 32-key tile staged [key][D] (xor swizzle): the chunks this thread copied, their NaN and infinite values set to 0.
+template <int D>
+__device__ __forceinline__ void clean(uint4* vs, int threads) {
+    for (int i = threadIdx.x; i < 32 * (D / 8); i += threads) {
+        uint4* p = vs + (i / (D / 8)) * (D / 8) + ((i % (D / 8)) ^ ((i / (D / 8)) & 7));
+        const uint4 x = *p;
+        *p = make_uint4(finite2(x.x), finite2(x.y), finite2(x.z), finite2(x.w));
+    }
+}
+
+// Keys k0..k1 (V row at vrow(k)) that clean() emptied: their NaN and infinite values to the rows that see them.
+template <int D, typename Row, typename Vis>
+__device__ __forceinline__ void add_back(float (&o)[D / 8][4], Row vrow, int k0, int k1, int lane, Vis vis, float l0,
+                                         float l1) {
+    const int g = lane / 4, qd = lane % 4;
+    for (int k = k0; k < k1; ++k) {
+        const bool v0 = l0 > 0.0f && vis(g, k), v1 = l1 > 0.0f && vis(g + 8, k); // l 0: no key weighs, 0 / 0 stays NaN
+        if (!v0 && !v1) continue;
+        const uint32_t* v = reinterpret_cast<const uint32_t*>(vrow(k));
+#pragma unroll
+        for (int i = 0; i < D / 8; ++i) {
+            const uint32_t x = v[4 * i + qd];
+            if (finite2(x) == x) continue;
+            const float lo = __uint_as_float(x << 16), hi = __uint_as_float(x & 0xffff0000u);
+            if ((x & 0x7f80u) == 0x7f80u) {
+                if (v0) o[i][0] = __fadd_rn(o[i][0], lo);
+                if (v1) o[i][2] = __fadd_rn(o[i][2], lo);
+            }
+            if ((x & 0x7f800000u) == 0x7f800000u) {
+                if (v0) o[i][1] = __fadd_rn(o[i][1], hi);
+                if (v1) o[i][3] = __fadd_rn(o[i][3], hi);
+            }
+        }
+    }
+}
+
+}  // namespace attn_odd

zig/kernels/cuda/block_attention.cu

@@ -0,0 +1,227 @@
+// DFlash-style block attention for every CUDA family (float64 reference: core/attention_ref.zig): context, then block.
+
+// Block (stream, key head): a warp a query head, its 16 mma rows the block's rows; key tiles of 32 shared by the warps.
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <math.h>
+#include <stdint.h>
+
+#include "attn_odd.cuh"
+#include "qmm_frag.cuh"
+
+namespace tf_block {
+
+using qmm_frag::commit;
+using qmm_frag::cp16z;
+using qmm_frag::ldmatrix4;
+using qmm_frag::mma;
+using qmm_frag::mma0;
+
+constexpr int D = 128, C = D / 8, K16 = D / 16, NB = D / 8, KT = 32;
+constexpr float LOG2E = 1.4426950408889634f;
+
+__device__ __forceinline__ void wait_all() { asm volatile("cp.async.wait_group 0;\n" ::); }
+
+__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) {
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(qmm_frag::smem(p)));
+}
+
+__device__ __forceinline__ float texp(float x) {
+    float y;
+    asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E)));
+    return y;
+}
+
+__device__ __forceinline__ uint32_t pack(float lo, float hi) {
+    __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi);
+    return *reinterpret_cast<uint32_t*>(&v);
+}
+
+__device__ __forceinline__ int sw(int r, int c) { return r * C + (c ^ (r & 7)); }
+
+// One warp's running fold over a 32-key tile: scores masked by `vis(row, key)`, online softmax per row, P x V.
+struct Fold {
+    float o[NB][4];
+    float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.f, l1 = 0.f;
+
+    __device__ __forceinline__ Fold() {
+#pragma unroll
+        for (int i = 0; i < NB; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f;
+    }
+
+    template <typename Vis>
+    __device__ __forceinline__ void tile(const uint32_t (&qa)[K16][4], const uint4* ks, const uint4* vs, float scale,
+                                         int lane, int k0, Vis vis) {
+        const int g8 = lane / 4, qd = lane % 4, mt = lane / 8;
+        float sc[4][4];
+#pragma unroll
+        for (int k = 0; k < K16; ++k) {
+#pragma unroll
+            for (int n = 0; n < 4; n += 2) {
+                uint32_t b[4];
+                ldmatrix4(b, ks + sw(8 * (n + mt / 2) + lane % 8, 2 * k + mt % 2));
+                if (k == 0) {
+                    mma0(sc[n], qa[k], b[0], b[1]);
+                    mma0(sc[n + 1], qa[k], b[2], b[3]);
+                } else {
+                    mma(sc[n], qa[k], b[0], b[1]);
+                    mma(sc[n + 1], qa[k], b[2], b[3]);
+                }
+            }
+        }
+        float tm0 = -INFINITY, tm1 = -INFINITY;
+#pragma unroll
+        for (int n = 0; n < 4; ++n) {
+            const int c = k0 + 8 * n + 2 * qd;
+            sc[n][0] = vis(g8, c) ? __fmul_rn(sc[n][0], scale) : -INFINITY;
+            sc[n][1] = vis(g8, c + 1) ? __fmul_rn(sc[n][1], scale) : -INFINITY;
+            sc[n][2] = vis(g8 + 8, c) ? __fmul_rn(sc[n][2], scale) : -INFINITY;
+            sc[n][3] = vis(g8 + 8, c + 1) ? __fmul_rn(sc[n][3], scale) : -INFINITY;
+            tm0 = fmaxf(tm0, fmaxf(sc[n][0], sc[n][1]));
+            tm1 = fmaxf(tm1, fmaxf(sc[n][2], sc[n][3]));
+        }
+        for (int x = 1; x <= 2; x *= 2) {
+            tm0 = fmaxf(tm0, __shfl_xor_sync(0xffffffffu, tm0, x));
+            tm1 = fmaxf(tm1, __shfl_xor_sync(0xffffffffu, tm1, x));
+        }
+        const float n0 = fmaxf(m0, tm0), n1 = fmaxf(m1, tm1);
+        const float a0 = tm0 == -INFINITY ? 1.f : (m0 == -INFINITY ? 0.f : texp(__fsub_rn(m0, n0)));
+        const float a1 = tm1 == -INFINITY ? 1.f : (m1 == -INFINITY ? 0.f : texp(__fsub_rn(m1, n1)));
+        float c0 = 0.f, c1 = 0.f;
+#pragma unroll
+        for (int n = 0; n < 4; ++n) {
+            sc[n][0] = sc[n][0] != -INFINITY ? texp(__fsub_rn(sc[n][0], n0)) : 0.f;
+            sc[n][1] = sc[n][1] != -INFINITY ? texp(__fsub_rn(sc[n][1], n0)) : 0.f;
+            sc[n][2] = sc[n][2] != -INFINITY ? texp(__fsub_rn(sc[n][2], n1)) : 0.f;
+            sc[n][3] = sc[n][3] != -INFINITY ? texp(__fsub_rn(sc[n][3], n1)) : 0.f;
+            c0 = __fadd_rn(c0, __fadd_rn(sc[n][0], sc[n][1]));
+            c1 = __fadd_rn(c1, __fadd_rn(sc[n][2], sc[n][3]));
+        }
+        for (int x = 1; x <= 2; x *= 2) {
+            c0 = __fadd_rn(c0, __shfl_xor_sync(0xffffffffu, c0, x));
+            c1 = __fadd_rn(c1, __shfl_xor_sync(0xffffffffu, c1, x));
+        }
+        l0 = __fmaf_rn(l0, a0, c0);
+        l1 = __fmaf_rn(l1, a1, c1);
+        m0 = n0;
+        m1 = n1;
+#pragma unroll
+        for (int i = 0; i < NB; ++i) {
+            o[i][0] = __fmul_rn(o[i][0], a0);
+            o[i][1] = __fmul_rn(o[i][1], a0);
+            o[i][2] = __fmul_rn(o[i][2], a1);
+            o[i][3] = __fmul_rn(o[i][3], a1);
+        }
+        uint32_t pa[2][4];
+#pragma unroll
+        for (int j = 0; j < 2; ++j) {
+            pa[j][0] = pack(sc[2 * j][0], sc[2 * j][1]);
+            pa[j][1] = pack(sc[2 * j][2], sc[2 * j][3]);
+            pa[j][2] = pack(sc[2 * j + 1][0], sc[2 * j + 1][1]);
+            pa[j][3] = pack(sc[2 * j + 1][2], sc[2 * j + 1][3]);
+        }
+#pragma unroll
+        for (int j = 0; j < 2; ++j)
+#pragma unroll
+            for (int i = 0; i < NB; i += 2) {
+                uint32_t b[4];
+                ldmatrix4t(b, vs + sw(16 * j + 8 * (mt % 2) + lane % 8, i + mt / 2));
+                mma(o[i], pa[j], b[0], b[1]);
+                mma(o[i + 1], pa[j], b[2], b[3]);
+            }
+    }
+};
+
+}  // namespace tf_block
+
+using namespace tf_block;
+
+// Block rows r see context key k when lens[j] + r - k <= window, then block keys c < L (c <= r if causal).
+extern "C" __global__ void __launch_bounds__(256) tf_block_attention(
+    const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, const __nv_bfloat16* __restrict__ vb,
+    const long long* __restrict__ table, const int* __restrict__ lens, __nv_bfloat16* __restrict__ out, float scale,
+    int window, int rows, int h, int hk, int L, int causal) {
+    extern __shared__ uint4 smem[];
+    uint4* ks = smem;
+    uint4* vs = ks + KT * C;
+    const int warp = threadIdx.x / 32, lane = threadIdx.x % 32, g8 = lane / 4, qd = lane % 4;
+    const int j = blockIdx.x, g = blockIdx.y, G = h / hk;
+    const int head = g * G + warp;
+    const int s = lens[j];
+    const __nv_bfloat16* kc = reinterpret_cast<const __nv_bfloat16*>(table[2 * j]) + (int64_t)g * s * D;
+    const __nv_bfloat16* vc = reinterpret_cast<const __nv_bfloat16*>(table[2 * j + 1]) + (int64_t)g * s * D;
+    uint32_t qa[K16][4];
+    {
+        const int r0 = g8, r1 = g8 + 8;
+        const uint32_t* q0 = reinterpret_cast<const uint32_t*>(q + ((int64_t)head * rows + j * L + min(r0, L - 1)) * D);
+        const uint32_t* q1 = reinterpret_cast<const uint32_t*>(q + ((int64_t)head * rows + j * L + min(r1, L - 1)) * D);
+#pragma unroll
+        for (int k = 0; k < K16; ++k) {
+            qa[k][0] = r0 < L ? q0[8 * k + qd] : 0u;
+            qa[k][1] = r1 < L ? q1[8 * k + qd] : 0u;
+            qa[k][2] = r0 < L ? q0[8 * k + qd + 4] : 0u;
+            qa[k][3] = r1 < L ? q1[8 * k + qd + 4] : 0u;
+        }
+    }
+    Fold f;
+    const int lo = max(0, s - window); // the earliest context key any row sees (row 0: s - k <= window)
+    const auto vis = [&](int r, int key) { return r < L && key < s && s + r - key <= window; };
+    // the context's edge tiles, where some row cannot see some key: e0..e1
+    const int e0 = lo / KT * KT, e1 = min(s, e0 + (max(0, s + L - 1 - window - e0) + KT - 1) / KT * KT);
+    auto context = [&](bool odd) { // odd: the edge tiles' NaN and infinite values leave the mma (added back below)
+        for (int k0 = e0; k0 < s; k0 += KT) {
+            for (int i = threadIdx.x; i < KT * C; i += blockDim.x) {
+                const int key = k0 + i / C;
+                const bool in = key < s;
+                cp16z(ks + sw(i / C, i % C), kc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in);
+                cp16z(vs + sw(i / C, i % C), vc + (int64_t)(in ? key : 0) * D + (i % C) * 8, in);
+            }
+            commit();
+            wait_all();
+            if (odd && k0 < e1) attn_odd::clean<D>(vs, blockDim.x);
+            __syncthreads();
+            f.tile(qa, ks, vs, scale, lane, k0, vis);
+            __syncthreads();
+        }
+    };
+    auto block_keys = [&] { // the block's own keys
+        for (int i = threadIdx.x; i < KT * C; i += blockDim.x) {
+            const int c = i / C;
+            const bool in = c < L;
+            const int64_t at = ((int64_t)g * rows + j * L + (in ? c : 0)) * D + (i % C) * 8;
+            cp16z(ks + sw(c, i % C), kb + at, in);
+            cp16z(vs + sw(c, i % C), vb + at, in);
+        }
+        commit();
+        wait_all();
+    };
+    const auto own = [&](int r, int c) { return r < L && c < L && (!causal || c <= r); };
+    context(false);
+    block_keys();
+    __syncthreads();
+    f.tile(qa, ks, vs, scale, lane, 0, own);
+    // a NaN or infinite output: again, values some row cannot see out of the mma, added back (finite values: same bits)
+    if (__syncthreads_or(attn_odd::nonfinite_out<NB>(f.o, g8 < L, g8 + 8 < L))) {
+        f = Fold();
+        context(true);
+        attn_odd::add_back<D>(f.o, [&](int k) { return vc + (int64_t)k * D; }, e0, e1, lane, vis, f.l0, f.l1);
+        block_keys();
+        if (causal) attn_odd::clean<D>(vs, blockDim.x); // row r cannot see keys past r
+        __syncthreads();
+        f.tile(qa, ks, vs, scale, lane, 0, own);
+        const auto row = [&](int c) { return vb + ((int64_t)g * rows + j * L + c) * D; };
+        if (causal) attn_odd::add_back<D>(f.o, row, 0, L, lane, own, f.l0, f.l1);
+    }
+    if (g8 < L)
+#pragma unroll
+        for (int i = 0; i < NB; ++i)
+            *reinterpret_cast<uint32_t*>(out + ((int64_t)(j * L + g8) * h + head) * D + 8 * i + 2 * qd) =
+                pack(__fdiv_rn(f.o[i][0], f.l0), __fdiv_rn(f.o[i][1], f.l0));
+    if (g8 + 8 < L)
+#pragma unroll
+        for (int i = 0; i < NB; ++i)
+            *reinterpret_cast<uint32_t*>(out + ((int64_t)(j * L + g8 + 8) * h + head) * D + 8 * i + 2 * qd) =
+                pack(__fdiv_rn(f.o[i][2], f.l1), __fdiv_rn(f.o[i][3], f.l1));
+}

zig/kernels/cuda/dense_rows.cu

@@ -0,0 +1,70 @@
+// bf16 rows times bf16 weight rows into fp32 for every CUDA family: routers and output heads (core/dense_ref.zig).
+
+// A warp an output column, RT rows a launch tile; each row's sum runs in one order whatever RT and the row count.
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+namespace tf_dense {
+
+constexpr int WARPS = 8;
+constexpr int VEC = 8; // bf16 a 16-byte load
+
+__device__ __forceinline__ void unpack(const uint4& u, float* f) {
+    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&u);
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+        const float2 t = __bfloat1622float2(h[i]);
+        f[2 * i] = t.x;
+        f[2 * i + 1] = t.y;
+    }
+}
+
+// out[r, n] = x[r] . w[n] for a tile's rows: lane l sums k = 8 l + 256 j in order, then the fixed shuffle tree.
+template <int RT>
+__device__ __forceinline__ void tile(const __nv_bfloat16* x, int64_t x_stride, const __nv_bfloat16* w, float* out,
+                                     int64_t out_stride, int rows, int n, int k) {
+    const int lane = threadIdx.x & 31;
+    const int col = blockIdx.x * WARPS + (threadIdx.x >> 5);
+    const int r0 = blockIdx.y * RT;
+    if (col >= n) return;
+    float acc[RT];
+#pragma unroll
+    for (int r = 0; r < RT; ++r) acc[r] = 0.f;
+    const __nv_bfloat16* wr = w + static_cast<int64_t>(col) * k;
+#pragma unroll 4 // several 16-byte weight loads in flight a lane (memory-level parallelism)
+    for (int k0 = lane * VEC; k0 < k; k0 += 32 * VEC) {
+        float wf[VEC];
+        unpack(*reinterpret_cast<const uint4*>(wr + k0), wf);
+#pragma unroll
+        for (int r = 0; r < RT; ++r) {
+            if (r0 + r >= rows) break;
+            float xf[VEC];
+            unpack(*reinterpret_cast<const uint4*>(x + (r0 + r) * x_stride + k0), xf);
+#pragma unroll
+            for (int i = 0; i < VEC; ++i) acc[r] = __fmaf_rn(xf[i], wf[i], acc[r]);
+        }
+    }
+#pragma unroll
+    for (int r = 0; r < RT; ++r) {
+        float v = acc[r];
+        for (int m = 16; m > 0; m >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xFFFFFFFFu, v, m));
+        if (lane == 0 && r0 + r < rows) out[(r0 + r) * out_stride + col] = v;
+    }
+}
+
+}  // namespace tf_dense
+
+using namespace tf_dense;
+
+#define TF_DENSE_ROWS(RT)                                                                                              \
+    extern "C" __global__ void __launch_bounds__(WARPS * 32) tf_dense_rows_##RT(                                     \
+        const __nv_bfloat16* x, long long x_stride, const __nv_bfloat16* w, float* out, long long out_stride,          \
+        int rows, int n, int k) {                                                                                      \
+        tile<RT>(x, x_stride, w, out, out_stride, rows, n, k);                                                         \
+    }
+
+TF_DENSE_ROWS(1)
+TF_DENSE_ROWS(4)
+TF_DENSE_ROWS(16)
+TF_DENSE_ROWS(32)