From b430a1969cc58ff90955c6afb1d47732f8335fa1 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 11:59:52 +0100 Subject: [PATCH 1/2] dory: KV pages follow-ups: max_slots from the round geometry, attach refuses a context off the page size, window_pages only with paging (#1284) --- zig/src/families/dory/cuda_kernels.zig | 5 ++-- zig/src/families/dory/cuda_model.zig | 4 ++- zig/src/families/dory/cuda_native.zig | 2 +- zig/src/families/dory/cuda_pages.zig | 35 ++++++++++++++++++++++++-- zig/src/families/dory/cuda_run.zig | 2 +- zig/tests/cuda/dory_pages_tests.zig | 16 ++++++++++++ zig/tests/cuda/dory_tests.zig | 2 +- 7 files changed, 58 insertions(+), 8 deletions(-) diff --git a/zig/src/families/dory/cuda_kernels.zig b/zig/src/families/dory/cuda_kernels.zig index e6e7b2450a..d9dbf1ad3f 100644 --- a/zig/src/families/dory/cuda_kernels.zig +++ b/zig/src/families/dory/cuda_kernels.zig @@ -17,7 +17,8 @@ pub const Kernels = struct { to_bf16_fn: cuda.Function = undefined, draws: cuda.sample.Draws = undefined, - pub fn load(ctx: *const cuda.Context) !Kernels { + /// `paged`: also window_pages.cu (the paged caches' attention); else that module is not loaded. + pub fn load(ctx: *const cuda.Context, paged: bool) !Kernels { const d = ctx.d; var k: Kernels = undefined; k.served = false; @@ -37,7 +38,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 (paged) 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 0b24b7b174..764c6ac317 100644 --- a/zig/src/families/dory/cuda_model.zig +++ b/zig/src/families/dory/cuda_model.zig @@ -109,6 +109,7 @@ pub const Model = struct { /// 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 { + if (k.win.pages == null) return error.PagesNotLoaded; var m = blank(gpa, d, w, context, 1); m.paged = true; errdefer m.deinit(); @@ -118,6 +119,7 @@ pub const Model = struct { /// 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 { + if (k.win.pages == null) return error.PagesNotLoaded; var m = blank(gpa, d, w, context, 1); defer m.deinit(); m.paged = true; @@ -140,7 +142,7 @@ pub const Model = struct { 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 per_slot = try pages.perSlot(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); diff --git a/zig/src/families/dory/cuda_native.zig b/zig/src/families/dory/cuda_native.zig index 4965bac488..425791220f 100644 --- a/zig/src/families/dory/cuda_native.zig +++ b/zig/src/families/dory/cuda_native.zig @@ -104,7 +104,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.w = try wts.load(gpa, io, ctx.d, dir, if (o.loops) |l| @as(usize, l) else null); errdefer own.w.deinit(); if (context > own.w.config.positions + 1024) return error.PastContext; - own.k = try fwd.Kernels.load(ctx); + own.k = try fwd.Kernels.load(ctx, pagesOn()); errdefer own.k.deinit(); if (!own.k.served) return error.BuiltWithoutKernels; own.s = try cuda.Stream.init(ctx.d, false); diff --git a/zig/src/families/dory/cuda_pages.zig b/zig/src/families/dory/cuda_pages.zig index a4d5ea47d5..cc5b0f3a5a 100644 --- a/zig/src/families/dory/cuda_pages.zig +++ b/zig/src/families/dory/cuda_pages.zig @@ -3,9 +3,15 @@ const std = @import("std"); const kv = @import("core").kv_pool; const config = @import("config.zig"); const memory = @import("memory.zig"); +const fwd = @import("cuda_model.zig"); -/// Slots a round holds: its rows over the widest window (cuda_model's round_rows / window_rows). -pub const max_slots = 8; +/// Slots a round of `round` rows holds over windows of `window` rows. +pub fn slotsOf(round: usize, window: usize) usize { + return round / window; +} + +/// Slots a round holds: its rows over the widest window. +pub const max_slots = slotsOf(fwd.round_rows, fwd.window_rows); /// 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; @@ -32,6 +38,13 @@ pub fn windowOf(ids: []const u32, per: usize) ?u32 { return @intCast(ids[0] / per); } +/// Table entries a slot of `context` positions takes: whole pages only, else error.BadAttach. +pub fn perSlot(context: usize) !usize { + const n = pagesOf(context); + if (n == 0 or n * tokens != context) return error.BadAttach; + return n; +} + /// Pages that hold `positions`. pub fn pagesOf(positions: usize) usize { return (positions + tokens - 1) / tokens; @@ -78,3 +91,21 @@ test "a window: consecutive ids from a multiple of the slot's pages; a gap, a st try std.testing.expectEqual(@as(?u32, null), windowOf(&.{ 9, 8 }, 4)); try std.testing.expectEqual(@as(?u32, null), windowOf(&.{}, 4)); } + +test "a round's slots are its rows over the widest window, not a number kept beside them" { + try std.testing.expectEqual(@as(usize, 8), slotsOf(128, 16)); + try std.testing.expectEqual(@as(usize, 4), slotsOf(128, 32)); + try std.testing.expectEqual(@as(usize, 16), slotsOf(256, 16)); + try std.testing.expectEqual(fwd.round_rows / fwd.window_rows, max_slots); + try std.testing.expect(max_slots * fwd.window_rows <= fwd.round_rows); +} + +test "a slot's table is whole pages: a context off the page size refuses to attach" { + try std.testing.expectEqual(@as(usize, 132), try perSlot(33792)); + try std.testing.expectEqual(@as(usize, 4), try perSlot(1024)); + try std.testing.expectError(error.BadAttach, perSlot(33793)); + try std.testing.expectError(error.BadAttach, perSlot(1000)); + try std.testing.expectError(error.BadAttach, perSlot(255)); + try std.testing.expectError(error.BadAttach, perSlot(0)); + for (1..200) |k| try std.testing.expectEqual(k, try perSlot(k * tokens)); +} diff --git a/zig/src/families/dory/cuda_run.zig b/zig/src/families/dory/cuda_run.zig index ff2a18c50f..2945eef10d 100644 --- a/zig/src/families/dory/cuda_run.zig +++ b/zig/src/families/dory/cuda_run.zig @@ -28,7 +28,7 @@ pub fn run(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [] if (o.tokens.len == 0) return error.EmptyPrompt; var w = try wts.load(gpa, io, ctx.d, dir, o.loops); defer w.deinit(); - var k = try fwd.Kernels.load(ctx); + var k = try fwd.Kernels.load(ctx, false); defer k.deinit(); if (!k.served) return error.BuiltWithoutKernels; var s = try cuda.Stream.init(ctx.d, false); diff --git a/zig/tests/cuda/dory_pages_tests.zig b/zig/tests/cuda/dory_pages_tests.zig index 0e436fbca5..059f62bbf7 100644 --- a/zig/tests/cuda/dory_pages_tests.zig +++ b/zig/tests/cuda/dory_pages_tests.zig @@ -49,6 +49,22 @@ pub fn pages(gpu: Gpu, rest: []const []const u8) !void { var a: Rig = undefined; try a.initMode(gpu, rest[0], loops, context, slots, false); defer a.deinit(); + // paging off loads no window_pages module, and a paged model refuses kernels without it; paging on loads it + try check.expect(a.k.win.pages == null, "dory window_pages: paging off loaded the module", .{}); + var refused = true; + if (fwd.Model.initPaged(gpa, gpu.d, a.s, &a.w, &a.k, context)) |opened| { + var m = opened; + m.deinit(); + refused = false; + } else |e| if (e != error.PagesNotLoaded) return e; + try check.expect(refused, "dory window_pages: a paged model opened over kernels without the module", .{}); + { + var on = try fwd.Kernels.load(gpu.ctx, true); + defer on.deinit(); + try check.expect(on.win.pages != null, "dory window_pages: paging on did not load the module", .{}); + } + check.pass("dory window_pages: paging off loads no module and a paged model refuses it; paging on loads it", .{}); + try a.k.win.loadPages(gpu.d); 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); diff --git a/zig/tests/cuda/dory_tests.zig b/zig/tests/cuda/dory_tests.zig index 3dc7dd024f..cd9b61d3c8 100644 --- a/zig/tests/cuda/dory_tests.zig +++ b/zig/tests/cuda/dory_tests.zig @@ -27,7 +27,7 @@ pub const Rig = struct { r.gpu = gpu; r.w = try dory.cuda_weights.load(gpu.gpa, gpu.io, gpu.d, model, loops); errdefer r.w.deinit(); - r.k = try fwd.Kernels.load(gpu.ctx); + r.k = try fwd.Kernels.load(gpu.ctx, pages); errdefer r.k.deinit(); r.s = try cuda.Stream.init(gpu.d, false); errdefer r.s.deinit(); From f759da5c47f409f7b7291dd4fd8efbac1e39eecc Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:29:06 +0000 Subject: [PATCH 2/2] changelog: Dory KV pages follow-ups (#1387) --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 34636e72ec..9c4e361db8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Dory on CUDA with `TENSORFOLD_DORY_KV_PAGES=1` refuses to attach a context that is not a whole number of 256-position pages, takes its slot count from the round geometry instead of a fixed 8, and loads the paged attention kernels only when paging is on, so the default path no longer loads that module (#1387, for #1284). Thanks to @jschmied. + - Gemma 4 on CUDA serves structured output: `response_format` (`json_schema`, `json_object`) and `tool_choice: "required"` or a named function now constrain its replies, with or without a drafter (#1386, for #1384). Sampled rows are masked on the GPU before the draw; greedy rows pick the first largest allowed token on the host, the same rule as the GPU pick. Requests without a grammar take the old path unchanged. Thanks to @ThomasHalwax. - Qwen3.6-35B-A3B and Qwen3.8-27B on CUDA: `TF_QWEN_KV_SHARE=n` (off by default) keeps the prompt cache's KV rows as pool pages on the device, which a later turn's stream adopts, instead of copying them to pinned host memory and back; only the recurrent state, at most `n` snapshots of it, still goes through the host (#1357, for #1156). It needs `TF_QWEN_KV_PAGES` set to a divisor of 4096, runs every stream on pool pages, refuses `--learn-conversations`, and warns at startup when `n` is below 4 times `--parallel`. Replies are the same tokens with it on or off. On a GB10 at `--parallel 8` with Qwen3.6-35B-A3B and MTP, a kept 4,096-token state takes about 68 MiB of pinned host memory instead of 148 MiB, a warm 9k-token turn reaches its first token in the same time (0.21 s), and a cold one in 1.68 s against 1.77 s. With every stream on pages (the share, or `TF_QWEN_KV_PAGES_ALL`), `--parallel N` now serves N streams at once where it served N - 1. Thanks to @jschmied.