dory: full-attention KV in pages of the core pool

docs/kernel-contract.md

@@ -155,6 +155,9 @@ output within float64's range over the spread its fp32 steps and tensor-core sum
 seven controls, each one planted fault).
 CUDA also checks Dory's nGPT rows: the L2 norm rows with and without their bf16 add, and the normalized lerp residual
 (`norms_l2.cuh`), each output within float64's range at its bf16 points (three controls, each dropping one point).
+CUDA also checks Dory's paged caches (`window_pages.cu`): `tf_wrows_put` places rows in a shuffled arena bit for bit, and the
+head dim 256 decode and prompt kernels read them through a page table, each output within float64's range (two controls:
+swapped stream tables, a stale page id).
 CUDA also checks Dory's attention inputs (`qk_l2.cuh`): RoPE first, each head's L2 norm, q times its scale vector, v as is, at
 head dim 256 and positions past 199,800 (two controls, each dropping one rounding point).
 CUDA also checks Dory's activation rows (`act_l2.cuh`): the scaled sigmoid output gate over all heads and the scaled SwiGLU,

zig/build/cuda_kernels.zig

@@ -89,6 +89,7 @@ pub const kernels = [_]Kernel{
     .{ .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 = "window_rows", .flags = &.{"-O3"} }, // ours: window attention at head dims 256 and 512 (cuda/window_attention.zig)
+    .{ .name = "window_pages", .flags = &.{"-O3"} }, // ours: window_rows.cuh over Dory's paged caches (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

zig/kernels/cuda/window_pages.cu

@@ -0,0 +1,44 @@
+// window_rows.cuh over paged caches (Dory's full-attention layers, head dim 256): the kernels, and the rows' writer.
+
+#include "window_rows.cuh"
+
+using namespace tf_wrows;
+
+// Dory's full-attention caches in pages (head dim 256, no window): kc/vc a layer's part of page 0, tab a row of ids.
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_part_256_pg(
+    const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, const int* pos, const int* slot,
+    float* po, float* pm, float* pl, int h, int hk, int win, long long ring, float scale, const int* tab,
+    long long page, int tok, int stride) {
+    part<256, NONE, true>(q, kc, vc, pos, slot, po, pm, pl, h, hk, win, ring, scale, Pg{tab, page, tok, stride});
+}
+
+extern "C" __global__ void __launch_bounds__(NT) tf_wrows_merge_256_pg(
+    const float* po, const float* pm, const float* pl, __nv_bfloat16* out, int h, int hk, int chunks,
+    const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, const int* pos, const int* slot, int win,
+    long long ring, float scale, const int* tab, long long page, int tok, int stride) {
+    merge<256, NONE, true>(po, pm, pl, out, h, hk, chunks, q, kc, vc, pos, slot, win, ring, scale,
+                           Pg{tab, page, tok, stride});
+}
+
+// The prompt kernel's tab is its stream's own row of page ids.
+extern "C" __global__ void __launch_bounds__(NT, 1) tf_wrows_prompt_256_pg(
+    const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, __nv_bfloat16* out, const int* hi,
+    int p0, int W, int H, int HK, int G, float scale, int win, int ring, const int* tab, long long page, int tok) {
+    prompt::kernel<256, false, NONE, true>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring,
+                                           Pg{tab, page, tok, 0});
+}
+
+// Rows k, v [n, row] (row = hk * D elements) into the pages at their positions: a block a row, 16 bytes a thread.
+extern "C" __global__ void __launch_bounds__(128) tf_wrows_put(const __nv_bfloat16* k, const __nv_bfloat16* v,
+                                                               __nv_bfloat16* kc, __nv_bfloat16* vc, const int* pos,
+                                                               const int* slot, const int* tab, long long page,
+                                                               int tok, int stride, int row) {
+    const int64_t r = blockIdx.x;
+    const int p = pos[r];
+    const int64_t id = tab[slot[r] * stride + p / tok];
+    const int64_t at = id * page + static_cast<int64_t>(p % tok) * row;
+    for (int e = threadIdx.x * 8; e < row; e += 128 * 8) {
+        *reinterpret_cast<uint4*>(kc + at + e) = *reinterpret_cast<const uint4*>(k + r * row + e);
+        *reinterpret_cast<uint4*>(vc + at + e) = *reinterpret_cast<const uint4*>(v + r * row + e);
+    }
+}

