From b5eadd6a1ea4c61e3ab5f8edaafabd6a405a92eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 25 Oct 2026 03:55:07 +0100 Subject: [PATCH 1/5] cuda window_pages: Dory's full-attention rows through a page table, and the writer that places rows window_rows.cuh's decode partials, merge and prompt tiles take a Pg (a stream's row of page ids, elements a page, positions a page) as a template switch. With the switch off every kernel in window_rows.cu and contract_controls.cu compiles to the same SASS on sm_80/86/89/120/121 as before; the paged head dim 256 kernels are the new image window_pages.cu: - tf_wrows_part_256_pg, tf_wrows_merge_256_pg, tf_wrows_prompt_256_pg: a decode chunk and a prompt tile lie in one page (pages hold whole 256-key chunks), so a page id is read once a chunk or tile. - tf_wrows_put: k and v rows into the pages at their positions, a block a row, 16 bytes a thread. WindowAttention.loadPages loads the image on demand (only Dory asks); Window.pg selects the paged kernels, and the unpaged calls are unchanged. All four are registered with the kernel contract: tf_wrows_put places the caches of each case in a shuffled arena bit for bit, then the attention kernels read them through the table, each output within float64's range, at pages of 256 and 1,024 positions. Two controls (two streams' tables swapped, a stale page id) go red at gate 3. Launch records from nvcc 13.0's sm_80/86/89/120/121 images. --- docs/kernel-contract.md | 3 + zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/window_pages.cu | 44 +++++++++++ zig/kernels/cuda/window_rows.cuh | 61 ++++++++++++---- zig/src/core/contract/window_rows_records.zig | 6 +- zig/src/cuda/kernels.zig | 1 + zig/src/cuda/window_attention.zig | 69 ++++++++++++++++-- zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/contract_tests.zig | 1 + zig/tests/cuda/contract_window_rows.zig | 73 +++++++++++++++++-- 10 files changed, 231 insertions(+), 30 deletions(-) create mode 100644 zig/kernels/cuda/window_pages.cu 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/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/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..c304631c5c 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,49 @@ 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 => ids[0] = ids[1], + 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 +262,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 +336,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 +392,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: ")); }; } From 09852121d844788da9ed1e3640e88598ade20264 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 25 Oct 2026 03:55:19 +0100 Subject: [PATCH 2/5] dory: the full-attention caches in pages of the core pool, every stream leasing, a window set aside for the first At --loops 2 and 256K a stream's 14 full-attention caches take 14 GiB; the 42 sliding rings are a fixed 3,072 keys (504 MiB) whatever the context, so only the full caches page. A page holds 256 positions of every full cache (K then V each), 14 MiB at 2 loops, 24 MiB at 4; the sliding rings stay one buffer a slot, because they do not grow and their writes wrap. - Loaded.layout with engine_paged, as GLM's: the first stream's rows are pages too, and the plan sets one window of them aside, so one stream can always reach the context. Every stream leases its pages from kv_pool.Pool at its prompt pass (prompt, reply and one verify window) and gives them back when it ends; a request that cannot lease waits (LaneMemory), and F.fits is the host's admission. --parallel counts a slot's rings and table row and a page each; the prompt cache gets the rest (kept states are copies of varying size: snap_bytes 0). - One table row a slot, so a decode round mixes streams in one launch. Model.attach makes the arena, the tables and the rings of the slots the plan names; forward refuses rows past a slot's pages. Positions are written by tf_wrows_put and read through the table by the paged window kernels; the bits do not change. - Prompt-cache states copy page runs in the contiguous layout, so a kept state moves between the two forms. - TENSORFOLD_DORY_KV_PAGES=0 keeps every stream's caches contiguous and allocated at open (what `tensorfold run` and the checks' contiguous side use). tf-cuda-test dory-pages runs one model contiguous and one paged over four slots on shuffled pages (prompts to 3,100 rows, page edges at 255/256/257, 64 rounds of 1 to 16 rows) and then the lane core over 1/3/8 streams, greedy, sampled and drafted: logits bit-equal, replies equal, the pool idle. DORY_PAGED=1 runs dory-serve and dory-states on pages. --- zig/src/families/dory/cuda_backend.zig | 79 +++++++++++++-- zig/src/families/dory/cuda_kernels.zig | 1 + zig/src/families/dory/cuda_model.zig | 129 +++++++++++++++++++++++-- zig/src/families/dory/cuda_native.zig | 39 +++++++- zig/src/families/dory/cuda_pages.zig | 63 ++++++++++++ zig/src/families/dory/cuda_snap.zig | 14 +++ zig/src/families/dory/dory.zig | 2 + zig/tests/cuda/dory_pages_tests.zig | 122 +++++++++++++++++++++++ zig/tests/cuda/dory_serve_tests.zig | 41 +++++++- zig/tests/cuda/dory_states_tests.zig | 2 + zig/tests/cuda/dory_tests.zig | 28 +++++- zig/tests/cuda/main.zig | 2 + 12 files changed, 496 insertions(+), 26 deletions(-) create mode 100644 zig/src/families/dory/cuda_pages.zig create mode 100644 zig/tests/cuda/dory_pages_tests.zig diff --git a/zig/src/families/dory/cuda_backend.zig b/zig/src/families/dory/cuda_backend.zig index 7a50677565..12d984b6bc 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,7 @@ 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) /// 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 +67,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 +76,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 +87,25 @@ 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, so tables are not ascending (a check's hook). + pub fn shufflePages(x: *Dory, seed: u64) void { + 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 +123,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 +167,41 @@ 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. + 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 }; + try p.cover(lease, n, .{ .ctx = &map, .write = Map.write }); + return slot; } /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. @@ -157,12 +211,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 +382,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..b7c19676d0 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,7 +41,10 @@ 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 blocks in page 0 of the arena (the pages' table maps the rest) + 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, @@ -71,6 +75,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 +105,85 @@ 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)); + 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 * pages.tokens * c.kvDim() * 2; + vc.* = kc.* + pages.tokens * c.kvDim() * 2; + 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; + } + + /// 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; + } + /// 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 +202,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) @@ -126,6 +225,9 @@ pub const Model = struct { for (&m.graphs) |*g| if (g.*) |*e| e.deinit(); 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); } @@ -146,9 +248,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 = p.page_elems, .tokens = pages.tokens, .stride = p.per_slot }; + }; + return w; } fn scratch(m: *Model, k: *const Kernels) !void { @@ -164,7 +271,7 @@ 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), @@ -197,6 +304,7 @@ 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; } @@ -295,10 +403,13 @@ 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 pg = full and m.paging != null; // written by `put` into the pages, not by the rows kernel + 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, x.slot, out, n); + const w = m.window(full, slot); + 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 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; 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..0c53c5ffd2 --- /dev/null +++ b/zig/src/families/dory/cuda_pages.zig @@ -0,0 +1,63 @@ +//! Dory's full-attention caches in pages: one arena of the pool's pages and a table of page 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 + + pub fn pageBytes(p: Paging) usize { + return p.page_elems * 2; + } +}; + +/// Pages that hold `positions`. +pub fn pagesOf(positions: usize) usize { + return (positions + tokens - 1) / tokens; +} + +/// bf16 elements of a page across `full` caches: each cache's K block, then its V block, of `tokens` rows of `kv_dim`. +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)); +} diff --git a/zig/src/families/dory/cuda_snap.zig b/zig/src/families/dory/cuda_snap.zig index 2f58b1c69e..6ae728ad41 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,18 @@ 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, a run in each page they lie in + var q: usize = 0; + while (q < at) { + const row = q % pages.tokens; + const run = @min(at - q, pages.tokens - row); + const id = pg.ids[slot * pg.per_slot + q / pages.tokens]; + try f(ctx, cache[i] + @as(usize, id) * pg.pageBytes() + row * 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 +121,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/dory_pages_tests.zig b/zig/tests/cuda/dory_pages_tests.zig new file mode 100644 index 0000000000..d35f714ef9 --- /dev/null +++ b/zig/tests/cuda/dory_pages_tests.zig @@ -0,0 +1,122 @@ +//! 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); + var mapped: [slots][]u32 = undefined; // slots 0..3 on shuffled runs of the pool; the rest idle + const all = try gpa.alloc(u32, slots * per); + defer gpa.free(all); + for (all, 0..) |*id, i| id.* = @intCast(i); + var prng = std.Random.DefaultPrng.init(3); + prng.random().shuffle(u32, all); + for (&mapped, 0..) |*mp, i| mp.* = all[i * per ..][0..per]; + for (0..4) |s| try b.mapPages(a.s, s, mapped[s], 0); + 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; + } + } + std.debug.print("RESULT dory L={d} paged == contiguous: {d} rows' fp32 logits over 4 slots (prompts to 3100, {d} rounds of 1-16 rows, 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; + const rep = try gpa.alloc(u32, 1800); + defer gpa.free(rep); + const cases = serve.mixed(ids, &buf, rep); + var differ: usize = 0; + for ([_]bool{ false, true }) |drafts| for ([_]usize{ 1, 3, 8 }) |k| { + 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}, {d} of {d} differ ({d} tokens)\n", .{ loops, k, if (drafts) " drafted" else "", n, k, ra.tokens }); + }; + const pool = db.pool.?; + try check.expect(differ == 0 and pool.idle(), "dory L={d} pages: {d} replies differ or the pool is not idle", .{ loops, differ }); + check.pass("dory L={d} pages: paged == contiguous, logits bit-equal and replies equal (greedy, sampled, drafted, 1/3/8 streams); pool idle", .{loops}); +} diff --git a/zig/tests/cuda/dory_serve_tests.zig b/zig/tests/cuda/dory_serve_tests.zig index c1de5e1307..ef3b2f5d9e 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,13 @@ 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 their pages from a shuffled pool (DORY_POOL pages: else a window each) + 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); + x.db.shufflePages(7); + } } pub fn teardown(x: *Setup, gpa: std.mem.Allocator) void { @@ -96,11 +103,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 +186,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 +241,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 From 9f00cc6cd8a989506243cd1bdc94563704d835e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:33:33 +0100 Subject: [PATCH 3/5] kv_pool: coverRun leases a stream's pages as one run from a multiple of a step while one is free findAligned finds the highest such run; claimRun's body moves to claimAt, unchanged. A lease that already holds pages, or finds no free aligned run, falls back to cover, so nothing waits on it. --- zig/src/core/kv_pool.zig | 28 +++++++++++++++++++++++++++- zig/src/core/kv_pool_test.zig | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) 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; From 95640589625bb0ec3bff46b555107d5e558b5bd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:33:33 +0100 Subject: [PATCH 4/5] dory: a stream whose pages are one aligned run keeps main's contiguous path, the table only past that Review of #1284: 2k prefill read -0.6 to -1 % against main. The arena is now cache-major (each full cache's K over every page, then V), so page ids p..p+n-1 are rows p*256.. of one buffer, and a lease taken as one run from a multiple of pagesOf(context) is exactly main's slot w of a contiguous cache. Streams lease that way while a run is free (kv_pool.coverRun), the first stream always. - A batch whose every stream sits in such a window runs main's kernels on the window index as the slot: the fused row writes, tf_wrows_prompt_256 and the contiguous decode, no tf_wrows_put, its own graphs (flat_graphs). Any stream on scattered pages sends the whole batch through the table, which is valid for the windowed streams too: one launch set, no split, graphs keyed by rows as before. - Prompt-cache copies coalesce consecutive page ids into one copy. dory-pages maps slots 0, 1 on windows and 2, 3 on shuffled pages and adds rounds of a lone stream and a windowed pair; the lane-core replies run scattered and windowed. DORY_SCATTER=0 runs dory-serve and dory-states on windows. --- zig/src/families/dory/cuda_backend.zig | 9 ++-- zig/src/families/dory/cuda_model.zig | 69 +++++++++++++++++--------- zig/src/families/dory/cuda_pages.zig | 21 +++++++- zig/src/families/dory/cuda_snap.zig | 9 ++-- zig/tests/cuda/dory_pages_tests.zig | 47 ++++++++++++++---- zig/tests/cuda/dory_serve_tests.zig | 5 +- 6 files changed, 116 insertions(+), 44 deletions(-) diff --git a/zig/src/families/dory/cuda_backend.zig b/zig/src/families/dory/cuda_backend.zig index 12d984b6bc..f0d2c1a0f8 100644 --- a/zig/src/families/dory/cuda_backend.zig +++ b/zig/src/families/dory/cuda_backend.zig @@ -59,6 +59,7 @@ pub const Dory = struct { 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 { @@ -93,8 +94,9 @@ pub const Dory = struct { 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, so tables are not ascending (a check's hook). + /// 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]); } @@ -191,7 +193,7 @@ pub const Dory = struct { } }; - /// A free slot for the request, its pages leased now when paged. + /// 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; @@ -200,7 +202,8 @@ pub const Dory = struct { try p.promise(lease, n); errdefer p.finish(lease); var map: Map = .{ .x = x, .slot = slot }; - try p.cover(lease, n, .{ .ctx = &map, .write = Map.write }); + 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; } diff --git a/zig/src/families/dory/cuda_model.zig b/zig/src/families/dory/cuda_model.zig index b7c19676d0..0b24b7b174 100644 --- a/zig/src/families/dory/cuda_model.zig +++ b/zig/src/families/dory/cuda_model.zig @@ -41,7 +41,7 @@ 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 = &.{}, // paged: a full cache's K and V blocks in page 0 of the arena (the pages' table maps the rest) + 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 @@ -51,9 +51,11 @@ pub const Model = struct { 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| { @@ -152,11 +154,12 @@ pub const Model = struct { 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 * pages.tokens * c.kvDim() * 2; - vc.* = kc.* + pages.tokens * c.kvDim() * 2; + kc.* = p.arena + block * region; + vc.* = kc.* + region; block += 2; } else { const bytes = m.slots * m.ring * c.kvDim() * 2; @@ -177,11 +180,15 @@ pub const Model = struct { @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; + 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. @@ -222,7 +229,7 @@ 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(); @@ -238,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 { @@ -253,7 +264,7 @@ pub const Model = struct { const c = m.c; 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 = p.page_elems, .tokens = pages.tokens, .stride = p.per_slot }; + 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; } @@ -276,6 +287,7 @@ pub const Model = struct { .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), @@ -310,7 +322,11 @@ pub const Model = struct { } 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| { @@ -318,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(); @@ -349,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; @@ -368,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: { @@ -390,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); @@ -403,20 +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); - const pg = full and m.paging != null; // written by `put` into the pages, not by the rows kernel + 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, x.slot, out, n); - const w = m.window(full, slot); + 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 = if (pg) 0 else 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_pages.zig b/zig/src/families/dory/cuda_pages.zig index 0c53c5ffd2..a4d5ea47d5 100644 --- a/zig/src/families/dory/cuda_pages.zig +++ b/zig/src/families/dory/cuda_pages.zig @@ -1,4 +1,4 @@ -//! Dory's full-attention caches in pages: one arena of the pool's pages and a table of page ids a slot. +//! 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"); @@ -18,18 +18,26 @@ pub const Paging = struct { 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 of a page across `full` caches: each cache's K block, then its V block, of `tokens` rows of `kv_dim`. +/// 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(); } @@ -61,3 +69,12 @@ test "a page holds 256 positions of every full cache: 9, 14 and 24 MiB at 1, 2 a 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 6ae728ad41..ef51f64e3d 100644 --- a/zig/src/families/dory/cuda_snap.zig +++ b/zig/src/families/dory/cuda_snap.zig @@ -42,13 +42,14 @@ 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, a run in each page they lie in + 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 row = q % pages.tokens; - const run = @min(at - q, pages.tokens - row); const id = pg.ids[slot * pg.per_slot + q / pages.tokens]; - try f(ctx, cache[i] + @as(usize, id) * pg.pageBytes() + row * rb, off, run * rb); + 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; } diff --git a/zig/tests/cuda/dory_pages_tests.zig b/zig/tests/cuda/dory_pages_tests.zig index d35f714ef9..0e436fbca5 100644 --- a/zig/tests/cuda/dory_pages_tests.zig +++ b/zig/tests/cuda/dory_pages_tests.zig @@ -53,14 +53,23 @@ pub fn pages(gpu: Gpu, rest: []const []const u8) !void { defer b.deinit(); const per = dory.cuda_pages.pagesOf(context); try b.attach(a.s, slots - 1, slots * per); - var mapped: [slots][]u32 = undefined; // slots 0..3 on shuffled runs of the pool; the rest idle + // 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); - for (all, 0..) |*id, i| id.* = @intCast(i); + 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); - for (&mapped, 0..) |*mp, i| mp.* = all[i * per ..][0..per]; - for (0..4) |s| try b.mapPages(a.s, s, mapped[s], 0); + 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| { @@ -86,9 +95,23 @@ pub fn pages(gpu: Gpu, rest: []const []const u8) !void { } 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 (prompts to 3100, {d} rounds of 1-16 rows, page edges at 255/256/257), {d} rounds differ\n", .{ loops, rows_checked, 2 * count, 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(); @@ -97,11 +120,14 @@ pub fn pages(gpu: Gpu, rest: []const []const u8) !void { 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{ false, true }) |drafts| for ([_]usize{ 1, 3, 8 }) |k| { + 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; @@ -114,9 +140,10 @@ pub fn pages(gpu: Gpu, rest: []const []const u8) !void { 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}, {d} of {d} differ ({d} tokens)\n", .{ loops, k, if (drafts) " drafted" else "", n, k, ra.tokens }); + 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(), "dory L={d} pages: {d} replies differ or the pool is not idle", .{ loops, differ }); - check.pass("dory L={d} pages: paged == contiguous, logits bit-equal and replies equal (greedy, sampled, drafted, 1/3/8 streams); pool idle", .{loops}); + 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 ef3b2f5d9e..3e32cb59aa 100644 --- a/zig/tests/cuda/dory_serve_tests.zig +++ b/zig/tests/cuda/dory_serve_tests.zig @@ -84,11 +84,12 @@ pub fn setup(x: *Setup, gpu: Gpu, rest: []const []const u8, need: usize, 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 their pages from a shuffled pool (DORY_POOL pages: else a window each) + 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); - x.db.shufflePages(7); + 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); } } From 79ea7b2281f80413035759fefe73d9b6fe942db5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Mon, 26 Oct 2026 02:53:51 +0100 Subject: [PATCH 5/5] kernel contract: the stale-page control plants only where a prompt has a later page At 256 rows and no prefix the table holds one page, so ids[1] read past it: a Debug panic, undefined in ReleaseFast. With one page the control is planted from the next served shape on, where it stays red. --- zig/tests/cuda/contract_window_rows.zig | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/zig/tests/cuda/contract_window_rows.zig b/zig/tests/cuda/contract_window_rows.zig index c304631c5c..6af7aff9b4 100644 --- a/zig/tests/cuda/contract_window_rows.zig +++ b/zig/tests/cuda/contract_window_rows.zig @@ -196,7 +196,9 @@ pub const Rows = struct { 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 => ids[0] = ids[1], + .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));