diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 7777f80686..0d8241d5c2 100644 --- a/docs/kernel-contract.md +++ b/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, diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 268c3ce405..8ee6998300 100644 --- a/zig/build/cuda_kernels.zig +++ b/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 diff --git a/zig/kernels/cuda/window_pages.cu b/zig/kernels/cuda/window_pages.cu new file mode 100644 index 0000000000..5d34a6e814 --- /dev/null +++ b/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(p % tok) * row; + for (int e = threadIdx.x * 8; e < row; e += 128 * 8) { + *reinterpret_cast(kc + at + e) = *reinterpret_cast(k + r * row + e); + *reinterpret_cast(vc + at + e) = *reinterpret_cast(v + r * row + e); + } +} diff --git a/zig/kernels/cuda/window_rows.cuh b/zig/kernels/cuda/window_rows.cuh index 566ad08206..37402546d5 100644 --- a/zig/kernels/cuda/window_rows.cuh +++ b/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 +template __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(pg.tab[ts + k0 / pg.tok]) * pg.page + static_cast(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(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 +template __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(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(key % pg.tok) * hk + kh) * D; + } + return c + ((base + key % ring) * hk + kh) * D; }; const int c0 = blockIdx.z * NT % D; slow_head(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 +template __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; 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(pg.tab[key0 / pg.tok]) * pg.page + static_cast(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(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in); + if constexpr (PG) { + cp16z(dst + sw(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(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 +template __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; 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 +template __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(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring)) - odd(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring); + float scale, int win, int ring, Pg pg = Pg{}) { + if (rows(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring, pg)) + odd(q, kc, vc, out, hi, p0, W, H, HK, G, scale, win, ring, pg); } } // namespace prompt diff --git a/zig/src/core/contract/window_rows_records.zig b/zig/src/core/contract/window_rows_records.zig index c4c9e0db06..8ac7b01272 100644 --- a/zig/src/core/contract/window_rows_records.zig +++ b/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 } } }, }; diff --git a/zig/src/core/kv_pool.zig b/zig/src/core/kv_pool.zig index cc439d93e9..a1783f19b5 100644 --- a/zig/src/core/kv_pool.zig +++ b/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. diff --git a/zig/src/core/kv_pool_test.zig b/zig/src/core/kv_pool_test.zig index 20b60dd873..61f6117ce9 100644 --- a/zig/src/core/kv_pool_test.zig +++ b/zig/src/core/kv_pool_test.zig @@ -300,6 +300,38 @@ test "a run: consecutive pages from the top, from the stream's promise, back who try expectEqual(@as(usize, 10), p.available()); } +test "coverRun: a lease's pages as one run from a multiple of the alignment while one is free, else scattered" { + const a = std.testing.allocator; + var p = try Pool.init(a, 10, 1); + defer p.deinit(); + var x: Lease = .{}; + try p.promise(&x, 3); + try p.coverRun(&x, 3, 4, null); + try expect(std.mem.eql(u32, x.ids.items, &.{ 4, 5, 6 })); // 8 + 3 is past the pool: the window at 4 + try expectEqual(@as(usize, 0), x.promised); + var y: Lease = .{}; + try p.promise(&y, 2); + try p.coverRun(&y, 2, 4, null); + try expect(std.mem.eql(u32, y.ids.items, &.{ 8, 9 })); + var z: Lease = .{}; + try p.promise(&z, 1); + try p.cover(&z, 1, null); // page 0: no window of 4 is free now + var w: Lease = .{}; + try p.promise(&w, 2); + try p.coverRun(&w, 2, 4, null); + try expect(std.mem.eql(u32, w.ids.items, &.{ 1, 2 })); // from the free stack + p.finish(&w); + p.finish(&x); + try p.promise(&w, 4); + try p.coverRun(&w, 4, 1, null); // align 1: any run, the highest + try expect(std.mem.eql(u32, w.ids.items, &.{ 4, 5, 6, 7 })); + p.finish(&w); + p.finish(&y); + p.finish(&z); + try expect(p.idle()); + try expectEqual(@as(usize, 10), p.available()); +} + test "the plan with a run: each stream pays its run's pages beside its first page, the window bound adds them" { const g: u64 = 1 << 30; var l = nemotron; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 66b29d264a..dec012894e 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -57,6 +57,7 @@ pub const tree_attention: []const u8 = if (available) &Blob("fatbin_tree_attenti pub const block_attention: []const u8 = if (available) &Blob("fatbin_block_attention").bytes else &.{}; pub const window_attention: []const u8 = if (available) &Blob("fatbin_window_attention").bytes else &.{}; pub const window_rows: []const u8 = if (available) &Blob("fatbin_window_rows").bytes else &.{}; +pub const window_pages: []const u8 = if (available) &Blob("fatbin_window_pages").bytes else &.{}; pub const moe_glue: []const u8 = if (available) &Blob("fatbin_moe_glue").bytes else &.{}; pub const dense_rows: []const u8 = if (available) &Blob("fatbin_dense_rows").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; diff --git a/zig/src/cuda/window_attention.zig b/zig/src/cuda/window_attention.zig index 48e6bc534d..f806a038bc 100644 --- a/zig/src/cuda/window_attention.zig +++ b/zig/src/cuda/window_attention.zig @@ -10,8 +10,19 @@ const Module = module.Module; const Function = module.Function; const Stream = stream_.Stream; +/// A Window's `pg` (win 0, head dim 256): cache pointers name a layer in page 0; `table` rows of ids, `stride` a slot. +pub const Pages = struct { table: u64, page_elems: usize, tokens: usize, stride: usize = 0 }; + /// A row's heads and keys (pos - win, pos] (win 0: all) at cache row slot * ring + key % ring; p32: wide P in halves. -pub const Window = struct { h: usize, hk: usize, d: usize, win: usize, ring: usize, scale: f32, p32: bool = false }; +pub const Window = struct { h: usize, hk: usize, d: usize, win: usize, ring: usize, scale: f32, p32: bool = false, pg: ?Pages = null }; + +/// Positions a page holds at least (a multiple of it): the decode chunk, so a chunk and a prompt tile stay in a page. +pub const page_unit = chunk_keys; + +fn pagedShape(w: Window) bool { + const p = w.pg orelse return true; + return w.d == 256 and w.win == 0 and p.tokens > 0 and p.tokens % page_unit == 0 and p.page_elems <= std.math.maxInt(i63); +} /// Head dims window_rows.cu serves, with any query heads a key head up to 16. pub fn wide(d: usize) bool { @@ -42,10 +53,15 @@ pub const WindowAttention = struct { tc_fn: Function, prompt_fns: [2]Function, // 4, then 8 query heads a block rows: Module, // window_rows.cu + pages: ?Module = null, // window_pages.cu (loadPages): window_rows.cuh over paged caches + put_fn: ?Function = null, // tf_wrows_put: rows into paged caches wide_fns: [rows_smem.len]Wide, /// window_rows.cu's kernels at one head dim: decode partials and merge, prompt rows (bf16 P, then P32). - pub const Wide = struct { part: Function, merge: Function, prompt: [2]Function }; + pub const Wide = struct { part: Function, merge: Function, prompt: [2]Function, paged: ?Paged = null }; + + /// The head dim 256 kernels over paged caches: decode partials, their merge, prompt rows (P as bf16). + pub const Paged = struct { part: Function, merge: Function, prompt: Function }; pub fn load(d: *const Driver) !WindowAttention { if (!kernels.available) return error.BuiltWithoutKernels; @@ -71,7 +87,20 @@ pub const WindowAttention = struct { }, .prompt_fns = p, .rows = rows, .wide_fns = wf }; } + /// Loads window_pages.cu: decode, prompt rows and `put` over caches in pages (head dim 256). + pub fn loadPages(g: *WindowAttention, d: *const Driver) !void { + if (g.pages != null) return; + var pages = try Module.load(d, kernels.window_pages); + errdefer pages.unload(); + const p: Paged = .{ .part = try pages.function("tf_wrows_part_256_pg"), .merge = try pages.function("tf_wrows_merge_256_pg"), .prompt = try pages.function("tf_wrows_prompt_256_pg") }; + try p.prompt.allowDynamicShared(rows_smem[0]); + g.put_fn = try pages.function("tf_wrows_put"); + g.wide_fns[0].paged = p; + g.pages = pages; + } + pub fn unload(g: *WindowAttention) void { + if (g.pages) |*p| p.unload(); g.rows.unload(); g.mod.unload(); } @@ -147,7 +176,8 @@ pub const WindowAttention = struct { /// Wide decode rows by `f` (wideFns, or the contract's controls): chunk partials (scratchBytes), their merge. pub fn decodeWith(f: *const Wide, s: Stream, q: u64, kc: u64, vc: u64, out: u64, pos: u64, slot: ?u64, w: Window, scratch: u64, n: usize) !void { - if (!wideShape(w) or w.ring == 0 or w.ring > std.math.maxInt(i32)) return error.WindowShape; + if (!wideShape(w) or w.ring == 0 or w.ring > std.math.maxInt(i32) or !pagedShape(w)) return error.WindowShape; + const pf = if (w.pg != null) f.paged orelse return error.WindowShape else null; if (n == 0) return; const t = chunks(w); const parts = n * w.hk * t * (w.h / w.hk); @@ -161,8 +191,9 @@ pub const WindowAttention = struct { a.add(@as(i32, @intCast(w.win))); a.add(@as(i64, @intCast(w.ring))); a.add(w.scale); + if (w.pg) |p| addPages(&a, p, true); const pairs = if (w.d == 512) (w.h / w.hk + 3) / 4 else 1; // head dim 512: a block 4 query heads - try launch_.launch(f.part, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk * pairs), .z = @intCast(t) }, .block = .{ .x = 256 }, .shared = if (w.d == 512) staged_smem else 0 }, s, &a); + try launch_.launch(if (pf) |x| x.part else f.part, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk * pairs), .z = @intCast(t) }, .block = .{ .x = 256 }, .shared = if (w.d == 512) staged_smem else 0 }, s, &a); var b: launch_.Args = .{}; for ([_]u64{ po, pm, pl, out }) |v| b.add(v); b.add(@as(i32, @intCast(w.h))); @@ -172,7 +203,30 @@ pub const WindowAttention = struct { b.add(@as(i32, @intCast(w.win))); b.add(@as(i64, @intCast(w.ring))); b.add(w.scale); - try launch_.launch(f.merge, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = @intCast(w.h / w.hk * w.d / 256) }, .block = .{ .x = 256 } }, s, &b); + if (w.pg) |p| addPages(&b, p, true); + try launch_.launch(if (pf) |x| x.merge else f.merge, .{ .grid = .{ .x = @intCast(n), .y = @intCast(w.hk), .z = @intCast(w.h / w.hk * w.d / 256) }, .block = .{ .x = 256 } }, s, &b); + } + + /// A paged kernel's trailing arguments: the table, elements a page, positions a page, and (decode) entries a slot. + fn addPages(a: *launch_.Args, p: Pages, strided: bool) void { + a.add(p.table); + a.add(@as(i64, @intCast(p.page_elems))); + a.add(@as(i32, @intCast(p.tokens))); + if (strided) a.add(@as(i32, @intCast(p.stride))); + } + + /// n rows k, v [n, hk * d] bf16 into the paged caches kc/vc (a layer's part of page 0) at pos/slot (int32 [n]). + pub fn put(g: *const WindowAttention, s: Stream, k: u64, v: u64, kc: u64, vc: u64, pos: u64, slot: u64, w: Window, n: usize) !void { + const p = w.pg orelse return error.WindowShape; + if (!pagedShape(w)) return error.WindowShape; + if (n == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ k, v, kc, vc, pos, slot, p.table }) |x| a.add(x); + a.add(@as(i64, @intCast(p.page_elems))); + a.add(@as(i32, @intCast(p.tokens))); + a.add(@as(i32, @intCast(p.stride))); + a.add(@as(i32, @intCast(w.hk * w.d))); + try launch_.launch(g.put_fn orelse return error.WindowShape, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 128 } }, s, &a); } /// n prompt rows of one stream at p0, p0 + 1, ... (d 128); kc/vc that stream's ring, key p at row p % ring. @@ -202,7 +256,7 @@ pub const WindowAttention = struct { /// promptRows by `f` (wideFns, or the contract's controls). pub fn promptWith(f: *const Wide, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, n: usize, w: Window, hi: ?u64) !void { const grp = if (w.hk > 0) w.h / w.hk else 0; - if (!wideShape(w) or p0 + n > std.math.maxInt(i32)) return error.WindowShape; + if (!wideShape(w) or p0 + n > std.math.maxInt(i32) or !pagedShape(w) or (w.pg != null and (w.p32 or f.paged == null))) return error.WindowShape; if (n == 0) return; const ng = 1024 / w.d; // 16-row groups a block: 8 warps, a 128-dim slice each const hb = if (grp % ng == 0) ng else 1; @@ -213,8 +267,9 @@ pub const WindowAttention = struct { a.add(w.scale); a.add(@as(i32, @intCast(w.win))); a.add(@as(i32, @intCast(w.ring))); + if (w.pg) |p| addPages(&a, p, false); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((n + rows - 1) / rows), .y = @intCast(w.hk * (grp / hb)), .z = 1 }, .block = .{ .x = 256 }, .shared = promptShared(w.d) }; - try launch_.launch(f.prompt[@intFromBool(w.p32)], cfg, s, &a); + try launch_.launch(if (w.pg != null) f.paged.?.prompt else f.prompt[@intFromBool(w.p32)], cfg, s, &a); } }; diff --git a/zig/src/families/dory/cuda_backend.zig b/zig/src/families/dory/cuda_backend.zig index 7a50677565..f0d2c1a0f8 100644 --- a/zig/src/families/dory/cuda_backend.zig +++ b/zig/src/families/dory/cuda_backend.zig @@ -4,6 +4,8 @@ const cuda = @import("cuda"); const lanes = @import("lanes"); const fwd = @import("cuda_model.zig"); const snap = @import("cuda_snap.zig"); +const pages = @import("cuda_pages.zig"); +const kv = @import("core").kv_pool; const be = lanes.backend; const Stream = lanes.Stream; @@ -29,6 +31,7 @@ const Lane = struct { rule: u64, // sample.cu's Rule, when the stream samples sampled: bool = false, fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + lease: kv.Lease = .{}, // a paged slot's pages }; /// The physical positions a slot holds for a served window: the window, the verify reserve, in whole 1,024s. @@ -55,6 +58,8 @@ pub const Dory = struct { meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot next: u64 = 0, + pool: ?kv.Pool = null, // the pages every slot leases (attach) + scatter: bool = false, // a check's: leases take the free pages in their order, never a run /// The device bytes `init` allocates: bf16 logits, the token ring, metas, each slot's last row and rule. pub fn openBytes(vocab: usize, slots: usize) usize { @@ -63,7 +68,7 @@ pub const Dory = struct { pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *fwd.Model, k: *const Kernels, s: cuda.Stream) !Dory { const v = m.c.vocab; - const free = try gpa.alloc(bool, m.slots); + const free = try gpa.alloc(bool, m.slotCap()); errdefer gpa.free(free); @memset(free, true); var logits16 = try cuda.DeviceBuffer.alloc(d, round_rows * v * 2); @@ -72,7 +77,7 @@ pub const Dory = struct { errdefer tok.free(); var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); errdefer meta.free(); - const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slots); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slotCap()); errdefer gpa.free(bufs); var made: usize = 0; errdefer for (bufs[0..made]) |*b| b.free(); @@ -83,9 +88,26 @@ pub const Dory = struct { return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .logits16 = logits16, .tok = tok, .meta = meta, .slot_bufs = bufs }; } + /// Makes a paged model's arena of `pool_pages` pages for `extra + 1` slots, and the pool that leases them. + pub fn attach(x: *Dory, extra: usize, pool_pages: usize) !void { + if (x.m.paging == null) try x.m.attach(x.s, extra, pool_pages); // a check may have attached it already + x.pool = try kv.Pool.init(x.gpa, x.m.paging.?.pool, x.m.paging.?.pageBytes()); + } + + /// The pool's free pages in a seeded random order, leased in that order, so tables are not runs (a check's hook). + pub fn shufflePages(x: *Dory, seed: u64) void { + x.scatter = true; + var prng = std.Random.DefaultPrng.init(seed); + if (x.pool) |*p| prng.random().shuffle(u32, p.free[0..p.nfree]); + } + pub fn deinit(x: *Dory) void { var it = x.lanes.valueIterator(); - while (it.next()) |l| l.rows.deinit(x.gpa); + while (it.next()) |l| { + if (x.pool) |*p| p.finish(&l.lease); + l.rows.deinit(x.gpa); + } + if (x.pool) |*p| p.deinit(); x.lanes.deinit(x.gpa); for (x.slot_bufs) |*b| b.free(); x.gpa.free(x.slot_bufs); @@ -103,7 +125,7 @@ pub const Dory = struct { /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a forward. pub fn facts(x: *const Dory) lanes.Model { - return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slots, round_rows)) }; + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slotCap(), round_rows)) }; } fn self(ptr: *anyopaque) *Dory { @@ -147,7 +169,42 @@ pub const Dory = struct { } fn freeSlot(x: *const Dory) ?usize { - return std.mem.indexOfScalar(bool, x.free, true); + return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); + } + + /// Pages a request of this prompt and reply holds. + fn pagesFor(prompt: usize, max_new: usize) usize { + return pages.need(prompt, max_new, reserve); + } + + /// Whether a request could start now: a free slot and (paged) its pages; one that never fits passes. + pub fn admits(x: *const Dory, prompt: usize, max_new: usize) bool { + if (!fits(x.m.context, prompt, max_new)) return true; + const p = x.pool orelse return x.freeSlot() != null; + return x.freeSlot() != null and p.fits(pagesFor(prompt, max_new)); + } + + const Map = struct { + x: *Dory, + slot: usize, + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const m: *const Map = @ptrCast(@alignCast(ctx)); + try m.x.m.mapPages(m.x.s, m.slot, ids[from..], from); + } + }; + + /// A free slot for the request, its pages leased now when paged: a window's run (contiguous) when one is free. + fn takeSlot(x: *Dory, lease: *kv.Lease, prompt: usize, max_new: usize) !usize { + const slot = x.freeSlot() orelse return error.NoFreeSlot; + if (x.pool == null) return slot; + const p = &x.pool.?; + const n = pagesFor(prompt, max_new); + try p.promise(lease, n); + errdefer p.finish(lease); + var map: Map = .{ .x = x, .slot = slot }; + const table: kv.Table = .{ .ctx = &map, .write = Map.write }; + try (if (x.scatter) p.cover(lease, n, table) else p.coverRun(lease, n, x.m.paging.?.per_slot, table)); + return slot; } /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. @@ -157,12 +214,13 @@ pub const Dory = struct { if (!fits(x.m.context, p.len, @as(usize, s.max_new))) return error.PastContext; const got = try x.lanes.getOrPut(x.gpa, s); if (!got.found_existing) { - const slot = x.freeSlot() orelse { + var lease: kv.Lease = .{}; + const slot = x.takeSlot(&lease, p.len, s.max_new) catch |e| { x.lanes.removeByPtr(got.key_ptr); - return error.NoFreeSlot; + return e; }; x.free[slot] = false; - got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr }; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .lease = lease }; } const l = got.value_ptr; l.sampled = false; @@ -327,8 +385,12 @@ pub const Dory = struct { } fn dropLane(x: *Dory, s: *const Stream) void { - if (x.lanes.fetchRemove(s)) |kv| { - var l = kv.value; + if (x.lanes.fetchRemove(s)) |entry| { + var l = entry.value; + if (x.pool) |*p| { + p.finish(&l.lease); + x.m.unmapPages(l.slot); + } x.free[l.slot] = true; l.rows.deinit(x.gpa); } diff --git a/zig/src/families/dory/cuda_kernels.zig b/zig/src/families/dory/cuda_kernels.zig index 7c88141d56..e6e7b2450a 100644 --- a/zig/src/families/dory/cuda_kernels.zig +++ b/zig/src/families/dory/cuda_kernels.zig @@ -37,6 +37,7 @@ pub const Kernels = struct { errdefer k.dense.unload(); k.win = try cuda.window_attention.WindowAttention.load(d); errdefer k.win.unload(); + try k.win.loadPages(d); if (!cuda.kernels.available) return k; const kk = cuda.kernels; const images = [_][]const u8{ kk.sample, kk.torch_argmax, kk.torch_pointwise }; diff --git a/zig/src/families/dory/cuda_model.zig b/zig/src/families/dory/cuda_model.zig index 768126fa40..0b24b7b174 100644 --- a/zig/src/families/dory/cuda_model.zig +++ b/zig/src/families/dory/cuda_model.zig @@ -5,6 +5,7 @@ const config = @import("config.zig"); const schedule = @import("schedule.zig"); const memory = @import("memory.zig"); const wts = @import("cuda_weights.zig"); +const pages = @import("cuda_pages.zig"); const Config = config.Config; const DeviceBuffer = cuda.DeviceBuffer; @@ -40,16 +41,21 @@ pub const Model = struct { ring: usize, slots: usize, k: []u64 = &.{}, // per cache: [slots, len, HK, D] bf16 (len: context for full caches, the ring for sliding ones) - v: []u64 = &.{}, + v: []u64 = &.{}, // paged: a full cache's K and V regions of the arena, page id p at rows [p * tokens, (p + 1) * tokens) + paged: bool = false, // full caches live in pages (attach makes them); sliding rings stay a buffer a slot + paging: ?pages.Paging = null, + rings: std.ArrayList(DeviceBuffer) = .empty, // the sliding caches' buffers, made again for the slots attach names inv: [2]u64 = .{ 0, 0 }, // fp32 [D / 2] inverse frequencies: sliding, full s: Scratch = undefined, buffers: std.ArrayList(DeviceBuffer) = .empty, planned: ?usize = null, // a plan: the bytes counted, nothing allocated /// Decode forwards of n rows as CUDA graphs at this model's loop count, captured at first use ([n]). graphs: [round_rows + 1]?cuda.graph.Exec = @splat(null), + /// Paged: decode forwards of n rows whose streams all hold a window of pages (the contiguous kernels). + flat_graphs: [round_rows + 1]?cuda.graph.Exec = @splat(null), graphed: bool = true, - const Scratch = struct { ids: u64, pos: u64, slot: u64, h: u64, h2: u64, state: u64, input: u64, out: u64, q: u64, kk: u64, vv: u64, g: u64, qn: u64, kn: u64, vn: u64, att: u64, gated: u64, mg: u64, mu: u64, m: u64, part: u64, win: u64, logits: u64 }; + const Scratch = struct { ids: u64, pos: u64, slot: u64, fslot: u64, h: u64, h2: u64, state: u64, input: u64, out: u64, q: u64, kk: u64, vv: u64, g: u64, qn: u64, kn: u64, vn: u64, att: u64, gated: u64, mg: u64, mu: u64, m: u64, part: u64, win: u64, logits: u64 }; fn alloc(m: *Model, bytes: usize) !u64 { if (m.planned) |*p| { @@ -71,6 +77,19 @@ pub const Model = struct { return p; } + /// A zeroed buffer for a sliding cache (paged models make these again for the slots attach names). + fn ringZeros(m: *Model, s: cuda.Stream, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += @max(bytes, 256); + return 0; + } + var b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer b.free(); + try m.d.check(m.d.api.cuMemsetD8Async(b.ptr, 0, bytes, s.handle), "cuMemsetD8Async"); + try m.rings.append(m.gpa, b); + return b.ptr; + } + /// Cache `i`'s positions a slot: the context (full) or the ring (sliding). pub fn cacheLen(m: *const Model, i: usize) usize { return if (m.sched.cache_kind[i] == .full) m.context else m.ring; @@ -88,6 +107,90 @@ pub const Model = struct { return m; } + /// A model whose full caches are pages: one slot's rings now; `attach` makes the arena, the other slots' rings. + pub fn initPaged(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const Kernels, context: usize) !Model { + var m = blank(gpa, d, w, context, 1); + m.paged = true; + errdefer m.deinit(); + try m.build(s, k); + return m; + } + + /// The device bytes `initPaged` allocates. + pub fn openBytesPaged(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const Kernels, context: usize) !usize { + var m = blank(gpa, d, w, context, 1); + defer m.deinit(); + m.paged = true; + m.planned = 0; + try m.build(s, k); + return m.planned.?; + } + + /// What the host's pool plans for this model's pages (see cuda_pages.layout). + pub fn layout(m: *const Model) @import("core").kv_pool.Layout { + return pages.layout(m.c, m.sched.cacheCount(.full), m.sched.cacheCount(.sliding), m.context, m.ring); + } + + /// Slots a backend holds: the model's now, or (paged) as many as a round takes. + pub fn slotCap(m: *const Model) usize { + return if (m.paged) pages.max_slots else m.slots; + } + + /// Makes the arena of `pool` pages to lease, the tables and the rings of `extra + 1` slots. + pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void { + if (!m.paged or m.paging != null or extra >= pages.max_slots) return error.BadAttach; + const c = m.c; + const per_slot = pages.pagesOf(m.context); + const elems = pages.pageElems(c, m.sched.cacheCount(.full)); + const ids = try m.gpa.alloc(u32, pages.max_slots * per_slot); + errdefer m.gpa.free(ids); + @memset(ids, 0); + for (m.rings.items) |*b| b.free(); + m.rings.clearRetainingCapacity(); + m.slots = extra + 1; + try m.rings.ensureTotalCapacity(m.gpa, 2 * m.sched.caches); + var p: pages.Paging = .{ .pool = pool, .per_slot = per_slot, .page_elems = elems, .arena = 0, .table = 0, .ids = ids }; + p.arena = try m.alloc(pool * elems * 2); + p.table = try m.alloc(pages.max_slots * per_slot * 4); + try m.d.check(m.d.api.cuMemsetD8Async(p.arena, 0, pool * elems * 2, s.handle), "cuMemsetD8Async"); + try m.upload(s, p.table, std.mem.sliceAsBytes(ids)); + const region = pool * pages.tokens * c.kvDim() * 2; // one full cache's K (or V) rows over every page + var block: usize = 0; + for (m.sched.cache_kind[0..m.sched.caches], m.k, m.v) |kind, *kc, *vc| { + if (kind == .full) { + kc.* = p.arena + block * region; + vc.* = kc.* + region; + block += 2; + } else { + const bytes = m.slots * m.ring * c.kvDim() * 2; + kc.* = try m.ringZeros(s, bytes); + vc.* = try m.ringZeros(s, bytes); + } + } + m.paging = p; + m.dropGraphs(); + try s.synchronize(); + } + + /// Slot `slot`'s pages are `ids` (arena page numbers, in position order from entry `from`): table and host copy. + pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void { + const p = if (m.paging) |*x| x else return error.NotPaged; + if (slot >= m.slots or from + ids.len > p.per_slot) return error.BadChain; + const at = slot * p.per_slot + from; + @memcpy(p.ids[at..][0..ids.len], ids); + try m.upload(s, p.table + at * 4, std.mem.sliceAsBytes(ids)); + p.caps[slot] = (from + ids.len) * pages.tokens; + p.wins[slot] = pages.windowOf(p.ids[slot * p.per_slot ..][0 .. from + ids.len], p.per_slot); + } + + /// Slot `slot` holds no pages (its table entries stay; no forward may reach them). + pub fn unmapPages(m: *Model, slot: usize) void { + if (m.paging) |*p| { + p.caps[slot] = 0; + p.wins[slot] = null; + } + } + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, k: *const Kernels, context: usize, slots: usize) !usize { var m = blank(gpa, d, w, context, slots); @@ -106,10 +209,13 @@ pub const Model = struct { const c = m.c; m.k = try m.gpa.alloc(u64, m.sched.caches); m.v = try m.gpa.alloc(u64, m.sched.caches); + @memset(m.k, 0); + @memset(m.v, 0); for (m.k, m.v, 0..) |*kc, *vc, i| { + if (m.paged and m.sched.cache_kind[i] == .full) continue; const bytes = m.slots * m.cacheLen(i) * c.kvDim() * 2; - kc.* = try m.zeros(s, bytes); - vc.* = try m.zeros(s, bytes); + kc.* = if (m.paged) try m.ringZeros(s, bytes) else try m.zeros(s, bytes); + vc.* = if (m.paged) try m.ringZeros(s, bytes) else try m.zeros(s, bytes); } for ([_]f64{ c.theta_sliding, c.theta_full }, &m.inv) |theta, *p| { var inv: [256]f32 = undefined; // vLLM's inv_freq in fp32: 1 / theta^(2j / D) @@ -123,9 +229,12 @@ pub const Model = struct { } pub fn deinit(m: *Model) void { - for (&m.graphs) |*g| if (g.*) |*e| e.deinit(); + m.dropGraphs(); for (m.buffers.items) |*b| b.free(); m.buffers.deinit(m.gpa); + for (m.rings.items) |*b| b.free(); + m.rings.deinit(m.gpa); + if (m.paging) |p| m.gpa.free(p.ids); m.gpa.free(m.k); m.gpa.free(m.v); } @@ -136,6 +245,10 @@ pub const Model = struct { e.deinit(); g.* = null; }; + for (&m.flat_graphs) |*g| if (g.*) |*e| { + e.deinit(); + g.* = null; + }; } fn upload(m: *const Model, s: cuda.Stream, dst: u64, bytes: []const u8) !void { @@ -146,9 +259,14 @@ pub const Model = struct { try m.d.check(m.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, s.handle), "cuMemcpyDtoDAsync"); } - fn window(m: *const Model, full: bool) cuda.window_attention.Window { + /// A call's window; full caches of a paged model read through the table (a prompt's slot: its own row). + fn window(m: *const Model, full: bool, prompt_slot: ?usize) cuda.window_attention.Window { const c = m.c; - return .{ .h = c.heads, .hk = c.kv_heads, .d = c.head_dim, .win = if (full) 0 else c.window, .ring = if (full) m.context else m.ring, .scale = c.factors().attn }; + var w: cuda.window_attention.Window = .{ .h = c.heads, .hk = c.kv_heads, .d = c.head_dim, .win = if (full) 0 else c.window, .ring = if (full) m.context else m.ring, .scale = c.factors().attn }; + if (full) if (m.paging) |p| { + w.pg = .{ .table = p.table + (prompt_slot orelse 0) * p.per_slot * 4, .page_elems = pages.tokens * c.kvDim(), .tokens = pages.tokens, .stride = p.per_slot }; + }; + return w; } fn scratch(m: *Model, k: *const Kernels) !void { @@ -164,11 +282,12 @@ pub const Model = struct { part = @max(part, k.lane.partBytes(rows, q.nvfp4)); }, }; - const win = @max(WindowAttention.scratchBytes(m.window(true), attend_rows), WindowAttention.scratchBytes(m.window(false), attend_rows)); + const win = @max(WindowAttention.scratchBytes(m.window(true, null), attend_rows), WindowAttention.scratchBytes(m.window(false, null), attend_rows)); m.s = .{ .ids = try m.alloc(r * 4), .pos = try m.alloc(r * 4), .slot = try m.alloc(r * 4), + .fslot = if (m.paged) try m.alloc(r * 4) else 0, // a row's window of pages, as a contiguous slot .h = try m.alloc(r * d * 2), .h2 = try m.alloc(r * d * 2), .state = try m.alloc(r * d * 2), @@ -197,12 +316,17 @@ pub const Model = struct { var rows: usize = 0; for (chains) |ch| { if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; + if (m.paging) |p| if (ch.p0 + ch.tokens.len > p.caps[ch.slot]) return error.PastPages; if (!prompt and ch.tokens.len > window_rows) return error.WindowTooWide; rows += ch.tokens.len; } if (rows == 0 or rows > @as(usize, if (prompt) prompt_rows else round_rows) or (prompt and chains.len != 1)) return error.TooManyRows; if (heads > round_rows or head0 + heads > rows) return error.TooManyRows; - const ids = try m.gpa.alloc(i32, 3 * rows); + // paged: when every stream holds a window of pages, its full caches are main's contiguous slot of that window + const flat = if (m.paging) |p| for (chains) |ch| { + if (p.wins[ch.slot] == null) break false; + } else true else false; + const ids = try m.gpa.alloc(i32, 4 * rows); defer m.gpa.free(ids); var at: usize = 0; for (chains) |ch| { @@ -210,25 +334,29 @@ pub const Model = struct { ids[at + i] = @intCast(t); ids[rows + at + i] = @intCast(ch.p0 + i); ids[2 * rows + at + i] = @intCast(ch.slot); + if (flat) ids[3 * rows + at + i] = @intCast(m.paging.?.wins[ch.slot].?); } at += ch.tokens.len; } try m.upload(s, m.s.ids, std.mem.sliceAsBytes(ids[0..rows])); try m.upload(s, m.s.pos, std.mem.sliceAsBytes(ids[rows .. 2 * rows])); - try m.upload(s, m.s.slot, std.mem.sliceAsBytes(ids[2 * rows ..])); + try m.upload(s, m.s.slot, std.mem.sliceAsBytes(ids[2 * rows .. 3 * rows])); + if (flat) try m.upload(s, m.s.fslot, std.mem.sliceAsBytes(ids[3 * rows ..])); const p0 = if (prompt) chains[0].p0 else 0; const slot = if (prompt) chains[0].slot else 0; - // a decode forward's launches depend on its row count only (positions and slots are read on the device) - if (prompt or !m.graphed or head0 != 0 or heads != rows) return m.launch(k, s, rows, prompt, p0, slot, head0, heads); - if (m.graphs[rows] == null) m.graphs[rows] = try m.capture(k, s, rows); - try m.graphs[rows].?.launchOn(s); + const fslot = if (flat and prompt) m.paging.?.wins[slot].? else slot; + // a decode forward's launches depend on its row count and `flat` only (the rest is read on the device) + if (prompt or !m.graphed or head0 != 0 or heads != rows) return m.launch(k, s, rows, prompt, p0, slot, fslot, flat, head0, heads); + const g = if (flat) &m.flat_graphs[rows] else &m.graphs[rows]; + if (g.* == null) g.* = try m.capture(k, s, rows, flat); + try g.*.?.launchOn(s); } /// A decode forward of `rows` rows, every row's logits, captured as a graph. - fn capture(m: *Model, k: *const Kernels, s: cuda.Stream, rows: usize) !cuda.graph.Exec { + fn capture(m: *Model, k: *const Kernels, s: cuda.Stream, rows: usize, flat: bool) !cuda.graph.Exec { try s.synchronize(); try cuda.graph.beginCapture(s, .thread_local); - m.launch(k, s, rows, false, 0, 0, 0, rows) catch |e| { + m.launch(k, s, rows, false, 0, 0, 0, flat, 0, rows) catch |e| { if (cuda.graph.endCapture(s)) |g| { var gg = g; gg.deinit(); @@ -241,7 +369,7 @@ pub const Model = struct { } /// The forward's launches over staged ids, positions and slots: no host copies, so a graph can hold them. - fn launch(m: *Model, k: *const Kernels, s: cuda.Stream, n: usize, prompt: bool, p0: usize, slot: usize, head0: usize, heads: usize) !void { + fn launch(m: *Model, k: *const Kernels, s: cuda.Stream, n: usize, prompt: bool, p0: usize, slot: usize, fslot: usize, flat: bool, head0: usize, heads: usize) !void { const c = m.c; const lin: cuda.qlinear.Linear = .{ .lane = &k.lane, .nvfp4_prompt = &k.prompt }; const d = c.hidden; @@ -260,7 +388,7 @@ pub const Model = struct { } const alpha = switch (m.w.subs[step.sub]) { .attn => |a| blk: { - try m.attention(k, s, lin, a, step, h, n, prompt, p0, slot); + try m.attention(k, s, lin, a, step, h, n, prompt, p0, slot, fslot, flat); break :blk a.alpha; }, .mlp => |p| blk: { @@ -282,7 +410,7 @@ pub const Model = struct { try k.dense.rowsScaled(s, h + head0 * d * 2, d, m.w.head, m.w.logit, x.logits, c.vocab, heads, c.vocab, d); } - fn attention(m: *Model, k: *const Kernels, s: cuda.Stream, lin: cuda.qlinear.Linear, a: wts.Attn, step: schedule.Step, h: u64, n: usize, prompt: bool, p0: usize, slot: usize) !void { + fn attention(m: *Model, k: *const Kernels, s: cuda.Stream, lin: cuda.qlinear.Linear, a: wts.Attn, step: schedule.Step, h: u64, n: usize, prompt: bool, p0: usize, slot: usize, fslot: usize, flat: bool) !void { const c = m.c; const x = m.s; try m.proj(lin, s, a.q, h, x.q, n, prompt); @@ -295,17 +423,23 @@ pub const Model = struct { const src: cuda.qk_rows.Src = .{ .q = x.q, .q_ld = @intCast(c.qDim()), .q_hs = dh, .k = x.kk, .k_ld = @intCast(c.kvDim()), .k_hs = dh, .v = x.vv, .v_ld = @intCast(c.kvDim()), .v_hs = dh }; const shape: cuda.qk_rows.Shape = .{ .l2 = true, .q_scale = a.qs, .inv = m.inv[@intFromBool(full)], .rot = c.head_dim, .eps = eps }; const len = m.cacheLen(cache); - try k.qk.rowsFrom(s, src, .{ .h = c.heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, x.slot, .{ .q = x.qn, .k = x.kn, .v = x.vn, .kc = m.k[cache], .vc = m.v[cache], .len = len }, n); - const w = m.window(full); + const flat_full = full and flat; // a full cache's window of pages w is main's slot w: its kernels and addresses + const pg = full and m.paging != null and !flat; // written by `put` into the pages, not by the rows kernel + const sl = if (flat_full) x.fslot else x.slot; + const out: cuda.qk_rows.Out = .{ .q = x.qn, .k = x.kn, .v = x.vn, .kc = if (pg) null else m.k[cache], .vc = if (pg) null else m.v[cache], .len = len }; + try k.qk.rowsFrom(s, src, .{ .h = c.heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, sl, out, n); + var w = m.window(full, slot); + if (flat_full) w.pg = null; + if (pg) try k.win.put(s, x.kn, x.vn, m.k[cache], m.v[cache], x.pos, x.slot, m.window(true, null), n); // the whole table if (prompt) { - const off = slot * len * c.kvDim() * 2; + const off = if (pg) 0 else (if (flat_full) fslot else slot) * len * c.kvDim() * 2; try k.win.promptRows(s, x.qn, m.k[cache] + off, m.v[cache] + off, x.att, p0, n, w, null); } else { const qw = c.qDim() * 2; var r: usize = 0; while (r < n) : (r += attend_rows) { const rows = @min(attend_rows, n - r); - try k.win.decode(s, x.qn + r * qw, m.k[cache], m.v[cache], x.att + r * qw, x.pos + r * 4, x.slot + r * 4, w, x.win, rows); + try k.win.decode(s, x.qn + r * qw, m.k[cache], m.v[cache], x.att + r * qw, x.pos + r * 4, sl + r * 4, w, x.win, rows); } } try k.act.gateL2(s, x.att, x.g, c.qDim(), a.sg, x.gated, n, c.qDim(), eps); diff --git a/zig/src/families/dory/cuda_native.zig b/zig/src/families/dory/cuda_native.zig index bc536d9b95..5ed13bc24d 100644 --- a/zig/src/families/dory/cuda_native.zig +++ b/zig/src/families/dory/cuda_native.zig @@ -7,6 +7,8 @@ const fwd = @import("cuda_model.zig"); const be = @import("cuda_backend.zig"); const snap = @import("cuda_snap.zig"); const config = @import("config.zig"); +const pages = @import("cuda_pages.zig"); +const kv = @import("core").kv_pool; pub const model_type = "dory"; pub const formats: []const []const u8 = &.{"modelopt-nvfp4"}; @@ -33,6 +35,10 @@ pub const Loaded = struct { ctx: *anyopaque, deinit: *const fn (*anyopaque) void, target: *be.Dory, // the prompt cache keeps this backend's states (snap_*) + /// The full-attention caches as pages of the core pool (null: every stream allocated at open). + layout: ?kv.Layout = null, + /// With a layout: the first stream's rows are pages too, a window of them set aside by the plan. + engine_paged: bool = true, }; const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: fwd.Kernels, s: cuda.Stream, m: fwd.Model, db: be.Dory }; @@ -49,6 +55,11 @@ fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { return n; } +/// TENSORFOLD_DORY_KV_PAGES=0 keeps every stream's caches contiguous and allocated at open. +fn pagesOn() bool { + return !envOff("TENSORFOLD_DORY_KV_PAGES"); +} + /// The device bytes `n` streams take: their caches by loop count and context, the scratch, the backend's buffers. const Plan = struct { own: *const Owned, @@ -57,10 +68,29 @@ const Plan = struct { fn bytes(p: Plan, n: usize) !u64 { const o = p.own; + if (pagesOn()) return try fwd.Model.openBytesPaged(o.gpa, p.d, o.s, &o.w, &o.k, p.context) + be.Dory.openBytes(o.w.config.vocab, pages.max_slots); return try fwd.Model.openBytes(o.gpa, p.d, o.s, &o.w, &o.k, p.context, n) + be.Dory.openBytes(o.w.config.vocab, n); } }; +/// Paged: the first stream's rings and scratch must fit; the pool is planned by the host after. +fn fitFirst(room: u64, plan: anytype) !usize { + if (try plan.bytes(1) > room) return error.NoStreamFits; + return 1; +} + +/// The pool: `streams` more slots beside the first (rings made now) and an arena of `n` pages. +pub fn attach(p: *anyopaque, streams: u32, n: usize) !void { + const own: *Owned = @ptrCast(@alignCast(p)); + try own.db.attach(streams, n); +} + +/// Whether a request's caches could start now (the lane host's admission, on its engine thread). +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const own: *const Owned = @ptrCast(@alignCast(p)); + return own.db.admits(prompt, max_tokens); +} + /// The checkpoint at its loop count, the shared kernels and every stream's caches. pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { if (kernels != null) return error.CapturedKernels; @@ -80,18 +110,20 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer own.s.deinit(); const plan: Plan = .{ .own = own, .d = ctx.d, .context = context }; const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); - const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + const n = if (pagesOn()) try fitFirst(room, plan) else fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { if (e == error.ParallelPastMemory) { const most = fit(room, o.lanes, false, plan) catch 0; std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and {d} loops, and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, own.w.config.loops, gib(room), most, if (most == 1) "" else "s" }); } return e; }; - own.m = try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, n); + own.m = if (pagesOn()) try fwd.Model.initPaged(gpa, ctx.d, own.s, &own.w, &own.k, context) else try fwd.Model.init(gpa, ctx.d, own.s, &own.w, &own.k, context, n); errdefer own.m.deinit(); own.m.graphed = !envOff("TENSORFOLD_DORY_GRAPHS"); // =0: every decode forward eager own.db = try be.Dory.init(gpa, ctx.d, &own.m, &own.k, own.s); - std.log.info("Dory at {d} loops (per_loop KV): {d} caches a stream, {d} streams, CUDA graphs {s}", .{ own.w.config.loops, own.m.sched.caches, n, if (own.m.graphed) "on" else "off" }); + if (own.m.paged) { + std.log.info("Dory at {d} loops (per_loop KV): {d} caches a stream, the full ones ({d}) in pages of {d} positions, CUDA graphs {s}", .{ own.w.config.loops, own.m.sched.caches, own.m.sched.cacheCount(.full), pages.tokens, if (own.m.graphed) "on" else "off" }); + } else std.log.info("Dory at {d} loops (per_loop KV): {d} caches a stream, {d} streams, CUDA graphs {s}", .{ own.w.config.loops, own.m.sched.caches, n, if (own.m.graphed) "on" else "off" }); return .{ .backend = own.db.backend(), .facts = own.db.facts(), @@ -101,6 +133,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .ctx = own, .deinit = release, .target = &own.db, + .layout = if (own.m.paged) own.m.layout() else null, }; } diff --git a/zig/src/families/dory/cuda_pages.zig b/zig/src/families/dory/cuda_pages.zig new file mode 100644 index 0000000000..a4d5ea47d5 --- /dev/null +++ b/zig/src/families/dory/cuda_pages.zig @@ -0,0 +1,80 @@ +//! Dory's full-attention caches in pages: an arena (each cache's K, then V, over every page) and a table of ids a slot. +const std = @import("std"); +const kv = @import("core").kv_pool; +const config = @import("config.zig"); +const memory = @import("memory.zig"); + +/// Slots a round holds: its rows over the widest window (cuda_model's round_rows / window_rows). +pub const max_slots = 8; +/// Positions a page holds: window_rows.cu's decode chunk, so a chunk and a prompt tile stay inside one page. +pub const tokens = 256; + +/// The arena of `pool` pages, every slot's leases among them, and the slots' page tables. +pub const Paging = struct { + pool: usize, + per_slot: usize, // table entries a slot + page_elems: usize, // bf16 elements a page, over every full cache's K and V + arena: u64, + table: u64, + ids: []u32, // the table on the host, [max_slots * per_slot] + caps: [max_slots]usize = @splat(0), // positions the slot's pages cover + wins: [max_slots]?u32 = @splat(null), // w when the slot's pages are ids from w * per_slot on (contiguous slot w) + + pub fn pageBytes(p: Paging) usize { + return p.page_elems * 2; + } +}; + +/// w when `ids` are consecutive from w * per (a contiguous slot of the arena's caches), else null. +pub fn windowOf(ids: []const u32, per: usize) ?u32 { + if (ids.len == 0 or ids[0] % per != 0) return null; + for (ids, ids[0]..) |id, want| if (id != want) return null; + return @intCast(ids[0] / per); +} + +/// Pages that hold `positions`. +pub fn pagesOf(positions: usize) usize { + return (positions + tokens - 1) / tokens; +} + +/// bf16 elements a page holds across `full` caches: `tokens` rows of `kv_dim` in each cache's K and V. +pub fn pageElems(c: config.Config, full: usize) usize { + return 2 * full * tokens * c.kvDim(); +} + +/// What the host's pool plans: a slot costs its rings and table row, a token its full caches' rows. +pub fn layout(c: config.Config, full: usize, sliding: usize, context: usize, ring: usize) kv.Layout { + return .{ + .fixed_bytes = sliding * ring * memory.tokenBytes(c) + pagesOf(context) * 4, + .token_bytes = full * memory.tokenBytes(c), + .page_tokens = tokens, + .snap_bytes = 0, // kept states are copies of varying size: no slots + }; +} + +/// Pages a request holds: its prompt, its reply and one last verify window (positions written, from 0). +pub fn need(prompt: usize, max_new: usize, window: usize) usize { + return pagesOf(prompt + max_new + window); +} + +test "a page holds 256 positions of every full cache: 9, 14 and 24 MiB at 1, 2 and 4 loops" { + const text = try config.pinnedText(std.testing.allocator); + defer std.testing.allocator.free(text); + var why: config.Why = .{}; + const c = try config.parse(std.testing.allocator, text, &why); + for ([_]usize{ 1, 2, 4 }, [_]usize{ 9, 14, 24 }) |loops, full| { + const l = layout(try c.withLoops(loops), full, 3 * full, 32768, 3072); + try std.testing.expectEqual(@as(u64, full << 20), l.pageBytes()); + } + try std.testing.expectEqual(@as(usize, 129), pagesOf(32768 + 1)); + try std.testing.expectEqual(@as(usize, 1), need(1, 1, 16)); +} + +test "a window: consecutive ids from a multiple of the slot's pages; a gap, a step back, an offset or none is not" { + try std.testing.expectEqual(@as(?u32, 2), windowOf(&.{ 8, 9, 10 }, 4)); + try std.testing.expectEqual(@as(?u32, 0), windowOf(&.{0}, 4)); + try std.testing.expectEqual(@as(?u32, null), windowOf(&.{ 7, 8 }, 4)); + try std.testing.expectEqual(@as(?u32, null), windowOf(&.{ 8, 10 }, 4)); + try std.testing.expectEqual(@as(?u32, null), windowOf(&.{ 9, 8 }, 4)); + try std.testing.expectEqual(@as(?u32, null), windowOf(&.{}, 4)); +} diff --git a/zig/src/families/dory/cuda_snap.zig b/zig/src/families/dory/cuda_snap.zig index 2f58b1c69e..ef51f64e3d 100644 --- a/zig/src/families/dory/cuda_snap.zig +++ b/zig/src/families/dory/cuda_snap.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const config = @import("config.zig"); const fwd = @import("cuda_model.zig"); +const pages = @import("cuda_pages.zig"); /// What a state was made under: its loop count and KV mode (a state never crosses either). pub const Key = struct { loops: u8, mode: Mode }; @@ -41,6 +42,19 @@ fn each(m: *const fwd.Model, slot: usize, at: usize, ctx: anytype, comptime f: f for (0..m.sched.caches) |i| for ([_][]u64{ m.k, m.v }) |cache| { const full = m.sched.cache_kind[i] == .full; const len = m.cacheLen(i); + if (full) if (m.paging) |pg| { // a full cache's positions: one copy a run of consecutive page ids + var q: usize = 0; + while (q < at) { + const id = pg.ids[slot * pg.per_slot + q / pages.tokens]; + var end = (q / pages.tokens + 1) * pages.tokens; + while (end < at and pg.ids[slot * pg.per_slot + end / pages.tokens] == id + (end - q) / pages.tokens) end += pages.tokens; + const run = @min(at, end) - q; + try f(ctx, cache[i] + (@as(usize, id) * pages.tokens + q % pages.tokens) * rb, off, run * rb); + off += run * rb; + q += run; + } + continue; + }; const base = cache[i] + slot * len * rb; var p = if (full) 0 else at - tail(m.c.window, at); while (p < at) { @@ -108,6 +122,7 @@ test "a state holds every full cache's positions and each ring's last window - 1 m.c = c; m.sched = @import("schedule.zig").build(c); m.context = 8192; + m.paging = null; m.ring = fwd.ringSize(fwd.prompt_rows, c.window); var ks: [64]u64 = undefined; var vs: [64]u64 = undefined; diff --git a/zig/src/families/dory/dory.zig b/zig/src/families/dory/dory.zig index 0c5561f34a..59d107d779 100644 --- a/zig/src/families/dory/dory.zig +++ b/zig/src/families/dory/dory.zig @@ -11,6 +11,7 @@ pub const cuda_weights = @import("cuda_weights.zig"); pub const cuda_model = @import("cuda_model.zig"); pub const cuda_backend = @import("cuda_backend.zig"); pub const cuda_snap = @import("cuda_snap.zig"); +pub const cuda_pages = @import("cuda_pages.zig"); pub const cuda_run = @import("cuda_run.zig"); /// The native CUDA registry's entry (native/cuda.zig). pub const native = @import("cuda_native.zig"); @@ -26,6 +27,7 @@ test { _ = @import("cuda_kernels.zig"); _ = @import("cuda_backend.zig"); _ = @import("cuda_snap.zig"); + _ = @import("cuda_pages.zig"); _ = @import("cuda_run.zig"); _ = @import("cuda_native.zig"); } diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 2b987c2a58..fda3f13f8d 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -23,7 +23,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (@import("contract_deepseek_images.zig").image(name)) |bytes| return bytes; if (std.mem.eql(u8, name, "window_rows")) return kk.window_rows; if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index dfb8347c99..1b44c79a1a 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -409,6 +409,7 @@ pub fn run(gpu: check.Gpu) !void { for (&q35s) |*x| try kernels.append(a, x.kernel()); var window_fns = try cuda.window_attention.WindowAttention.load(gpu.d); defer window_fns.unload(); + try window_fns.loadPages(gpu.d); var windows: [window_k.cfgs.len]window_k.Window = undefined; for (window_k.cfgs, &windows) |f, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .g = &window_fns }; defer for (&windows) |*x| x.deinit(); diff --git a/zig/tests/cuda/contract_window_rows.zig b/zig/tests/cuda/contract_window_rows.zig index d4bf7b7eda..6af7aff9b4 100644 --- a/zig/tests/cuda/contract_window_rows.zig +++ b/zig/tests/cuda/contract_window_rows.zig @@ -16,7 +16,7 @@ const wa = cuda.window_attention; const WA = wa.WindowAttention; /// A planted fault, from contract_controls.cu (window_rows.cuh's Fault), or the served bf16-P kernel (p16). -pub const Ctl = enum { none, edge, rescale, trunc, prompt_edge, last, p16, rescale512 }; +pub const Ctl = enum { none, edge, rescale, trunc, prompt_edge, last, p16, rescale512, swap, stale }; /// A call's fixed heads, window and ring, and which kernels it runs. pub const Cfg = struct { @@ -31,6 +31,8 @@ pub const Cfg = struct { p32: bool = false, image: bool = false, ctl: Ctl = .none, + /// Positions a page, when the caches are pages (tf_wrows_*_pg after tf_wrows_put placed them in a shuffled arena). + pages: usize = 0, served: []const c.Shape, }; @@ -66,7 +68,20 @@ pub const cfgs = [_]Cfg{ } ++ ratio(12, 4) ++ ratio(20, 4) ++ ratio(32, 4) ++ ratio(32, 2) ++ [_]Cfg{ .{ .name = "control: decode merge without each chunk's exp(m_c - M), 20/4 heads (GQA 5)", .h = 20, .hk = 4, .ctl = .rescale, .served = decodeAt(20) }, .{ .name = "control: P32 prompt rows whose window starts a key early, 20/4 heads (GQA 5)", .h = 20, .hk = 4, .prompt = true, .p32 = true, .ctl = .prompt_edge, .served = promptAt(20) }, -} ++ @import("window_rows_shapes.zig").cases(Cfg); +} ++ @import("window_rows_shapes.zig").cases(Cfg) ++ paged; + +const paged_decode = &[_]c.Shape{ .{ .rows = 4, .n = 16, .k = 32768 }, .{ .rows = 16, .n = 16, .k = 3000 }, .{ .rows = 1, .n = 16, .k = 257 } }; +const paged_prompt = &[_]c.Shape{ .{ .rows = 256, .n = 16, .k = 0 }, .{ .rows = 128, .n = 16, .k = 4096 }, .{ .rows = 32, .n = 16, .k = 1000 } }; + +/// Dory's full-attention layers (16 query heads over 4 of 256, no window): caches in pages of 256 or 1,024 positions. +const paged = [_]Cfg{ + .{ .name = "cuda window_pages decode tf_wrows_part_256_pg + tf_wrows_merge_256_pg, 16/4 heads, no window, pages of 256", .h = 16, .hk = 4, .win = 0, .ring = 32768, .streams = 2, .pages = 256, .served = paged_decode }, + .{ .name = "cuda window_pages decode tf_wrows_part_256_pg + tf_wrows_merge_256_pg, 16/4 heads, no window, pages of 1024", .h = 16, .hk = 4, .win = 0, .ring = 32768, .streams = 2, .pages = 1024, .served = paged_decode }, + .{ .name = "cuda window_pages prompt tf_wrows_prompt_256_pg (P as bf16), 16/4 heads, no window, pages of 256", .h = 16, .hk = 4, .win = 0, .ring = 0, .prompt = true, .pages = 256, .served = paged_prompt }, + .{ .name = "cuda window_pages prompt tf_wrows_prompt_256_pg (P as bf16), 16/4 heads, no window, pages of 1024", .h = 16, .hk = 4, .win = 0, .ring = 0, .prompt = true, .pages = 1024, .served = paged_prompt }, + .{ .name = "control: paged decode whose first two streams' tables are swapped", .h = 16, .hk = 4, .win = 0, .ring = 32768, .streams = 2, .pages = 256, .ctl = .swap, .served = paged_decode }, + .{ .name = "control: paged prompt rows whose first page id is a later page's", .h = 16, .hk = 4, .win = 0, .ring = 0, .prompt = true, .pages = 256, .ctl = .stale, .served = paged_prompt }, +}; fn decodeAt(comptime h: usize) []const c.Shape { return &[_]c.Shape{ .{ .rows = 16, .n = h, .k = 3000 }, .{ .rows = 1, .n = h, .k = 1100 } }; @@ -89,7 +104,7 @@ fn ratio(comptime h: usize, comptime hk: usize) [3]Cfg { /// The control kernel's symbol in contract_controls.cu, if it has one. pub fn symbol(ctl: Ctl) ?[:0]const u8 { return switch (ctl) { - .none, .p16 => null, + .none, .p16, .swap, .stale => null, .edge => "tf_ctl_wrows_part_edge", .rescale => "tf_ctl_wrows_merge_rescale", .rescale512 => "tf_ctl_wrows_merge_rescale_512", @@ -121,6 +136,9 @@ pub const Rows = struct { call: wc.Call = undefined, fns: WA.Wide = undefined, dev: [6]u64 = @splat(0), + table: u64 = 0, // paged: the attention kernels' table (stream rows of `per` page ids) + per: usize = 0, + placed: ?[]const u8 = null, // paged: what tf_wrows_put got wrong, if anything scratch: u64 = 0, out: u64 = 0, half: []f64 = &.{}, @@ -139,7 +157,51 @@ pub const Rows = struct { fn win(self: *const Rows) wa.Window { const f = self.f; - return .{ .h = f.h, .hk = f.hk, .d = f.d, .win = f.win, .ring = f.ring, .scale = 1, .p32 = f.p32 and f.ctl != .p16 }; + const pg: ?wa.Pages = if (f.pages > 0) .{ .table = self.table, .page_elems = 2 * f.pages * f.hk * f.d, .tokens = f.pages, .stride = self.per } else null; + return .{ .h = f.h, .hk = f.hk, .d = f.d, .win = f.win, .ring = f.ring, .scale = 1, .p32 = f.p32 and f.ctl != .p16, .pg = pg }; + } + + /// Paged: tf_wrows_put places the caches in a shuffled arena (bit for bit, or `placed` says); the kernels read it. + fn place(self: *Rows, a: std.mem.Allocator, rng: std.Random, kc: []const u16, vc: []const u16) !void { + const f = self.f; + const row = f.hk * f.d; + const n = kc.len / row; + const span = if (f.prompt) n else f.ring; + const streams = if (f.prompt) 1 else f.streams; + self.per = (span + f.pages - 1) / f.pages; + const ids = try a.alloc(i32, streams * self.per); + for (ids, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, ids); + const pos = try a.alloc(i32, n); + const slot = try a.alloc(i32, n); + const elems = 2 * f.pages * row; + const want = try a.alloc(u16, ids.len * elems); + @memset(want, 0); + for (pos, slot, 0..) |*p, *sl, i| { + p.* = @intCast(i % span); + sl.* = @intCast(i / span); + const at = @as(usize, @intCast(ids[@as(usize, @intCast(sl.*)) * self.per + i % span / f.pages])) * elems + i % span % f.pages * row; + @memcpy(want[at..][0..row], kc[i * row ..][0..row]); + @memcpy(want[at + f.pages * row ..][0..row], vc[i * row ..][0..row]); + } + const arena = try self.rig.scratch(want.len * 2); + self.table = try self.rig.input(std.mem.sliceAsBytes(ids)); + const dpos = try self.rig.input(std.mem.sliceAsBytes(pos)); + const dslot = try self.rig.input(std.mem.sliceAsBytes(slot)); + const vcache = arena + f.pages * row * 2; + try self.g.put(self.rig.s, self.dev[1], self.dev[2], arena, vcache, dpos, dslot, self.win(), n); + const got = try self.rig.bytes(arena, want.len * 2); + self.placed = if (std.mem.eql(u8, got, std.mem.sliceAsBytes(want))) null else "tf_wrows_put placed rows differently from the table's pages"; + self.dev[1] = arena; + self.dev[2] = vcache; + switch (f.ctl) { // the attention reads a table that is not the one the rows were placed by + .swap => for (0..self.per) |i| std.mem.swap(i32, &ids[i], &ids[self.per + i]), + .stale => if (ids.len > 1) { // one page has no later page: planted from the next shape on + ids[0] = ids[1]; + } else return, + else => return, + } + self.table = try self.rig.input(std.mem.sliceAsBytes(ids)); } /// The kernels a run launches: the served ones, one of them swapped for its control. @@ -202,6 +264,8 @@ pub const Rows = struct { .stress => try self.stress(rng, kc, vc), } for ([_][]const u8{ std.mem.sliceAsBytes(q), std.mem.sliceAsBytes(kc), std.mem.sliceAsBytes(vc), std.mem.sliceAsBytes(pos), std.mem.sliceAsBytes(slot), std.mem.sliceAsBytes(hi) }, &self.dev) |bytes, *p| p.* = try self.rig.input(bytes); + self.placed = null; + if (f.pages > 0) try self.place(a, rng, kc, vc); self.scratch = if (f.prompt) 0 else try self.rig.scratch(WA.scratchBytes(self.win(), n)); self.out = try self.rig.output(n * f.h * f.d * 2); self.fns = try self.kernels(); @@ -274,6 +338,7 @@ pub const Rows = struct { /// Every output in its float64 range (gates 1 and 3); at gate 1 some, not all, of float64's outputs non-finite. pub fn check(self: *Rows, at: g.Gate) !?[]const u8 { + if (self.placed) |why_| return why_; const a = self.arena.allocator(); const out = try a.alloc(u16, self.count()); @memcpy(std.mem.sliceAsBytes(out), try self.rig.outputBytes()); @@ -329,7 +394,7 @@ fn mirrored(rng: std.Random, out: []f32) void { test "every control names a kernel of contract_controls.cu, or runs the served bf16-P kernel" { for (cfgs) |f| if (f.ctl != .none) { - try std.testing.expect(symbol(f.ctl) != null or f.ctl == .p16); + try std.testing.expect(symbol(f.ctl) != null or f.ctl == .p16 or f.pages > 0); try std.testing.expect(std.mem.startsWith(u8, f.name, "control: ")); }; } diff --git a/zig/tests/cuda/dory_pages_tests.zig b/zig/tests/cuda/dory_pages_tests.zig new file mode 100644 index 0000000000..0e436fbca5 --- /dev/null +++ b/zig/tests/cuda/dory_pages_tests.zig @@ -0,0 +1,149 @@ +//! Dory's full caches in pages against contiguous caches: every row's fp32 logits bit-equal, then served replies equal. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const dory = @import("dory"); +const check = @import("check.zig"); +const serve = @import("dory_serve_tests.zig"); +const Rig = @import("dory_tests.zig").Rig; +const Gpu = check.Gpu; + +const fwd = dory.cuda_model; +const be = dory.cuda_backend; + +const slots = 8; +/// Prompt lengths a slot takes in each wave: past a chunk, on and around a page edge, past a window. +const waves = [_][4]usize{ .{ 2600, 257, 1500, 3100 }, .{ 255, 256, 1024, 2049 } }; + +/// fp32 logits of the rows [head0, head0 + heads) of `m`'s last forward, on the host. +fn logits(gpu: Gpu, m: *const fwd.Model, heads: usize) ![]f32 { + const out = try gpu.gpa.alloc(f32, heads * m.c.vocab); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = m.s.logits, .len = out.len * 4, .borrowed = true }).download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// The same chains through both models: false when any row's logits differ at a bit. +fn both(x: *Rig, b: *fwd.Model, chains: []const fwd.Chain, prompt: bool, head0: usize, heads: usize, what: []const u8, bad: *usize) !void { + for ([_]*fwd.Model{ &x.m, b }) |m| try m.forward(&x.k, x.s, chains, prompt, head0, heads); + try x.s.synchronize(); + const gpa = x.gpu.gpa; + const l0 = try logits(x.gpu, &x.m, heads); + defer gpa.free(l0); + const l1 = try logits(x.gpu, b, heads); + defer gpa.free(l1); + if (!std.mem.eql(u8, std.mem.sliceAsBytes(l0), std.mem.sliceAsBytes(l1))) { + bad.* += 1; + if (bad.* < 12) std.debug.print("FAIL dory pages: {s} (slot {d} at {d}, {d} chains): {d} rows' logits differ from contiguous\n", .{ what, chains[0].slot, chains[0].p0, chains.len, heads }); + } +} + +/// tf-cuda-test dory-pages MODEL LOOPS TEXT N: lock-step rows over 4 slots for N rounds, then lane-core replies. +pub fn pages(gpu: Gpu, rest: []const []const u8) !void { + if (rest.len < 4) return error.MissingArguments; + const gpa = gpu.gpa; + const loops = try std.fmt.parseInt(usize, rest[1], 10); + const count = try std.fmt.parseInt(u32, rest[3], 10); + const ids = try serve.tokens(gpu, rest[0], rest[2], 4000); + defer gpa.free(ids); + const context = try be.capacity(4000 + count); + var a: Rig = undefined; + try a.initMode(gpu, rest[0], loops, context, slots, false); + defer a.deinit(); + var b = try fwd.Model.initPaged(gpa, gpu.d, a.s, &a.w, &a.k, context); + defer b.deinit(); + const per = dory.cuda_pages.pagesOf(context); + try b.attach(a.s, slots - 1, slots * per); + // slots 0 and 1 on windows 5 and 2 of the arena (contiguous kernels), slots 2 and 3 on the other pages shuffled + const all = try gpa.alloc(u32, slots * per); + defer gpa.free(all); + var next: usize = 2 * per; + for (all[0..per], all[per .. 2 * per], 0..) |*w5, *w2, i| { + w5.* = @intCast(5 * per + i); + w2.* = @intCast(2 * per + i); + } + for (0..slots * per) |id| if (id / per != 5 and id / per != 2) { + all[next] = @intCast(id); + next += 1; + }; + var prng = std.Random.DefaultPrng.init(3); + prng.random().shuffle(u32, all[2 * per ..]); + for (0..4) |s| try b.mapPages(a.s, s, all[s * per ..][0..per], 0); + const wins = b.paging.?.wins; + try check.expect(wins[0] == 5 and wins[1] == 2 and wins[2] == null and wins[3] == null, "dory pages: slots 0, 1 on windows 5, 2", .{}); + var bad: usize = 0; + var rows_checked: usize = 0; + for (waves, 0..) |wave, wv| { + var pos: [4]usize = wave; + for (wave, 0..) |len, s| { + var at: usize = 0; + while (at < len) : (at += fwd.prompt_rows) { + const end = @min(len, at + fwd.prompt_rows); + const heads = @min(16, end - at); + try both(&a, &b, &.{.{ .slot = s, .p0 = at, .tokens = ids[at..end] }}, true, end - at - heads, heads, "prompt chunk", &bad); + rows_checked += heads; + } + } + for (0..count) |r| { + var chains: [4]fwd.Chain = undefined; + var rows: usize = 0; + for (&chains, 0..) |*ch, s| { + const w = 1 + (r * 7 + s * 3 + wv) % fwd.window_rows; + const from = (pos[s] * 13 + r) % 3000; + ch.* = .{ .slot = s, .p0 = pos[s], .tokens = ids[from..][0..w] }; + pos[s] += w; + rows += w; + } + try both(&a, &b, &chains, false, 0, rows, "decode round", &bad); + rows_checked += rows; + const s = r % 4; // one stream alone: slots 0, 1 on their windows, 2, 3 through the table + const w = 1 + (r * 5 + wv) % fwd.window_rows; + const from = (pos[s] * 11 + r) % 3000; + try both(&a, &b, &.{.{ .slot = s, .p0 = pos[s], .tokens = ids[from..][0..w] }}, false, 0, w, "lone decode", &bad); + pos[s] += w; + rows_checked += w; + var two: [2]fwd.Chain = undefined; // slots 0 and 1 together: both on windows, the contiguous kernels + for (&two, 0..) |*ch, t| { + const v = 1 + (r * 3 + t * 5 + wv) % fwd.window_rows; + ch.* = .{ .slot = t, .p0 = pos[t], .tokens = ids[(pos[t] * 7 + r) % 3000 ..][0..v] }; + pos[t] += v; + rows_checked += v; + } + try both(&a, &b, &two, false, 0, two[0].tokens.len + two[1].tokens.len, "windowed pair", &bad); + } + } + std.debug.print("RESULT dory L={d} paged == contiguous: {d} rows' fp32 logits over 4 slots (slots 0, 1 on windows, prompts to 3100, {d} rounds each of 4 mixed, 1 alone and 2 windowed streams, page edges at 255/256/257), {d} rounds differ\n", .{ loops, rows_checked, 2 * count, bad }); + try check.expect(bad == 0, "dory L={d} pages: {d} forwards differ from contiguous", .{ loops, bad }); + var da = try be.Dory.init(gpa, gpu.d, &a.m, &a.k, a.s); + defer da.deinit(); + var db = try be.Dory.init(gpa, gpu.d, &b, &a.k, a.s); + defer db.deinit(); + try db.attach(slots - 1, slots * per); + db.shufflePages(9); + var buf: [8]serve.Case = undefined; + var flat: usize = 0; + const rep = try gpa.alloc(u32, 1800); + defer gpa.free(rep); + const cases = serve.mixed(ids, &buf, rep); + var differ: usize = 0; + for ([_]bool{ true, false }) |scatter| for ([_]bool{ false, true }) |drafts| for ([_]usize{ 1, 3, 8 }) |k| { + db.scatter = scatter; // scattered tables, then each lease one run of pages (the first stream's on an idle pool) + if (!scatter and !drafts and k == 1) db.m.dropGraphs(); // from here windowed streams capture flat graphs + var cs: [8]serve.Case = undefined; + for (cs[0..k], cases[0..k]) |*o, c| { + o.* = c; + o.drafts = drafts; + } + const ra = try serve.lanesRun(gpu, &da, cs[0..k], count); + defer serve.freeRun(gpa, ra); + const rb = try serve.lanesRun(gpu, &db, cs[0..k], count); + defer serve.freeRun(gpa, rb); + var n: usize = 0; + for (ra.out, rb.out) |oa, ob| n += @intFromBool(!std.mem.eql(u32, oa, ob)); + differ += n; + std.debug.print("RESULT dory L={d} paged == contiguous replies: {d} streams{s}, {s}, {d} of {d} differ ({d} tokens)\n", .{ loops, k, if (drafts) " drafted" else "", if (scatter) "scattered pages" else "windows of pages", n, k, ra.tokens }); + }; + for (db.m.flat_graphs) |g| flat += @intFromBool(g != null); + const pool = db.pool.?; + try check.expect(differ == 0 and pool.idle() and flat > 0, "dory L={d} pages: {d} replies differ, the pool is not idle or no windowed round ran ({d})", .{ loops, differ, flat }); + check.pass("dory L={d} pages: paged == contiguous, logits bit-equal and replies equal (greedy, sampled, drafted, 1/3/8 streams, scattered and windowed leases); pool idle", .{loops}); +} diff --git a/zig/tests/cuda/dory_serve_tests.zig b/zig/tests/cuda/dory_serve_tests.zig index c1de5e1307..3e32cb59aa 100644 --- a/zig/tests/cuda/dory_serve_tests.zig +++ b/zig/tests/cuda/dory_serve_tests.zig @@ -57,7 +57,7 @@ pub fn lanesRun(gpu: Gpu, db: *be.Dory, cases: []const Case, count: u32) !Run { } /// TEXT's tokens after , repeated until `need` of them. -fn tokens(gpu: Gpu, model: []const u8, path: []const u8, need: usize) ![]u32 { +pub fn tokens(gpu: Gpu, model: []const u8, path: []const u8, need: usize) ![]u32 { const gpa = gpu.gpa; var tok = try core.tokenizer.loadTokenizer(gpu.io, gpa, model); defer tok.deinit(); @@ -83,6 +83,14 @@ pub fn setup(x: *Setup, gpu: Gpu, rest: []const []const u8, need: usize, slots: try x.rig.initSlots(gpu, rest[0], x.loops, try be.capacity(need + x.count), slots); errdefer x.rig.deinit(); x.db = try be.Dory.init(gpu.gpa, gpu.d, &x.rig.m, &x.rig.k, x.rig.s); + errdefer x.db.deinit(); + if (x.rig.m.paged) { // slots lease from a shuffled pool (DORY_POOL pages: else a window each; DORY_SCATTER=0: runs) + const per = dory.cuda_pages.pagesOf(x.rig.m.context); + const pool = if (std.c.getenv("DORY_POOL")) |v| try std.fmt.parseInt(usize, std.mem.span(v), 10) else slots * per; + try x.db.attach(slots - 1, pool); + const runs = if (std.c.getenv("DORY_SCATTER")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + if (!runs) x.db.shufflePages(7); + } } pub fn teardown(x: *Setup, gpa: std.mem.Allocator) void { @@ -96,11 +104,13 @@ pub fn run(gpu: Gpu, cmd: []const u8, rest: []const []const u8) !void { if (std.mem.eql(u8, cmd, "dory-serve")) return serve(gpu, rest); if (std.mem.eql(u8, cmd, "dory-states")) return @import("dory_states_tests.zig").states(gpu, rest); if (std.mem.eql(u8, cmd, "dory-bench")) return bench(gpu, rest); + if (std.mem.eql(u8, cmd, "dory-prefill")) return prefill(gpu, rest); + if (std.mem.eql(u8, cmd, "dory-pages")) return @import("dory_pages_tests.zig").pages(gpu, rest); return error.UnknownCommand; } /// Eight mixed prompts (short, past the window, past a chunk; greedy and sampled; a repeating one for copy drafts). -fn mixed(ids: []const u32, buf: *[8]Case, rep: []u32) []const Case { +pub fn mixed(ids: []const u32, buf: *[8]Case, rep: []u32) []const Case { const sp: lanes.Sampling = .{ .seed = 11, .temperature = 0.8, .top_k = 40, .top_p = 0.95 }; const sq: lanes.Sampling = .{ .seed = 5, .temperature = 1.0, .top_k = 0, .top_p = 0.95 }; @memcpy(rep[0 .. rep.len / 2], ids[0 .. rep.len / 2]); @@ -177,6 +187,15 @@ fn roundLogits(x: *Setup, cases: []const Case) !usize { const gpa = x.rig.gpu.gpa; var chains: [8]fwd.Chain = undefined; var rows: usize = 0; + if (x.rig.m.paged) { // no stream is live: slot i takes a window of pool pages for these direct forwards + const per = dory.cuda_pages.pagesOf(x.rig.m.context); + const ids = try gpa.alloc(u32, per); + defer gpa.free(ids); + for (0..cases.len) |s| { + for (ids, 0..) |*id, j| id.* = @intCast(s * per + j); + try x.rig.m.mapPages(x.rig.s, s, ids, 0); + } + } for (cases, 0..) |c, i| { const n = 1 + i % 3; // windows of 1 to 3 rows: the prompt's own last tokens, refed past its end var at: usize = 0; @@ -223,3 +242,22 @@ fn bench(gpu: Gpu, rest: []const []const u8) !void { std.debug.print("RESULT dory bench L={d} streams {d}: {d:.1} tok/s decode together, {d:.1} a stream, {d:.2} ms a round ({d} tokens; prompts {d:.3} s)\n", .{ x.loops, k, rate, rate / @as(f64, @floatFromInt(k)), 1e3 * @as(f64, @floatFromInt(k)) / rate, r.tokens, p.secs }); } } + +/// Prompt time alone (one reply token) at 512, 2k, 8k and 16k rows: the best of three after a warm pass. +fn prefill(gpu: Gpu, rest: []const []const u8) !void { + const gpa = gpu.gpa; + var x: Setup = undefined; + try setup(&x, gpu, rest[0..4], 16384 + 1, 1); + defer teardown(&x, gpa); + for ([_]usize{ 512, 2048, 8192, 16384 }) |n| { + freeRun(gpa, try lanesRun(gpu, &x.db, &.{.{ .prompt = x.ids[0..n] }}, 1)); + var best: f64 = std.math.inf(f64); + for (0..3) |_| { + const t0 = check.now(gpu.io); // admission runs the prompt pass, before lanesRun's own clock starts + const r = try lanesRun(gpu, &x.db, &.{.{ .prompt = x.ids[0..n] }}, 1); + defer freeRun(gpa, r); + best = @min(best, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9); + } + std.debug.print("RESULT dory prefill L={d} paged={any} {d} rows: {d:.1} ms\n", .{ x.loops, x.rig.m.paged, n, best * 1e3 }); + } +} diff --git a/zig/tests/cuda/dory_states_tests.zig b/zig/tests/cuda/dory_states_tests.zig index 1a347d06d0..db7e763e2a 100644 --- a/zig/tests/cuda/dory_states_tests.zig +++ b/zig/tests/cuda/dory_states_tests.zig @@ -16,7 +16,9 @@ const past = 2600; /// Every cache of every slot overwritten with 0xFF bytes (bf16 NaN): only what a restore writes back is read. fn wipe(m: *const dory.cuda_model.Model, s: cuda.Stream) !void { + if (m.paging) |p| try m.d.check(m.d.api.cuMemsetD8Async(p.arena, 0xFF, p.pool * p.pageBytes(), s.handle), "cuMemsetD8Async"); for (m.k, m.v, 0..) |kc, vc, i| { + if (m.paging != null and m.sched.cache_kind[i] == .full) continue; // the arena above const bytes = m.slots * m.cacheLen(i) * m.c.kvDim() * 2; try m.d.check(m.d.api.cuMemsetD8Async(kc, 0xFF, bytes, s.handle), "cuMemsetD8Async"); try m.d.check(m.d.api.cuMemsetD8Async(vc, 0xFF, bytes, s.handle), "cuMemsetD8Async"); diff --git a/zig/tests/cuda/dory_tests.zig b/zig/tests/cuda/dory_tests.zig index e5038a3dcf..3dc7dd024f 100644 --- a/zig/tests/cuda/dory_tests.zig +++ b/zig/tests/cuda/dory_tests.zig @@ -20,6 +20,10 @@ pub const Rig = struct { } pub fn initSlots(r: *Rig, gpu: Gpu, model: []const u8, loops: usize, context: usize, slots: usize) !void { + return r.initMode(gpu, model, loops, context, slots, paged()); + } + + pub fn initMode(r: *Rig, gpu: Gpu, model: []const u8, loops: usize, context: usize, slots: usize, pages: bool) !void { r.gpu = gpu; r.w = try dory.cuda_weights.load(gpu.gpa, gpu.io, gpu.d, model, loops); errdefer r.w.deinit(); @@ -27,7 +31,26 @@ pub const Rig = struct { errdefer r.k.deinit(); r.s = try cuda.Stream.init(gpu.d, false); errdefer r.s.deinit(); - r.m = try fwd.Model.init(gpu.gpa, gpu.d, r.s, &r.w, &r.k, context, slots); + r.m = if (pages) try fwd.Model.initPaged(gpu.gpa, gpu.d, r.s, &r.w, &r.k, context) else try fwd.Model.init(gpu.gpa, gpu.d, r.s, &r.w, &r.k, context, slots); + } + + /// DORY_PAGED=1: the full caches are pages (the tests' other arguments stay the same). + pub fn paged() bool { + const v = std.c.getenv("DORY_PAGED") orelse return false; + return !std.mem.eql(u8, std.mem.span(v), "0"); + } + + /// A paged model's arena for `slots` slots, each slot's pages a shuffled run of the pool (no lease: a fixed map). + pub fn mapStatic(r: *Rig, slots: usize, seed: u64) !void { + if (!r.m.paged) return; + const per = dory.cuda_pages.pagesOf(r.m.context); + try r.m.attach(r.s, slots - 1, slots * per); + const ids = try r.gpu.gpa.alloc(u32, slots * per); + defer r.gpu.gpa.free(ids); + for (ids, 0..) |*id, i| id.* = @intCast(i); + var prng = std.Random.DefaultPrng.init(seed); + prng.random().shuffle(u32, ids); + for (0..slots) |s| try r.m.mapPages(r.s, s, ids[s * per ..][0..per], 0); } pub fn deinit(r: *Rig) void { @@ -53,7 +76,7 @@ pub const Rig = struct { } }; -fn argmax(x: []const f32) u32 { +pub fn argmax(x: []const f32) u32 { var best: usize = 0; for (x, 0..) |v, i| if (v > x[best]) { best = i; @@ -77,6 +100,7 @@ pub fn greedy(gpu: Gpu, model: []const u8, loops_text: []const u8, count_text: [ var rig: Rig = undefined; try rig.init(gpu, model, loops, 4096); defer rig.deinit(); + try rig.mapStatic(1, 1); const v = rig.w.config.vocab; std.debug.print("INFO dory: {d} loops, {d} prompt tokens, {d} weight bytes on the GPU\n", .{ loops, prompt.items.len, rig.w.bytes }); const n0 = prompt.items.len; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 9e3ad0738e..a4ba620f74 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -235,6 +235,8 @@ const usage = \\ dory-serve MODEL LOOPS TEXT N Dory served: eager == graphs, copy drafts == plain, 2/4/8 streams == solo \\ dory-states MODEL LOOPS TEXT N Dory kept prompt states at 8k and 32k resumed into wiped caches == fresh \\ dory-bench MODEL LOOPS TEXT N [STREAMS...] Dory greedy decode tok/s at 1/2/4/8 streams (graphs on) + \\ dory-pages MODEL LOOPS TEXT N Dory with its full caches in pages == contiguous: fp32 logits bytes, then replies + \\ dory-prefill MODEL LOOPS TEXT N Dory prompt time alone at 512, 2k, 8k and 16k rows (DORY_PAGED=1: on pages) \\ kolibri-reference MODEL IDS Kolibri-1's layer 0 and head against float64 from the checkpoint (prompt and decode rows) \\ kolibri-layers MODEL IDS Kolibri-1's window edge (513 keys) and its fifth layer's full reach, on the checkpoint \\ kolibri-full MODEL IDS N Kolibri-1: eager rows, steps, drafted, sampled, shared rounds, cuts and kept states alike