dory: KV pages follow-ups (#1284)

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.

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 };

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);

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);

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));
+}

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);

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);

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();