zig/kernels/cuda/window_rows.cuh

@@ -28,6 +28,14 @@ enum Fault : int {
     TRUNC = 4,   // the merge's output rounded toward zero
 };
 
+// A paged cache's address: a stream's table row of page ids, elements a page holds across its caches, positions a page.
+struct Pg {
+    const int* tab;
+    long long page;
+    int tok;
+    int stride; // table entries a slot (decode); a prompt kernel is given its stream's row
+};
+
 using qmm_frag::commit;
 using qmm_frag::cp16z;
 using qmm_frag::ldmatrix4;
@@ -101,11 +109,12 @@ __device__ __noinline__ void slow_head(const __nv_bfloat16* qh, KV kv, int lo, i
 }
 
 // Decode partials, a block a (row, key head, chunk): the chunk's max, weights and o; empty chunks write m = -inf.
-template <int D, int FAULT>
+template <int D, int FAULT, bool PG = false>
 __device__ __forceinline__ void part(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
                                      const __nv_bfloat16* __restrict__ vc, const int* __restrict__ pos,
                                      const int* __restrict__ slot, float* __restrict__ po, float* __restrict__ pm,
-                                     float* __restrict__ pl, int h, int hk, int win, long long ring, float scale) {
+                                     float* __restrict__ pl, int h, int hk, int win, long long ring, float scale,
+                                     Pg pg = Pg{}) {
     constexpr int DPL = D / 32;  // a lane's dims of q and k, and its columns of v
     constexpr int HW = 32 / DPL; // query heads a warp takes: 32 q values, then 32 accumulators, a lane
     constexpr int U = 4;         // keys a warp loads at once
@@ -126,12 +135,19 @@ __device__ __forceinline__ void part(const __nv_bfloat16* __restrict__ q, const
     }
     const int nk = k1 - k0 + 1;
     const int64_t r0 = k0 % ring;
+    // a paged chunk lies in one page (no window: chunks start at multiples of CK, pages hold whole chunks)
+    int64_t pb = 0;
+    if constexpr (PG) {
+        const int ts = (slot != nullptr ? slot[r] : 0) * pg.stride;
+        pb = static_cast<int64_t>(pg.tab[ts + k0 / pg.tok]) * pg.page + static_cast<int64_t>(k0 % pg.tok) * hk * D;
+    }
     // warp (hg, kg): heads hg HW.., keys [j0, j1), the kg-th of WG runs in key order
     const int HG = (G + HW - 1) / HW, WG = max(1, NW / HG);
     const int hg = warp % HG, kg = warp / HG, run = (nk + WG - 1) / WG;
     const bool busy = warp < HG * WG;
     const int j0 = min(kg * run, nk), j1 = min(j0 + run, nk);
     const auto rowp = [&](const __nv_bfloat16* cache, int jj) {
+        if constexpr (PG) return cache + pb + (static_cast<int64_t>(jj) * hk + kh) * D + lane * DPL;
         int64_t row = r0 + jj;
         if (row >= ring) row -= ring;
         return cache + ((base + row) * hk + kh) * D + lane * DPL;
@@ -240,13 +256,13 @@ __device__ __forceinline__ void part(const __nv_bfloat16* __restrict__ q, const
 }
 
 // out = sum_c a_c o_c / sum_c a_c l_c, a_c = expf(m_c - M), chunks in order; NaN or inf: the head again in float64.
-template <int D, int FAULT>
+template <int D, int FAULT, bool PG = false>
 __device__ __forceinline__ void merge(const float* __restrict__ po, const float* __restrict__ pm,
                                       const float* __restrict__ pl, __nv_bfloat16* __restrict__ out, int h, int hk,
                                       int chunks, const __nv_bfloat16* __restrict__ q,
                                       const __nv_bfloat16* __restrict__ kc, const __nv_bfloat16* __restrict__ vc,
                                       const int* __restrict__ pos, const int* __restrict__ slot, int win,
-                                      long long ring, float scale) {
+                                      long long ring, float scale, Pg pg = Pg{}) {
     const int64_t r = blockIdx.x;
     const int kh = blockIdx.y, G = h / hk;
     const int o = blockIdx.z * NT + threadIdx.x; // NT divides D: the block's outputs are one head's
@@ -269,7 +285,12 @@ __device__ __forceinline__ void merge(const float* __restrict__ po, const float*
         const int p = pos[r];
         const int64_t base = (slot != nullptr ? static_cast<int64_t>(slot[r]) : 0) * ring;
         const auto kv = [&](int key, int which) {
-            return (which == 0 ? kc : vc) + ((base + key % ring) * hk + kh) * D;
+            const __nv_bfloat16* c = which == 0 ? kc : vc;
+            if constexpr (PG) {
+                const int64_t pid = pg.tab[(slot != nullptr ? slot[r] : 0) * pg.stride + key / pg.tok];
+                return c + pid * pg.page + (static_cast<int64_t>(key % pg.tok) * hk + kh) * D;
+            }
+            return c + ((base + key % ring) * hk + kh) * D;
         };
         const int c0 = blockIdx.z * NT % D;
         slow_head<D>(q + (r * h + kh * G + hd) * D, kv, win > 0 ? max(p - win + 1, 0) : 0, p, (double)scale, oh, c0,
@@ -313,11 +334,11 @@ struct Cfg {
 };
 
 // Prompt rows p0..: row i sees keys [p0 + i - win + 1, hi[i]]; CS warps a 16-row band, slices summed; P32: P halves.
-template <int D, bool P32, int FAULT>
+template <int D, bool P32, int FAULT, bool PG = false>
 __device__ __forceinline__ bool rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
                                      const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out,
                                      const int* __restrict__ hi, int p0, int W, int H, int HK, int G, float scale,
-                                     int win, int ring) {
+                                     int win, int ring, Pg pg = Pg{}) {
     using K = Cfg<D>;
     constexpr int CS = K::CS, NG = K::NG, C = K::C, NS = K::NS, SL = K::SL, K16 = K::K16, NB = K::NB;
     extern __shared__ uint4 wsmem[];
@@ -345,11 +366,17 @@ __device__ __forceinline__ bool rows(const __nv_bfloat16* __restrict__ q, const
         const __nv_bfloat16* src = item % 2 == 0 ? kc : vc;
         const int key0 = (t0 + item / 2) * BN;
         uint4* dst = slots + (item % NS) * K::SLOT;
+        int64_t pb = 0; // a paged tile lies in one page (pages hold whole tiles)
+        if constexpr (PG) pb = static_cast<int64_t>(pg.tab[key0 / pg.tok]) * pg.page + static_cast<int64_t>(key0 % pg.tok) * HK * D;
         for (int i = threadIdx.x; i < BN * C; i += NT) {
             const int key = key0 + i / C;
             const bool in = key >= kfirst && key <= klast && key < limit; // outside every row's keys: zeros
-            const int row = in ? (ring > 0 ? key % ring : key) : 0;
-            cp16z(dst + sw<D>(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in);
+            if constexpr (PG) {
+                cp16z(dst + sw<D>(i / C, i % C), in ? src + pb + ((int64_t)(i / C) * HK + hk) * D + (i % C) * 8 : src, in);
+            } else {
+                const int row = in ? (ring > 0 ? key % ring : key) : 0;
+                cp16z(dst + sw<D>(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in);
+            }
         }
     };
     uint32_t qa[K16][4]; // this warp's SL dims of rows r0, r1
@@ -528,16 +555,18 @@ __device__ __forceinline__ bool rows(const __nv_bfloat16* __restrict__ q, const
 }
 
 // A block with a NaN or infinite output: recompute only affected heads, leaving every finite head's bits alone.
-template <int D>
+template <int D, bool PG = false>
 __device__ __noinline__ void odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc,
                                  __nv_bfloat16* out, const int* hi, int p0, int W, int H, int HK, int G, float scale,
-                                 int win, int ring) {
+                                 int win, int ring, Pg pg = Pg{}) {
     using K = Cfg<D>;
     const int HB = G % K::NG == 0 ? K::NG : 1, RB = K::NG / HB;
     const int hk = blockIdx.y / (G / HB), hg = blockIdx.y % (G / HB);
     const int row0 = blockIdx.x * 16 * RB;
     const auto kv = [&](int key, int which) {
-        return (which == 0 ? kc : vc) + ((int64_t)(ring > 0 ? key % ring : key) * HK + hk) * D;
+        const __nv_bfloat16* c = which == 0 ? kc : vc;
+        if constexpr (PG) return c + (int64_t)pg.tab[key / pg.tok] * pg.page + ((int64_t)(key % pg.tok) * HK + hk) * D;
+        return c + ((int64_t)(ring > 0 ? key % ring : key) * HK + hk) * D;
     };
     for (int i = 0; i < 16 * RB * HB; ++i) {
         const int r = row0 + i % (16 * RB), head = hk * G + hg * HB + i / (16 * RB);
@@ -552,12 +581,12 @@ __device__ __noinline__ void odd(const __nv_bfloat16* q, const __nv_bfloat16* kc
     }
 }
 
-template <int D, bool P32, int FAULT>
+template <int D, bool P32, int FAULT, bool PG = false>
 __device__ __forceinline__ void kernel(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc,
                                        __nv_bfloat16* out, const int* hi, int p0, int W, int H, int HK, int G,
-                                       float scale, int win, int ring) {
-    if (rows<D, P32, FAULT>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring))
-        odd<D>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring);
+                                       float scale, int win, int ring, Pg pg = Pg{}) {
+    if (rows<D, P32, FAULT, PG>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring, pg))
+        odd<D, PG>(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring, pg);
 }
 
 } // namespace prompt

zig/src/core/contract/window_rows_records.zig

@@ -1,4 +1,4 @@
-//! window_rows.cu's launches on CUDA (head dims 256 and 512), read by kernel-contract-limits from every SM.
+//! window_rows.cu's launches on CUDA (head dims 256 and 512) and window_pages.cu's, read by kernel-contract-limits.
 const registry = @import("registry.zig");
 
 /// nvcc 13.0's usage, the compiler of the release fatbins.
@@ -11,4 +11,8 @@ pub const cuda = [_]registry.CudaRecord{
     .{ .launch = .{ .name = "window_rows tf_wrows_merge_512", .symbol = "tf_wrows_merge_512", .image = "window_rows", .threads = 256 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 62, .static = 128, .bound = 256 }, .{ .regs = 62, .static = 128, .bound = 256 } } },
     .{ .launch = .{ .name = "window_rows tf_wrows_prompt_512", .symbol = "tf_wrows_prompt_512", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 241, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 242, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 } } },
     .{ .launch = .{ .name = "window_rows tf_wrows_prompt_512_p32", .symbol = "tf_wrows_prompt_512_p32", .image = "window_rows", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 249, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 251, .static = 144, .bound = 256 }, .{ .regs = 254, .static = 144, .bound = 256 }, .{ .regs = 254, .static = 144, .bound = 256 } } },
+    .{ .launch = .{ .name = "window_pages tf_wrows_part_256_pg", .symbol = "tf_wrows_part_256_pg", .image = "window_pages", .threads = 256 }, .usage = .{ .{ .regs = 98, .static = 46208, .bound = 256 }, .{ .regs = 104, .static = 46208, .bound = 256 }, .{ .regs = 104, .static = 46208, .bound = 256 }, .{ .regs = 108, .static = 46208, .bound = 256 }, .{ .regs = 108, .static = 46208, .bound = 256 } } },
+    .{ .launch = .{ .name = "window_pages tf_wrows_merge_256_pg", .symbol = "tf_wrows_merge_256_pg", .image = "window_pages", .threads = 256 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 64, .static = 128, .bound = 256 }, .{ .regs = 40, .static = 128, .bound = 256 }, .{ .regs = 40, .static = 128, .bound = 256 } } },
+    .{ .launch = .{ .name = "window_pages tf_wrows_prompt_256_pg", .symbol = "tf_wrows_prompt_256_pg", .image = "window_pages", .threads = 256, .dynamic = 81920, .opt_in = true }, .usage = .{ .{ .regs = 241, .static = 144, .bound = 256 }, .{ .regs = 244, .static = 144, .bound = 256 }, .{ .regs = 244, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 }, .{ .regs = 255, .static = 144, .bound = 256 } } },
+    .{ .launch = .{ .name = "window_pages tf_wrows_put", .symbol = "tf_wrows_put", .image = "window_pages", .threads = 128 }, .usage = .{ .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 38, .static = 0, .bound = 128 }, .{ .regs = 38, .static = 0, .bound = 128 } } },
 };

zig/src/core/kv_pool.zig

@@ -269,6 +269,17 @@ pub const Pool = struct {
         if (table) |t| if (l.ids.items.len > old) try t.write(t.ctx, l.ids.items, old);
     }
 
+    /// `cover` for an empty table: n consecutive ids from a multiple of `step` when they are free, else as `cover`.
+    pub fn coverRun(p: *Pool, l: *Lease, n: usize, step: usize, table: ?Table) !void {
+        if (l.ids.items.len > 0 or n == 0 or l.promised < n) return p.cover(l, n, table);
+        try l.ids.ensureTotalCapacity(p.gpa, n);
+        const start = p.findAligned(n, step) orelse return p.cover(l, n, table);
+        p.claimAt(l, start, n);
+        for (0..n) |i| l.ids.appendAssumeCapacity(@intCast(start + i)); // held as `take` holds a page: `finish` frees them
+        errdefer p.shrink(l, 0);
+        if (table) |t| try t.write(t.ctx, l.ids.items, 0);
+    }
+
     /// Pages past `n` in `l`'s table go back to its promise, as they were before it took them.
     fn shrink(p: *Pool, l: *Lease, n: usize) void {
         while (l.ids.items.len > n) {
@@ -294,10 +305,26 @@ pub const Pool = struct {
         return null;
     }
 
+    /// The highest first id of `k` free consecutive pages at a multiple of `step` (> 0), or null.
+    pub fn findAligned(p: *const Pool, k: usize, step: usize) ?u32 {
+        if (k == 0 or k > p.refs.len) return null;
+        var start = (p.refs.len - k) / step * step;
+        while (true) : (start -= step) {
+            if (std.mem.allEqual(u32, p.refs[start..][0..k], 0)) return @intCast(start);
+            if (start < step) return null;
+        }
+    }
+
     /// `k` consecutive pages from `l`'s promise, held as one run until releaseRun; null: no such run is free now.
     pub fn claimRun(p: *Pool, l: *Lease, k: usize) ?u32 {
         if (l.promised < k) return null;
         const start = p.findRun(k) orelse return null;
+        p.claimAt(l, start, k);
+        return start;
+    }
+
+    /// Free pages [start, start + k) held by `l`, from its promise.
+    fn claimAt(p: *Pool, l: *Lease, start: u32, k: usize) void {
         var n: usize = 0;
         for (p.free[0..p.nfree]) |id| if (id < start or id >= start + k) {
             p.free[n] = id;
@@ -309,7 +336,6 @@ pub const Pool = struct {
         @memset(p.leased[start..][0..k], 1); // held as a lease holds its pages: never the store's to reclaim
         l.promised -= k;
         p.promised -= k;
-        return start;
     }
 
     /// A run claimed at `start` goes back to the free pages.