From 9c61b9620cf52bde5696cb992aad3ab451e90097 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:17:40 +0900 Subject: [PATCH 1/4] glm cuda: TF_GLM_KV_SHARE keeps prompts as the pool's pages in prompt_radix's store With TF_GLM_KV_SHARE=n and TF_KV_PAGES=on, GLM keeps a pass's states at its marks on the planned grid (2,048 rows, the prompt step) over 512-row pages, with a lookahead of 1 (the MTP head's row at a mark reads the token after it), in the core's radix store, at most n states past the pages. A state holds its prefix's pages by reference; its snapshot keeps only what lies past them (Plan.shared): the arena's head, the rings, the two pooled-key rows each paged cache has past the mark, and the MTP head's rows past it. A resumed pass writes nothing below its mark, since the pooled keys are written in whole pools from the step's head, so nothing is copied on write. A resume adopts the deepest state's pages at the head of the stream's lease (Pages.openWith, where TF_KV_POISON's NaN spares them) and restores the rest through the prompt pass's Resume. With it on the server makes no flat store for GLM (own_prompt_cache), as for the Qwen hybrids. One rank here; unset, every path is as before. --- zig/build/cuda.zig | 5 +- zig/build/cuda_engines.zig | 7 +- zig/build/glm_family_host.zig | 8 +- zig/glm_pages_host.zig | 4 +- zig/src/families/glm/cuda.zig | 1 + zig/src/families/glm/cuda_engine.zig | 1 + zig/src/families/glm/cuda_engine_types.zig | 2 +- zig/src/families/glm/cuda_lanes.zig | 116 +++++------ zig/src/families/glm/cuda_lanes_test.zig | 42 ++++ zig/src/families/glm/cuda_native.zig | 26 ++- zig/src/families/glm/cuda_pages.zig | 31 ++- zig/src/families/glm/cuda_pages_test.zig | 181 +++++++++++++++++- zig/src/families/glm/cuda_prefix.zig | 117 +++++++++++ zig/src/families/glm/cuda_snapshot_layout.zig | 33 ++-- zig/src/families/glm/cuda_snapshots.zig | 13 +- zig/src/families/glm/dflash.zig | 6 + 16 files changed, 492 insertions(+), 101 deletions(-) create mode 100644 zig/src/families/glm/cuda_lanes_test.zig create mode 100644 zig/src/families/glm/cuda_prefix.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 7bab04f0a4..ca14e5a728 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -103,7 +103,8 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen.zig"), .target = target, .optimize = optimize })); runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); - runner.addImport("glm_cuda", glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat)); + const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat, mods.api); + runner.addImport("glm_cuda", glm_cuda); @import("minimax_cuda.zig").targets(b, target, optimize, strip, cuda, mods.core, mods.lanes, runner); for ([_]struct { []const u8, *std.Build.Module }{ .{ "kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes) }, .{ "dory", @import("dory.zig").module(b, target, optimize, cuda, mods.core, mods.lanes) } }) |m| runner.addImport(m[0], m[1]); runner.addImport("gemma4_cuda", @import("gemma4_cuda.zig").module(b, target, optimize, cuda, mods.core, mods.lanes)); @@ -449,7 +450,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st // the contract's attention reference (tests/cuda/contract_attn.zig): its allocations under failure const attn_ref = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_attn.zig"), .target = host, .optimize = .debug, .imports = &.{.{ .name = "core", .module = mods.core }} }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = attn_ref })).step); - const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); + const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat, mods.api); @import("glm_family_host.zig").steps(b, cuda, glm, mods, step); @import("kolibri.zig").host(b, kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), native_modules.api, step); @import("gemma4_cuda.zig").host(b, @import("gemma4_cuda.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); diff --git a/zig/build/cuda_engines.zig b/zig/build/cuda_engines.zig index e805fdb9a8..b8886f3602 100644 --- a/zig/build/cuda_engines.zig +++ b/zig/build/cuda_engines.zig @@ -1,8 +1,8 @@ //! The CUDA engines a native server opens (native/cuda.zig) and the family modules they import. const std = @import("std"); -/// The GLM-5.3-Flash CUDA family over the runtime, the core, the lane core, canon's shares and core/heat.zig. -pub fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, heat: *std.Build.Module) *std.Build.Module { +/// The GLM-5.3-Flash CUDA family over the runtime, the core, the lane core, canon's shares, core/heat.zig, engine_api. +pub fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, heat: *std.Build.Module, api: *std.Build.Module) *std.Build.Module { const canon = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon.zig"), .target = target, .optimize = optimize }); const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); m.addImport("cuda", cuda); @@ -11,6 +11,7 @@ pub fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu m.addImport("lanes", lanes); m.addImport("heat", heat); m.addImport("tp", @import("tp.zig").module(b, target, optimize, cuda)); // over the engine's own cuda module + m.addImport("engine_api", api); // TF_GLM_KV_SHARE's prompt_radix store: family()'s, the server's return m; } @@ -48,7 +49,7 @@ pub fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .{ .name = "lfm2", .module = lfm2Module(b, target, optimize, cuda, core) }, .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = flashnextModule(b, target, optimize, cuda, core, lanes, qwen_vision, api) }, - .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat, api) }, .{ .name = "minimax_m3_cuda", .module = @import("minimax_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "qwen35", .module = qwen35 }, .{ .name = "kolibri1", .module = @import("kolibri.zig").module(b, target, optimize, cuda, core, lanes) }, diff --git a/zig/build/glm_family_host.zig b/zig/build/glm_family_host.zig index 766db3392f..3f179abe7e 100644 --- a/zig/build/glm_family_host.zig +++ b/zig/build/glm_family_host.zig @@ -32,7 +32,7 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, glm: *std.Build.Module, mod const page_contract_run = b.addRunArtifact(b.addTest(.{ .root_module = page_contract, .filters = &.{"GLM page contract:"} })); b.step("test-glm-page-contract-host", "GLM paged contract case and planted-control inventory without devices").dependOn(&page_contract_run.step); step.dependOn(&page_contract_run.step); - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = mods.lanes }} }); + const api = glm.import_table.get("engine_api").?; // the family's own (TF_GLM_KV_SHARE's store): one module, one set of types const admission = b.createModule(.{ .root_source_file = b.path("zig/glm_admission_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm", .module = glm }, .{ .name = "lanes", .module = mods.lanes }, .{ .name = "engine_api", .module = api } } }); const admission_run = b.addRunArtifact(b.addTest(.{ .root_module = admission, .filters = &.{"GLM CUDA admission:"} })); b.step("test-glm-admission-host", "Real GLM admission capabilities through the lane worker without GPU work").dependOn(&admission_run.step); @@ -59,14 +59,14 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, glm: *std.Build.Module, mod } pub fn compile(b: *std.Build, target: std.Build.ResolvedTarget, cuda: *std.Build.Module, mods: anytype, comptime makeGlm: anytype) void { - const fp8_gate = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/glm_ref_fp8.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm_cuda", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat) }, .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } } }); + const fp8_gate = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/glm_ref_fp8.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm_cuda", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat, mods.api) }, .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } } }); const fp8_object = b.addObject(.{ .name = "glm-fp8-gate-compile", .root_module = b.createModule(.{ .root_source_file = b.path("zig/glm_fp8_gate_compile.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "gate", .module = fp8_gate }} }) }); _ = fp8_object.getEmittedBin(); b.step("compile-glm-fp8-gate-host", "Compile synthetic FP8 write/read and FP64 device checks with inert images, never execute").dependOn(&fp8_object.step); - const glm_open = b.addObject(.{ .name = "glm-open-check", .root_module = b.createModule(.{ .root_source_file = b.path("zig/glm_open_compile.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "glm", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat) }} }) }); + const glm_open = b.addObject(.{ .name = "glm-open-check", .root_module = b.createModule(.{ .root_source_file = b.path("zig/glm_open_compile.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "glm", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat, mods.api) }} }) }); _ = glm_open.getEmittedBin(); b.step("compile-glm-native-open", "Compile the actual GLM Linux open body with empty GPU images").dependOn(&glm_open.step); - const prompt_gate = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/glm_mm_fixed_tests.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm_cuda", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat) }, .{ .name = "cuda", .module = cuda } } }); + const prompt_gate = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/glm_mm_fixed_tests.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm_cuda", .module = makeGlm(b, target, .debug, cuda, mods.core, mods.lanes, mods.heat, mods.api) }, .{ .name = "cuda", .module = cuda } } }); const prompt_gate_compile = b.addObject(.{ .name = "glm-prompt-gate-compile", .root_module = b.createModule(.{ .root_source_file = b.path("zig/glm_prompt_gate_compile.zig"), .target = target, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "gate", .module = prompt_gate }} }) }); _ = prompt_gate_compile.getEmittedBin(); b.step("compile-glm-prompt-gate-host", "Compile the actual device prompt comparison gate with inert images").dependOn(&prompt_gate_compile.step); diff --git a/zig/glm_pages_host.zig b/zig/glm_pages_host.zig index 9f14dfed6b..2563fac518 100644 --- a/zig/glm_pages_host.zig +++ b/zig/glm_pages_host.zig @@ -264,13 +264,13 @@ test "GLM pages: a failed seat binding returns its pages and lane before retry" defer back.deinit(); var stream = try lanes.Stream.init(t.allocator, .{ .id = "page-test", .prompt = &.{100}, .max_new = 8 }); defer stream.deinit(t.allocator); - try t.expectError(error.OutOfMemory, back.seat(&stream, 1, 8)); + try t.expectError(error.OutOfMemory, back.seat(&stream, 1, 8, &.{})); try t.expectEqual(@as(usize, 0), back.lanes.count()); try t.expect(f.p.pool.idle()); try t.expectEqual(@as(usize, 8), f.p.pool.census().free); try t.expectEqual(@as(u32, 1), f.e.slot); v.gpa = t.allocator; - try t.expectEqual(@as(usize, 1), (try back.seat(&stream, 1, 8)).len); + try t.expectEqual(@as(usize, 1), (try back.seat(&stream, 1, 8, &.{})).len); try t.expectEqual(@as(usize, 1), back.lanes.count()); try t.expectEqual(@as(u32, 0), f.e.slot); f.e.release(0); diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 451e7bb146..3992fabf83 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -48,6 +48,7 @@ test { _ = @import("cuda_pages_test.zig"); _ = @import("cuda_ranks_test.zig"); _ = @import("cuda_roce_multi.zig"); + _ = @import("cuda_lanes_test.zig"); _ = experts; _ = exl3; _ = vision; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 9716f9f790..b97243aa93 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -95,6 +95,7 @@ pub const Engine = struct { pages: ?*pages_mod.Pages = null, // the other streams' KV pages and sequences (cuda_pages), once the server attached them index_split: bool = false, // the next prompt step's indexer rows split across the ranks: rank 0's choice, the chunk's flag 2 index_splits: u64 = 0, // prompt steps whose indexer rows split (glm-split-check counts them) + share_pages: bool = false, // TF_GLM_KV_SHARE: a kept state leaves its prefix's pages to the prompt cache that holds them /// Opens rank `rank` of `world` over `comm`; the ranks agree on settings (cuda_ranks) before the weights load. pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, ctx: *const cuda.Context, checkpoint: []const u8, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, rank: usize, world: usize, o: Options) !*Engine { diff --git a/zig/src/families/glm/cuda_engine_types.zig b/zig/src/families/glm/cuda_engine_types.zig index 489b4edca0..dded605953 100644 --- a/zig/src/families/glm/cuda_engine_types.zig +++ b/zig/src/families/glm/cuda_engine_types.zig @@ -18,7 +18,7 @@ pub const Seq = struct { }; /// One rank's copy of a kept prompt state: the state arena's bytes and the host fields beside them. -pub const Snap = struct { lease: @import("cuda_snapshot_pages.zig").Lease, len: usize, pos: usize, mtp_len: usize, cur: []u1, bases: [2]usize }; +pub const Snap = struct { lease: @import("cuda_snapshot_pages.zig").Lease, len: usize, pos: usize, mtp_len: usize, cur: []u1, bases: [2]usize, shared: usize = 0 }; /// Where a prompt pass resumes: a kept state (every rank holds `id`) at `at`, a multiple of the chunk rows. pub const Resume = struct { id: u64, at: usize }; diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 8a1ec7e185..f676ddcc57 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -7,6 +7,7 @@ const follow = @import("cuda_follow.zig"); const fw = @import("cuda_forward.zig"); const dflash = @import("dflash.zig"); const prompt_lp = @import("cuda_prompt_logprobs.zig"); +const prefix_mod = @import("cuda_prefix.zig"); const cuda = @import("cuda"); const be = lanes.backend; @@ -38,6 +39,7 @@ const Lane = struct { held_n: usize = 0, first: ?u64 = null, fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + marks: []const u32 = &.{}, // TF_GLM_KV_SHARE: where the prompt cache keeps this stream's pass (owned) }; pub const Lanes = struct { @@ -51,10 +53,15 @@ pub const Lanes = struct { next: u64 = 0, bits: ?cuda.HostBuffer = null, // structured output's grammar words (shared_rows rows), device-read, at first use bits_dev: u64 = 0, + prefix: ?*prefix_mod.Prefix = null, // TF_GLM_KV_SHARE: the prompt cache over the pool's pages (rank 0's) pub fn deinit(self: *Lanes) void { if (self.bits) |*b| b.free(); + var it = self.lanes.valueIterator(); + while (it.next()) |l| self.x.e.gpa.free(l.marks); self.lanes.deinit(self.x.e.gpa); + if (self.prefix) |x| x.deinit(); // its pages back to the pool before the pool goes + self.prefix = null; } /// A window's grammar rows (rows of the vocabulary's words) at row `row` of the bits, device-read; 0 for none. @@ -190,13 +197,14 @@ pub const Lanes = struct { try sink.prepare(); self.x.e.prompt_logprobs = sink.want(); defer self.x.e.prompt_logprobs = null; - const first = try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s)); + var keeps: prefix_mod.Keeps = undefined; + const first = try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, self.marksOf(s, &keeps)); self.lanes.getPtr(s).?.first = self.take(first); } fn prefillManyFn(ptr: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { const self = of(ptr); - if (!self.x.e.m.tiled or ss.len > engine.max_together) return false; + if (!self.x.e.m.tiled or ss.len > engine.max_together or self.prefix != null) return false; // the cache keeps a pass alone for (ss) |s| if (!canBatch(s)) return false; var opened: usize = 0; errdefer for (ss[0..opened]) |s| releaseFn(ptr, s); @@ -238,7 +246,8 @@ pub const Lanes = struct { try sink.prepare(); self.x.e.prompt_logprobs = sink.want(); defer self.x.e.prompt_logprobs = null; - if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, marksOf(s))) |first| { + var keeps: prefix_mod.Keeps = undefined; + if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, self.marksOf(s, &keeps))) |first| { l.first = self.take(first); l.fill = null; return true; @@ -255,7 +264,8 @@ pub const Lanes = struct { const at = l.fill orelse return; if (at <= s.cached or s.reuse.hook == null) return; try self.x.bind(l.slot); - _ = self.x.chunk(s.prompt(), at, null, .{ .ptr = null, .check = cancelNow }, marksOf(s)) catch |err| { + var keeps: prefix_mod.Keeps = undefined; + _ = self.x.chunk(s.prompt(), at, null, .{ .ptr = null, .check = cancelNow }, self.marksOf(s, &keeps)) catch |err| { if (err == error.Cancelled) return; return err; }; @@ -270,7 +280,16 @@ pub const Lanes = struct { fn opening(self: *Lanes, s: *lanes.Stream) !?engine.Resume { const ids = s.prompt(); if (ids.len + s.max_new + max_rows >= self.x.e.capacity) return error.PromptTooLong; - const pages = try self.seat(s, ids.len, s.max_new); + const gpa = self.x.e.gpa; + const share = self.cached(s); + var owner: prefix_mod.Owner = .{}; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(gpa); + // TF_GLM_KV_SHARE: the deepest kept state's pages, held for this call until the seat takes them + const marks: []const u32 = if (share) (try self.prefix.?.begin(ids, s.history_len, s.shared_prefixes, &owner, &adopt)).marks else &.{}; + var held = true; + errdefer if (held) gpa.free(marks); + const pages = try self.seat(s, ids.len, s.max_new, adopt.items); errdefer { self.x.e.release(self.lanes.get(s).?.slot); _ = self.lanes.remove(s); @@ -279,12 +298,29 @@ pub const Lanes = struct { try self.x.beginOn(s.sampling, pages); s.cached = 0; s.prompt_rows.clearRetainingCapacity(); + const l = self.lanes.getPtr(s).?; + gpa.free(l.marks); + l.marks = marks; + held = false; + if (share) { + const k = owner.pending orelse return null; + s.cached = k.at; + return .{ .id = k.id, .at = k.at }; + } return if (s.prompt_logprobs != null) null else self.resumeFrom(s, ids.len); // prompt rows need every row } - /// The stream's slot, its pages, then its sequence: a slot past the paged ones waits before binding makes a window. - pub fn seat(self: *Lanes, s: *const lanes.Stream, prompt: usize, max_new: usize) ![]const u32 { - const l = try self.claim(s); + /// Whether the prompt cache over the pool's pages keeps and resumes this pass: text alone, no prompt rows asked. + fn cached(self: *const Lanes, s: *const lanes.Stream) bool { + return self.prefix != null and s.images.len == 0 and s.prompt_logprobs == null; + } + + /// The stream's slot, its pages (the cache's `adopted` first, given back on a refusal), then its sequence. + pub fn seat(self: *Lanes, s: *const lanes.Stream, prompt: usize, max_new: usize, adopted: []const u32) ![]const u32 { + const l = self.claim(s) catch |err| { + if (self.x.e.pages) |p| p.give(adopted); + return err; + }; l.pending_rows = null; l.held_n = 0; l.fill = null; @@ -293,17 +329,20 @@ pub const Lanes = struct { if (self.x.e.pages) |p| p.close(l.slot); _ = self.lanes.remove(s); } - const pages = try self.pagesFor(l.slot, prompt, max_new); + const pages = try self.pagesFor(l.slot, prompt, max_new, adopted); try self.x.bind(l.slot); return pages; } /// A slot's pages for a prompt and its reply (none without a pool); LaneMemory: the pool lacks them now. - fn pagesFor(self: *Lanes, slot: u32, prompt: usize, max_new: usize) ![]const u32 { + fn pagesFor(self: *Lanes, slot: u32, prompt: usize, max_new: usize, adopted: []const u32) ![]const u32 { const p = self.x.e.pages orelse return &.{}; - if (!p.paged(slot)) return error.LaneMemory; // every paged sequence is taken + if (!p.paged(slot)) { // every paged sequence is taken + p.give(adopted); + return error.LaneMemory; + } p.close(slot); // a stream that opens again gives its earlier pages back first - return p.open(slot, prompt, max_new); + return p.openWith(slot, prompt, max_new, adopted); } /// The lane host's admission: with a pool, a paged sequence and the request's pages are free. @@ -321,7 +360,12 @@ pub const Lanes = struct { } const Sink = @import("cuda_prompt_sink.zig").Sink; - fn marksOf(s: *lanes.Stream) ?engine.Marks { + fn marksOf(self: *Lanes, s: *lanes.Stream, keeps: *prefix_mod.Keeps) ?engine.Marks { + if (self.cached(s)) { + const l = self.lanes.getPtr(s) orelse return null; + keeps.* = .{ .x = self.prefix.?, .slot = l.slot, .prompt = s.prompt() }; + return if (l.marks.len == 0) null else .{ .at = l.marks, .ctx = keeps, .call = prefix_mod.Keeps.at, .interrupted = prefix_mod.Keeps.at }; + } return if (s.reuse.hook == null) null else .{ .at = s.reuse.marks, .ctx = s, @@ -507,7 +551,7 @@ pub const Lanes = struct { held[i] = l; } try self.x.block(reqs[0..requests.len]); - const cut = confidence(); + const cut = dflash.confidence(); for (requests, held[0..requests.len]) |r, l| { l.held_n = @min(r.depth, max_held); var chain: f64 = 1; // glm_cuda's rule: past the first, a chain stops where its product falls under the cut @@ -519,12 +563,6 @@ pub const Lanes = struct { } } - /// TF_GLM_DRAFT_CONFIDENCE: where a DFlash2 chain's chance product stops it (glm_cuda's 0.7). - fn confidence() f64 { - const v = std.c.getenv("TF_GLM_DRAFT_CONFIDENCE") orelse return 0.7; - return std.fmt.parseFloat(f64, std.mem.span(v)) catch 0.7; - } - /// Each held draft's chance, DFlash2's or the MTP head's under --mtp-confidence (else the per-depth rates stand). fn chancesFn(ptr: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { const self = of(ptr); @@ -545,6 +583,7 @@ pub const Lanes = struct { fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { const self = of(ptr); const kv = self.lanes.fetchRemove(s) orelse return; + self.x.e.gpa.free(kv.value.marks); self.x.release(kv.value.slot) catch |err| std.log.err("releasing slot {d}: {s}", .{ kv.value.slot, @errorName(err) }); if (self.lanes.count() == 0) self.x.idle() catch |err| std.log.err("the other ranks' idle step: {s}", .{@errorName(err)}); } @@ -554,43 +593,6 @@ comptime { std.debug.assert(max_rows <= fw.window_rows); } -test "a lane's slot is chosen before its map entry exists" { - var storage: [8192]u8 = @splat(0); - var fixed = std.heap.FixedBufferAllocator.init(&storage); - var e: engine.Engine = undefined; - e.gpa = fixed.allocator(); - var self: Lanes = .{ .x = .{ .e = &e, .link = undefined } }; - defer self.lanes.deinit(e.gpa); - var streams: [3]lanes.Stream = undefined; - const Probe = struct { - var bad = false; - fn choose(l: *const Lanes) u32 { - bad = l.lanes.count() != 0; - return 0; - } - }; - _ = try self.claimUsing(&streams[0], Probe.choose); - try std.testing.expect(!Probe.bad); - try std.testing.expectEqual(@as(u32, 0), (try self.claim(&streams[0])).slot); - try std.testing.expectEqual(@as(u32, 1), (try self.claim(&streams[1])).slot); - _ = self.lanes.remove(&streams[0]); - try std.testing.expectEqual(@as(u32, 0), (try self.claim(&streams[2])).slot); - try std.testing.expectEqual(@as(u32, 1), (try self.claim(&streams[1])).slot); -} - -test "a failed lane insertion leaves no entry, and a retry takes slot zero" { - var fail = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); - var e: engine.Engine = undefined; - e.gpa = fail.allocator(); - var self: Lanes = .{ .x = .{ .e = &e, .link = undefined } }; - defer self.lanes.deinit(e.gpa); - var stream: lanes.Stream = undefined; - try std.testing.expectError(error.OutOfMemory, self.claim(&stream)); - try std.testing.expectEqual(@as(usize, 0), self.lanes.count()); - fail.fail_index = std.math.maxInt(usize); - try std.testing.expectEqual(@as(u32, 0), (try self.claim(&stream)).slot); -} - /// Eligibility for the maintained fresh-prompt batching path. pub fn canBatch(s: anytype) bool { return s.prompt_logprobs == null and s.images.len == 0 and s.reuse.saved == null and s.prompt_len <= engine.prompt_rows; diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig new file mode 100644 index 0000000000..14e3d15a69 --- /dev/null +++ b/zig/src/families/glm/cuda_lanes_test.zig @@ -0,0 +1,42 @@ +//! The lane backend's slots: chosen before a stream's map entry exists, and none left behind by a failed insertion. +const std = @import("std"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const Lanes = @import("cuda_lanes.zig").Lanes; + +test "a lane's slot is chosen before its map entry exists" { + var storage: [8192]u8 = @splat(0); + var fixed = std.heap.FixedBufferAllocator.init(&storage); + var e: engine.Engine = undefined; + e.gpa = fixed.allocator(); + var self: Lanes = .{ .x = .{ .e = &e, .link = undefined } }; + defer self.lanes.deinit(e.gpa); + var streams: [3]lanes.Stream = undefined; + const Probe = struct { + var bad = false; + fn choose(l: *const Lanes) u32 { + bad = l.lanes.count() != 0; + return 0; + } + }; + _ = try self.claimUsing(&streams[0], Probe.choose); + try std.testing.expect(!Probe.bad); + try std.testing.expectEqual(@as(u32, 0), (try self.claim(&streams[0])).slot); + try std.testing.expectEqual(@as(u32, 1), (try self.claim(&streams[1])).slot); + _ = self.lanes.remove(&streams[0]); + try std.testing.expectEqual(@as(u32, 0), (try self.claim(&streams[2])).slot); + try std.testing.expectEqual(@as(u32, 1), (try self.claim(&streams[1])).slot); +} + +test "a failed lane insertion leaves no entry, and a retry takes slot zero" { + var fail = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var e: engine.Engine = undefined; + e.gpa = fail.allocator(); + var self: Lanes = .{ .x = .{ .e = &e, .link = undefined } }; + defer self.lanes.deinit(e.gpa); + var stream: lanes.Stream = undefined; + try std.testing.expectError(error.OutOfMemory, self.claim(&stream)); + try std.testing.expectEqual(@as(usize, 0), self.lanes.count()); + fail.fail_index = std.math.maxInt(usize); + try std.testing.expectEqual(@as(u32, 0), (try self.claim(&stream)).slot); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index c985ead7a3..28feff0a4f 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -14,6 +14,7 @@ const shard = @import("cuda_shard.zig"); const dflash = @import("dflash.zig"); const dflash_load = @import("dflash_load.zig"); const pages_mod = @import("cuda_pages.zig"); +const prefix_mod = @import("cuda_prefix.zig"); const kv = @import("core").kv_pool; pub const model_type = "glm5_next"; @@ -66,6 +67,7 @@ pub const Loaded = struct { layout: ?kv.Layout = null, // every stream's memory as one pool of KV pages (null: a window each, by TF_KV_PAGES) engine_paged: bool = false, // with a layout, the engine's own sequence takes pages too: the plan sets a window aside host_snaps: bool = false, // TF_GLM_HOST_CACHE: kept states in pinned host memory, --prompt-cache-gib of it + own_prompt_cache: bool = false, // TF_GLM_KV_SHARE: the family keeps prompts over the pool's pages, no host store }; /// The pool behind every stream: the engine's own sequence and `streams` more on `pages` KV pages, on every rank. @@ -74,6 +76,16 @@ pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { const l = own.layout orelse return error.NoPages; try own.lanes.x.pool(streams + 1, pages, l.pageBytes()); own.e.pages = try pages_mod.Pages.init(own.gpa, own.e, l, streams + 1, pages, true); + if (own.share > 0) try share(own); +} + +/// TF_GLM_KV_SHARE: the pool's prompt cache, its `own.share` snapshots sized at the window's last mark (rings whole). +fn share(own: *Owned) !void { + const e = own.e; + e.share_pages = true; + // cc-defer: the snapshots take the room set aside for --prompt-cache-gib; from the plan once it plans them + try readySnapshots(own, own.share * e.stateBytes(e.capacity - e.capacity % snap_grid)); + own.lanes.prefix = try prefix_mod.Prefix.init(e.gpa, e, e.pages.?, snap_grid, own.share); } /// Whether a request's caches fit now (the lane host's admission, on its engine thread). @@ -134,6 +146,7 @@ const Owned = struct { link: follow.Link, lanes: glm_lanes.Lanes, layout: ?kv.Layout, + share: usize, // TF_GLM_KV_SHARE: states the prompt cache over the pool's pages keeps (0: the host store) }; const env = @import("sys").env; @@ -180,7 +193,8 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ return err; }; set.kv8 = o.kv8; - return startup.open(io, p, if (std.c.getenv("TF_KV_PAGES")) |v| std.mem.span(v) else null, set, cuda.tp_watch, StartArgs{ .gpa = gpa, .io = io, .ctx = ctx, .dir = dir, .o = o, .p = p, .roce_bytes = roce_bytes }, openRank); + const share_snaps = try prefix_mod.sharing(); + return startup.open(io, p, if (std.c.getenv("TF_KV_PAGES")) |v| std.mem.span(v) else null, set, cuda.tp_watch, StartArgs{ .gpa = gpa, .io = io, .ctx = ctx, .dir = dir, .o = o, .p = p, .roce_bytes = roce_bytes, .share = share_snaps }, openRank); } const StartArgs = struct { @@ -191,6 +205,7 @@ const StartArgs = struct { o: Options, p: Place, roce_bytes: usize, + share: usize, }; /// After liveness and page settings: NCCL, the engine and the link, then a follower's loop or rank zero's lanes. @@ -202,10 +217,15 @@ fn openRank(a: StartArgs, set: ranks.Settings) !Loaded { const o = a.o; const p = a.p; const roce_bytes = a.roce_bytes; + // TF_GLM_KV_SHARE keeps prompts as the pool's pages: refused before the weights load, every rank told (startup) + if (a.share > 0 and !set.pages) return error.GlmShareNeedsPages; + // cc-defer: one rank; a follower keeps its fixed parts and maps adopted pages when the TP stage lands (V9F_RADIX) + if (a.share > 0 and p.world > 1) return error.GlmShareOneRank; cuda.memory_receipt.emit(@intCast(p.rank), .load_begin, true); // the four phases tools/memory_receipts.py reads const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; + own.share = a.share; // every connection registered at init, while memory is free (lazily, a full GB10's ibv_reg_mr can refuse it) if (std.c.getenv("NCCL_RUNTIME_CONNECT") == null) _ = env.setenv("NCCL_RUNTIME_CONNECT", "0", 0); // NCCL's cuMem host buffers wore out GB10's driver pool under load (NV_ERR_NO_MEMORY; next ranks failed to start) @@ -332,6 +352,7 @@ fn finishRank(a: OpenArgs) !Loaded { .layout = own.layout, .engine_paged = set.pages, .host_snaps = snapshots.pages.onHost(), + .own_prompt_cache = own.share > 0, }; } @@ -430,6 +451,9 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { error.NoCommId => "GLM-5.3-Flash's ranks meet over NCCL: with TF_GLM_WORLD 2 or more, set NCCL_COMM_ID to rank 0's address:port on every host", error.RankMismatch => "the ranks' settings differ (each rank's log names the one): start them alike", error.Fp8OnPages => "--kv-cache fp8 keeps a window per stream: unset TF_KV_PAGES, or set it to auto or off", + error.BadGlmKvShare => "TF_GLM_KV_SHARE: give how many states the prompt cache over the KV pages keeps, or 0 or unset for off", + error.GlmShareNeedsPages => "TF_GLM_KV_SHARE keeps prompts as KV pool pages: on one rank set TF_KV_PAGES=on as well", + error.GlmShareOneRank => "TF_GLM_KV_SHARE runs on one rank for now: serve TF_GLM_WORLD=1, or unset it", error.TooManyRanks => std.fmt.comptimePrint( "GLM-5.3-Flash on CUDA runs on at most {d} ranks: set TF_GLM_WORLD to {d} or fewer", .{ ranks.max_world, ranks.max_world }, diff --git a/zig/src/families/glm/cuda_pages.zig b/zig/src/families/glm/cuda_pages.zig index 1ceaca78bc..a6ac11642d 100644 --- a/zig/src/families/glm/cuda_pages.zig +++ b/zig/src/families/glm/cuda_pages.zig @@ -133,15 +133,32 @@ pub const Pages = struct { /// Rank 0: paged slot `slot` takes the pages of a prompt, its reply and a last window; LaneMemory: not yet. pub fn open(p: *Pages, slot: u32, prompt: usize, max_tokens: usize) ![]const u32 { + return p.openWith(slot, prompt, max_tokens, &.{}); + } + + /// The prompt cache's pages a resume held, given back when no lease took them. + pub fn give(p: *Pages, adopted: []const u32) void { + for (adopted) |id| p.pool.release(id); + } + + /// `open` with the prompt cache's `adopted` pages (held for this call) first: their rows stay as the cache kept. + pub fn openWith(p: *Pages, slot: u32, prompt: usize, max_tokens: usize, adopted: []const u32) ![]const u32 { + var held = adopted.len > 0; + errdefer if (held) for (adopted) |id| p.pool.release(id); if (!p.paged(slot) or p.used[slot]) return error.SlotTaken; const n = need(p.layout, prompt, max_tokens); - if (n > tableLen(p.e.capacity)) return error.PromptTooLong; + if (n > tableLen(p.e.capacity) or adopted.len > n) return error.PromptTooLong; const l = &p.leases[slot]; - try p.pool.promise(l, n); + try p.pool.promise(l, n - adopted.len); errdefer p.pool.finish(l); + if (adopted.len > 0) { + held = false; + try p.pool.adopt(l, adopted); // the lease holds them now (its finish releases them) + } var up: Upload = .{ .p = p, .slot = slot }; + if (adopted.len > 0) try Upload.write(&up, l.ids.items, 0); try p.pool.cover(l, n, .{ .ctx = &up, .write = Upload.write }); - try p.mapped(slot, l.ids.items); + try p.mapped(slot, l.ids.items, adopted.len); return l.ids.items; } @@ -156,14 +173,14 @@ pub const Pages = struct { errdefer l.ids.clearRetainingCapacity(); var up: Upload = .{ .p = p, .slot = slot }; try Upload.write(&up, l.ids.items, 0); - try p.mapped(slot, l.ids.items); + try p.mapped(slot, l.ids.items, 0); } - /// The slot's sequence reaches its pages' rows (NaN first under TF_KV_POISON). - fn mapped(p: *Pages, slot: u32, ids: []const u32) !void { + /// The slot's sequence reaches its pages' rows (past the `kept` first, NaN first under TF_KV_POISON). + fn mapped(p: *Pages, slot: u32, ids: []const u32, kept: usize) !void { const st = &p.e.seqs.items[slot].?.st; const page = p.layout.pageBytes(); - if (p.poison) for (ids) |id| try p.e.d.check(p.e.d.api.cuMemsetD32Async(p.arena.?.ptr + id * page, 0xFFFF_FFFF, page / 4, p.e.stream.handle), "cuMemsetD32Async"); + if (p.poison) for (ids[kept..]) |id| try p.e.d.check(p.e.d.api.cuMemsetD32Async(p.arena.?.ptr + id * page, 0xFFFF_FFFF, page / 4, p.e.stream.handle), "cuMemsetD32Async"); st.pages = ids; st.capacity = @min(ids.len * state.page_rows, p.e.capacity); p.used[slot] = true; diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index 7acda05fb6..f49f22afda 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -4,6 +4,11 @@ const cuda = @import("cuda"); const weights = @import("cuda_weights.zig"); const shard = @import("cuda_shard.zig"); const pages_mod = @import("cuda_pages.zig"); +const Plan = @import("cuda_snapshot_layout.zig").Plan; +const page_rows = @import("cuda_state.zig").page_rows; +const sparse = @import("cuda_sparse.zig"); +const prefix_mod = @import("cuda_prefix.zig"); +const Kept = @import("cuda_engine.zig").Kept; const Engine = @import("cuda_engine.zig").Engine; const Lanes = @import("cuda_lanes.zig").Lanes; const lanes = @import("lanes"); @@ -202,17 +207,17 @@ test "TF_KV_PAGES: a stream past the paged sequences waits for one, and its slot var ln: Lanes = .{ .x = .{ .e = &f.e, .link = undefined } }; defer ln.lanes.deinit(f.e.gpa); var streams: [3]lanes.Stream = undefined; // as many at once as lanes: the engine's own sequence and one more - try expectEqual(@as(usize, 2), (try ln.seat(&streams[0], 600, 100)).len); + try expectEqual(@as(usize, 2), (try ln.seat(&streams[0], 600, 100, &.{})).len); try expectEqual(@as(u32, 0), ln.lanes.get(&streams[0]).?.slot); // the engine's own takes pages too - try expectEqual(@as(usize, 2), (try ln.seat(&streams[1], 600, 100)).len); + try expectEqual(@as(usize, 2), (try ln.seat(&streams[1], 600, 100, &.{})).len); const live = F.live; - try std.testing.expectError(error.LaneMemory, ln.seat(&streams[2], 600, 100)); // slot 2: past the paged ones + try std.testing.expectError(error.LaneMemory, ln.seat(&streams[2], 600, 100, &.{})); // slot 2: past the paged ones try expectEqual(live, F.live); // nothing allocated for it try expectEqual(@as(usize, 2), f.e.seqs.items.len); try expectEqual(@as(usize, 2), ln.lanes.count()); _ = ln.lanes.remove(&streams[0]); // a stream ends: the one that waited takes its paged sequence p.close(0); - _ = try ln.seat(&streams[2], 600, 100); + _ = try ln.seat(&streams[2], 600, 100, &.{}); try expectEqual(@as(u32, 0), ln.lanes.get(&streams[2]).?.slot); try expectEqual(@as(u32, 0), f.e.slot); } @@ -228,8 +233,8 @@ test "TF_KV_PAGES: the lane host admits a request by the pool alone, the engine' defer ln.lanes.deinit(f.e.gpa); try expect(ln.fits(3000, 100)); // 7 of the 8 pages var streams: [2]lanes.Stream = undefined; - _ = try ln.seat(&streams[0], 600, 100); // slot 0: 2 pages - _ = try ln.seat(&streams[1], 2000, 100); // slot 1: 5 pages + _ = try ln.seat(&streams[0], 600, 100, &.{}); // slot 0: 2 pages + _ = try ln.seat(&streams[1], 2000, 100, &.{}); // slot 1: 5 pages _ = ln.lanes.remove(&streams[0]); p.close(0); try expect(!ln.fits(3000, 100)); // the engine's own is free, but only 3 pages are @@ -247,6 +252,170 @@ test "TF_KV_POISON: every page a stream takes is filled with NaN before its rows try expectEqual(@as(usize, 3), F.fills); } +test "a stream resuming a kept prompt takes the prompt cache's pages first: its table starts with them, NaN spares them" { + var f: Fake = undefined; + try f.init(std.testing.allocator); + const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); + const p = try f.pool(l, 3, 8, true); + defer f.deinit(p); + p.poison = true; + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, try p.open(1, 1500, 100)); + p.pool.retain(0); // the prompt cache keeps the first two pages + p.pool.retain(1); + p.close(1); + p.pool.retain(0); // its resume holds them once more for the stream + p.pool.retain(1); + F.fills = 0; + const ids = try p.openWith(2, 1500, 100, &.{ 0, 1 }); + try expectEqual(@as(usize, 4), ids.len); + try std.testing.expectEqualSlices(u32, &.{ 0, 1 }, ids[0..2]); + try std.testing.expectEqualSlices(u32, ids, F.table(f.e.seqs.items[2].?.st.table, 4)); + try expectEqual(@as(usize, 2), F.fills); // the two fresh pages alone + try expectEqual(@as(u32, 2), p.pool.holders(0)); // the cache and the stream + p.close(2); + try expectEqual(@as(u32, 1), p.pool.holders(0)); // the cache's alone again + p.pool.release(0); + p.pool.release(1); + try expect(p.pool.idle() and p.pool.census().free == 8); +} + +test "a resume that cannot have its fresh pages gives the prompt cache's back, and nothing is mapped" { + var f: Fake = undefined; + try f.init(std.testing.allocator); + const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); + const p = try f.pool(l, 3, 8, true); + defer f.deinit(p); + _ = try p.open(1, 1500, 100); // pages 0 to 3 + p.pool.retain(0); + p.close(1); + _ = try p.open(0, 3000, 100); // 7 of the 7 free pages + p.pool.retain(0); // the resume's hold + try std.testing.expectError(error.LaneMemory, p.openWith(2, 1500, 100, &.{0})); + try expectEqual(@as(u32, 1), p.pool.holders(0)); // the cache's alone + try expectEqual(@as(usize, 0), f.e.seqs.items[2].?.st.capacity); + p.close(0); + p.pool.release(0); + try expect(p.pool.idle() and p.pool.census().free == 8); +} + +test "a kept state's copy leaves the prompt cache's pages out: past them, two pools a paged cache and the head's rows" { + var f: Fake = undefined; + try f.init(std.testing.allocator); + const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); + const p = try f.pool(l, 2, 8, true); + defer f.deinit(p); + const ids = try p.open(1, 3000, 100); // 7 pages + const st = &f.e.seqs.items[1].?.st; + const c = &f.e.w.config; + const at = 2048; + const held = at / page_rows; + const whole = try Plan.init(st, c, at, at, 0); + const kept = try Plan.init(st, c, at, at, held); + const latent = c.kv_lora * 2; // a row's + const pooled = c.i_dim * 2; // a pool's + // the DSA layer's and the MTP head's caches: their 2,048 latent rows and 512 pools stay on the held pages + try expectEqual(whole.len - kept.len, 2 * (at * latent + at / sparse.pool * pooled)); + var parts: usize = 0; + for (kept.parts[0..kept.count]) |part| if (part.span != 0) { + parts += 1; + try expectEqual(2 * pooled, part.len); // the pooled keys' two rows past the last whole pool + try expectEqual(part.live + @as(u64, ids[held]) * part.page, kept.piece(part, 0).live); // the next page's first + }; + try expectEqual(@as(usize, 2), parts); + // the head 8 rows past the mark: those rows are copied from the next page too + const ahead = try Plan.init(st, c, at, at + 8, held); + var rows: usize = 0; + for (ahead.parts[0..ahead.count]) |part| if (part.span == page_rows * latent) { + rows += part.len / latent; + try expectEqual(part.live + @as(u64, ids[held]) * part.page, ahead.piece(part, 0).live); + }; + try expectEqual(@as(usize, 8), rows); + try std.testing.expectError(error.SnapshotShape, Plan.init(st, c, at, at, held + 1)); // a page past the mark +} + +/// The prompt cache's states without a device: each a Kept, live until the store drops it. +const States = struct { + gpa: std.mem.Allocator, + next: u64 = 1, + live: usize = 0, + + fn of(p: *anyopaque) *States { + return @ptrCast(@alignCast(p)); + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(p: *anyopaque, _: ?*anyopaque, at: u32) anyerror!*anyopaque { + const x = of(p); + const k = try x.gpa.create(Kept); + k.* = .{ .id = x.next, .at = at }; + x.next += 1; + x.live += 1; + return k; + } + fn drop(p: *anyopaque, saved: *anyopaque) void { + const x = of(p); + x.gpa.destroy(@as(*Kept, @ptrCast(@alignCast(saved)))); + x.live -= 1; + } + const vtable: prefix_mod.radix.Snapshots.VTable = .{ .bytes = bytes, .save = save, .restore = prefix_mod.Prefix.restore, .drop = drop }; +}; + +test "TF_GLM_KV_SHARE: a pass's pages below its mark go to the prompt cache, and a later pass resumes on those pages" { + const gpa = std.testing.allocator; + var f: Fake = undefined; + try f.init(gpa); + const l = try pages_mod.layout(gpa, &f.e, 0); + const p = try f.pool(l, 3, 8, true); + defer f.deinit(p); + var states: States = .{ .gpa = gpa }; + const x = try prefix_mod.Prefix.initWith(gpa, p, 2048, 4, .{ .ptr = &states, .vtable = &States.vtable }); + defer x.deinit(); + var owner: prefix_mod.Owner = .{}; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(gpa); + var a: [3000]u32 = undefined; + for (&a, 1..) |*t, i| t.* = @intCast(i); + // a cold pass: a mark on the grid below its end, where the tree takes the pages below it and the state + const cold = try x.begin(&a, 0, &.{}, &owner, &adopt); + defer gpa.free(cold.marks); + try std.testing.expectEqualSlices(u32, &.{2048}, cold.marks); + try expect(cold.from == 0 and owner.pending == null and adopt.items.len == 0); + var kept: [4]u32 = undefined; + @memcpy(&kept, (try p.openWith(1, a.len, 100, adopt.items))[0..4]); + x.keep(1, &a, 2048); + try expectEqual(@as(usize, 1), states.live); + p.close(1); // its pages past the mark free again, those below it the tree's + for (kept) |id| try expectEqual(@as(u32, 1), p.pool.holders(id)); + // the same 2,049 tokens, another tail: the pass resumes at 2,048 on the tree's pages, its state pending + var b = a; + for (b[2049..]) |*t| t.* += 10_000; + const warm = try x.begin(&b, 0, &.{}, &owner, &adopt); + defer gpa.free(warm.marks); + try expectEqual(@as(u32, 2048), warm.from); + try expectEqual(@as(u32, 2048), owner.pending.?.at); + try std.testing.expectEqualSlices(u32, &kept, adopt.items); + try std.testing.expectEqualSlices(u32, &kept, (try p.openWith(2, b.len, 100, adopt.items))[0..4]); + for (kept) |id| try expectEqual(@as(u32, 2), p.pool.holders(id)); // the tree and the stream + adopt.clearRetainingCapacity(); + p.close(2); + // another token at the mark: the head's row 2,047 read it, so the last page is another page and nothing resumes + var c = a; + for (c[2048..]) |*t| t.* += 10_000; + owner = .{}; + const miss = try x.begin(&c, 0, &.{}, &owner, &adopt); + defer gpa.free(miss.marks); + try expect(miss.from == 0 and owner.pending == null and adopt.items.len == 0); + // a stream that needs the tree's pages: its promise evicts the state, and the pages are the stream's + _ = try p.openWith(0, c.len, 100, &.{}); // 7 pages: the 4 free and 3 of the tree's + try expectEqual(@as(usize, 0), states.live); + p.close(0); + const gone = try x.begin(&a, 0, &.{}, &owner, &adopt); + defer gpa.free(gone.marks); + try expect(gone.from == 0 and owner.pending == null and adopt.items.len == 0); + try expect(p.pool.idle() and p.pool.census().free == 8); +} + fn made(gpa: std.mem.Allocator) !void { var f: Fake = undefined; try f.init(gpa); diff --git a/zig/src/families/glm/cuda_prefix.zig b/zig/src/families/glm/cuda_prefix.zig new file mode 100644 index 0000000000..33bd4591f9 --- /dev/null +++ b/zig/src/families/glm/cuda_prefix.zig @@ -0,0 +1,117 @@ +//! TF_GLM_KV_SHARE: a prompt cache over the pool's KV pages (prompt_radix.zig); a resume adopts them and the rest. +const std = @import("std"); +pub const radix = @import("engine_api").prompt_radix; +const kv = @import("core").kv_pool; +const engine = @import("cuda_engine.zig"); +const pages_mod = @import("cuda_pages.zig"); +const page_rows = @import("cuda_state.zig").page_rows; + +const Engine = engine.Engine; + +/// TF_GLM_KV_SHARE=n: keep prompts as pool pages and at most n states past them (0 or unset: off). +pub fn sharing() !usize { + const v = std.c.getenv("TF_GLM_KV_SHARE") orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch error.BadGlmKvShare; +} + +/// The store's rules: 512-row pages, states on the 2,048-row prompt grid, the MTP head's row at-1 keyed by token at. +pub fn rules(grid: u32) radix.Rules { + return .{ .page = page_rows, .planned = true, .grid = grid, .lookahead = 1, .min_prompt = 0, .tail = true }; +} + +/// A resume's state: restored by the prompt pass's start on every rank (engine.Resume), once its pages are mapped. +pub const Owner = struct { pending: ?*engine.Kept = null }; + +/// A pass's marks (the engine's Marks.ctx): the mark it stands at, or the last it completed when cancelled, keeps. +pub const Keeps = struct { + x: *Prefix, + slot: u32, + prompt: []const u32, + + pub fn at(ctx: *anyopaque, mark: u32) void { + const k: *Keeps = @ptrCast(@alignCast(ctx)); + k.x.keep(k.slot, k.prompt, mark); + } +}; + +pub const Prefix = struct { + gpa: std.mem.Allocator, + pages: *pages_mod.Pages, + store: radix.Store, + + /// The store over `p`'s pool, keeping `e`'s states past the pages it holds (`e.share_pages`). + pub fn init(gpa: std.mem.Allocator, e: *Engine, p: *pages_mod.Pages, grid: u32, snaps: usize) !*Prefix { + return initWith(gpa, p, grid, snaps, .{ .ptr = e, .vtable = &family }); + } + + pub fn initWith(gpa: std.mem.Allocator, p: *pages_mod.Pages, grid: u32, snaps: usize, states: radix.Snapshots) !*Prefix { + const x = try gpa.create(Prefix); + errdefer gpa.destroy(x); + const pool_pages = radix.pagesOf(kv.Pool, &p.pool, p.layout.pageBytes()); + x.* = .{ .gpa = gpa, .pages = p, .store = try radix.Store.init(gpa, states, pool_pages, rules(grid)) }; + x.store.limit(.{ .pages = p.pool.total(), .snaps = snaps }); + p.pool.reclaim = .{ .ctx = x, .call = reclaim }; // a promise evicts the store's idle pages before it fails + return x; + } + + pub fn deinit(x: *Prefix) void { + x.pages.pool.reclaim = null; + x.store.deinit(); // drops every state and releases the tree's pages, before the pool goes + x.gpa.destroy(x); + } + + fn reclaim(ctx: *anyopaque, want: usize) bool { + const x: *Prefix = @ptrCast(@alignCast(ctx)); + return x.store.reclaim(want); + } + + /// Before a prompt pass: the deepest kept state's pages in `adopt` (held for the stream), the state in `owner`. + pub fn begin(x: *Prefix, prompt: []const u32, history_len: u32, shared: []const u32, owner: *Owner, adopt: *std.ArrayList(u32)) !radix.Plan { + return x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + } + + /// Slot `slot`'s pass at mark `at`: the tree takes the pages below it it lacks, and the state (no table swaps). + pub fn keep(x: *Prefix, slot: u32, prompt: []const u32, at: u32) void { + if (!x.store.on() or !x.pages.paged(slot)) return; + const n = at / page_rows; + const ids = x.pages.leases[slot].ids.items; + if (n > ids.len) return; + const path = x.gpa.alloc(u32, n) catch return; + defer x.gpa.free(path); + var owner: Owner = .{}; + _ = x.store.keep(prompt, at, &owner, ids[0..n], path); + } + + fn of(ptr: *anyopaque) *Engine { + return @ptrCast(@alignCast(ptr)); + } + + fn snapBytes(ptr: *anyopaque, at: u32) u64 { + return of(ptr).stateBytes(at); + } + + /// The bound slot's state at `at` (a mark at a chunk start), its pages below `at` left to the tree. + fn save(ptr: *anyopaque, _: ?*anyopaque, at: u32) anyerror!radix.Saved { + const e = of(ptr); + if (e.pos() != at) return error.NotAtMark; + const k = try e.gpa.create(engine.Kept); + errdefer e.gpa.destroy(k); + k.* = .{ .id = try e.keepState(), .at = at }; + return @ptrCast(k); + } + + /// The store's restore: the state waits in `owner` for the pass's start (Resume), its pages for the stream's lease. + pub fn restore(_: *anyopaque, owner: ?*anyopaque, saved: radix.Saved) anyerror!void { + const o: *Owner = @ptrCast(@alignCast(owner orelse return error.NoGlmOwner)); + o.pending = @ptrCast(@alignCast(saved)); + } + + fn drop(ptr: *anyopaque, saved: radix.Saved) void { + const e = of(ptr); + const k: *engine.Kept = @ptrCast(@alignCast(saved)); + e.dropState(k.id); + e.gpa.destroy(k); + } + + const family: radix.Snapshots.VTable = .{ .bytes = snapBytes, .save = save, .restore = restore, .drop = drop }; +}; diff --git a/zig/src/families/glm/cuda_snapshot_layout.zig b/zig/src/families/glm/cuda_snapshot_layout.zig index b01d943e1f..5454cd38df 100644 --- a/zig/src/families/glm/cuda_snapshot_layout.zig +++ b/zig/src/families/glm/cuda_snapshot_layout.zig @@ -6,30 +6,33 @@ const Config = @import("config.zig").Config; const cache = @import("cuda_dsa_block.zig"); const sparse = @import("cuda_sparse.zig"); -/// `len` bytes from `live`, or on a paged sequence's pages `span` bytes a page, the pages `page` bytes apart. -pub const Part = struct { live: u64, len: usize, span: usize = 0, page: u64 = 0 }; +/// `len` bytes from `live`, or on a paged sequence's pages `span` bytes a page `page` bytes apart, from byte `from`. +pub const Part = struct { live: u64, len: usize, span: usize = 0, page: u64 = 0, from: usize = 0 }; pub const Plan = struct { parts: [5 * (@import("config.zig").max_layers + 1) + 2]Part = undefined, count: usize = 0, len: usize = 0, ids: []const u32 = &.{}, // a paged sequence's pages in row order sizing: bool = false, // `size`: the rows counted, not read, so no page need be mapped + shared: usize = 0, // a paged sequence's first pages a page-sharing store holds: their rows stay on them - pub fn init(st: *const State, c: *const Config, at: usize, head: ?usize) !Plan { - return make(st, c, at, head, false); + /// The state at `at`; past `shared` pages only, on a paged sequence whose first pages the store holds. + pub fn init(st: *const State, c: *const Config, at: usize, head: ?usize, shared: usize) !Plan { + return make(st, c, at, head, shared, false); } /// A prefix's bytes at `at`, rows to `capacity` counted held or not (a paged sequence holds its request's alone). - pub fn size(st: *const State, c: *const Config, capacity: usize, at: usize, head: ?usize) !usize { + pub fn size(st: *const State, c: *const Config, capacity: usize, at: usize, head: ?usize, shared: usize) !usize { var whole = st.*; whole.capacity = capacity; - return (try make(&whole, c, at, head, true)).len; + return (try make(&whole, c, at, head, shared, true)).len; } - fn make(st: *const State, c: *const Config, at: usize, head: ?usize, sizing: bool) !Plan { + fn make(st: *const State, c: *const Config, at: usize, head: ?usize, shared: usize, sizing: bool) !Plan { if (at > st.capacity or st.caches.len > @import("config.zig").max_layers) return error.SnapshotShape; if (st.rings.len != 2 * (st.caches.len + 1)) return error.SnapshotShape; - var p: Plan = .{ .ids = st.pages, .sizing = sizing }; + if (shared > 0 and (st.table == 0 or shared * page_rows > at)) return error.SnapshotShape; // pages below `at` only + var p: Plan = .{ .ids = st.pages, .sizing = sizing, .shared = shared }; // a paged sequence's caches are its pages: its fields up to them end at its first ring const first = if (st.table != 0) st.rings[0] else if (st.caches.len > 0) st.caches[0].lc else st.mtp_caches.lc; if (first < st.arena.buf.ptr) return error.SnapshotShape; @@ -58,21 +61,23 @@ pub const Plan = struct { try p.addCache(st, x, x.pk, pooled, page_rows / sparse.pool * @as(usize, c.i_dim) * 2); } - /// A cache's rows: in the arena, or on the sequence's pages, `span` bytes a page. + /// A cache's rows: in the arena, or on the sequence's pages, `span` bytes a page (past the store's `shared`). fn addCache(p: *Plan, st: *const State, x: cache.Caches, live: u64, len: usize, span: usize) !void { if (x.pt == 0) return p.add(st, live, len); - if (len == 0) return; + const from = @min(len, p.shared * span); + if (len == from) return; if ((!p.sizing and (len - 1) / span >= p.ids.len) or p.count == p.parts.len) return error.SnapshotShape; - p.len = try std.math.add(usize, p.len, len); - p.parts[p.count] = .{ .live = live, .len = len, .span = span, .page = x.ps }; + p.len = try std.math.add(usize, p.len, len - from); + p.parts[p.count] = .{ .live = live, .len = len - from, .span = span, .page = x.ps, .from = from }; p.count += 1; } /// Where a part's bytes from `done` on sit, as many as lie in a row there. pub fn piece(p: *const Plan, part: Part, done: usize) Part { if (part.span == 0) return .{ .live = part.live + done, .len = part.len - done }; - const off = done % part.span; - return .{ .live = part.live + @as(u64, p.ids[done / part.span]) * part.page + off, .len = @min(part.span - off, part.len - done) }; + const at = part.from + done; + const off = at % part.span; + return .{ .live = part.live + @as(u64, p.ids[at / part.span]) * part.page + off, .len = @min(part.span - off, part.len - done) }; } fn add(p: *Plan, st: *const State, live: u64, len: usize) !void { diff --git a/zig/src/families/glm/cuda_snapshots.zig b/zig/src/families/glm/cuda_snapshots.zig index d3a9f6a5a6..9b1adeb3e7 100644 --- a/zig/src/families/glm/cuda_snapshots.zig +++ b/zig/src/families/glm/cuda_snapshots.zig @@ -27,7 +27,7 @@ pub fn agree(e: anytype, budget: u64) !void { pub fn bytes(e: anytype, at: usize) u64 { const head: ?usize = if (e.mtp) @min(e.capacity, std.math.add(usize, at, @import("cuda_state.zig").window_rows) catch return std.math.maxInt(u64)) else null; // a paged sequence holds its request's rows: its prefix is counted to the capacity, in a window's bytes - const n = layout.Plan.size(&e.cur.st, &e.c, e.capacity, at, head) catch return std.math.maxInt(u64); + const n = layout.Plan.size(&e.cur.st, &e.c, e.capacity, at, head, heldPages(e, at)) catch return std.math.maxInt(u64); return pages.rounded(n) catch std.math.maxInt(u64); } @@ -45,7 +45,7 @@ pub fn keepAs(e: anytype, id: u64) !void { if (e.cur.st.pending != 0 or e.cur.st.mtp_len > e.cur.st.pos + @import("cuda_state.zig").window_rows) return error.NotAtMark; const pool = if (e.snap_pages) |*p| p else return error.NoSnapshotCache; - const plan = try layout.Plan.init(&e.cur.st, &e.c, e.cur.st.pos, if (e.mtp) e.cur.st.mtp_len else null); + const plan = try layout.Plan.init(&e.cur.st, &e.c, e.cur.st.pos, if (e.mtp) e.cur.st.mtp_len else null, heldPages(e, e.cur.st.pos)); var lease = try pool.take(plan.len); errdefer { cuda.completion.drain(e.d, e.stream, "a failed GLM snapshot copy did not drain"); @@ -55,13 +55,13 @@ pub fn keepAs(e: anytype, id: u64) !void { errdefer e.gpa.free(cur); try pool.move(e.stream, &lease, &plan, true); try e.stream.synchronize(); - e.snaps.putAssumeCapacityNoClobber(id, .{ .lease = lease, .len = plan.len, .pos = e.cur.st.pos, .mtp_len = e.cur.st.mtp_len, .cur = cur, .bases = .{ e.cur.st.base, e.cur.st.mtp_base } }); + e.snaps.putAssumeCapacityNoClobber(id, .{ .lease = lease, .len = plan.len, .pos = e.cur.st.pos, .mtp_len = e.cur.st.mtp_len, .cur = cur, .bases = .{ e.cur.st.base, e.cur.st.mtp_base }, .shared = plan.shared }); } pub fn restore(e: anytype, id: u64) !void { const s = e.snaps.getPtr(id) orelse return error.NoKeptState; const pool = if (e.snap_pages) |*p| p else return error.NoSnapshotCache; - const plan = try layout.Plan.init(&e.cur.st, &e.c, s.pos, if (e.mtp) s.mtp_len else null); + const plan = try layout.Plan.init(&e.cur.st, &e.c, s.pos, if (e.mtp) s.mtp_len else null, s.shared); if (plan.len != s.len) return error.SnapshotSize; errdefer cuda.completion.drain(e.d, e.stream, "a failed GLM snapshot restore did not drain"); try pool.move(e.stream, &s.lease, &plan, false); @@ -70,6 +70,11 @@ pub fn restore(e: anytype, id: u64) !void { try e.stream.synchronize(); } +/// The pages below `at` a page-sharing prompt cache holds (TF_GLM_KV_SHARE): a kept state's copy leaves them out. +fn heldPages(e: anytype, at: usize) usize { + return if (e.share_pages) at / @import("cuda_state.zig").page_rows else 0; +} + pub fn holds(e: anytype, id: u64) bool { return e.snaps.contains(id); } diff --git a/zig/src/families/glm/dflash.zig b/zig/src/families/glm/dflash.zig index 6b67d76489..355a794d98 100644 --- a/zig/src/families/glm/dflash.zig +++ b/zig/src/families/glm/dflash.zig @@ -34,6 +34,12 @@ pub const TAP_ROW: usize = taps_mod.layers.len * D * 2; // bytes of one row's ta /// Device bytes a stream's context ring takes (K and V of every layer). pub const ring_bytes: u64 = LAYERS * 2 * @as(u64, @intCast(WIN)) * KVW * 2; +/// TF_GLM_DRAFT_CONFIDENCE: where a DFlash2 chain's chance product stops it (glm_cuda's 0.7). +pub fn confidence() f64 { + const v = std.c.getenv("TF_GLM_DRAFT_CONFIDENCE") orelse return 0.7; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch 0.7; +} + /// Device bytes of the loader's raw scratch: a 4-bit linear's words, scales and biases, for the largest at `world`. pub fn rawBytes(bits: u8, world: usize) usize { const ff: usize = if (world == 2) FF / 2 else FF; // the MLP width a rank holds, as dflash_load splits it From 43a26d0b77fe0c6da985683733a917360c804919 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:17:40 +0900 Subject: [PATCH 2/4] glm cuda: TF_GLM_KV_SHARE at TP: rank 0's begin step names the cache's pages, its snapshots step their absence The tree and the pages stay rank 0's. The begin step carries how many of the table's first pages are the cache's, so another rank's map leaves their rows out of TF_KV_POISON's fill and refuses a count past the table. The snapshots step carries whether states leave the pool's pages out, so every rank's snapshot has rank 0's shape. Rank 0 sends its drops before the begin step that maps a freed page, as beginOn already flushes first. The one-rank refusal goes. --- zig/glm_pages_host.zig | 8 ++--- zig/src/families/glm/cuda_follow.zig | 23 ++++++++------- zig/src/families/glm/cuda_lanes.zig | 2 +- zig/src/families/glm/cuda_native.zig | 5 +--- zig/src/families/glm/cuda_pages.zig | 8 ++--- zig/src/families/glm/cuda_pages_test.zig | 37 ++++++++++++++++++++++-- zig/tests/cuda/glm_paged_tests.zig | 2 +- 7 files changed, 58 insertions(+), 27 deletions(-) diff --git a/zig/glm_pages_host.zig b/zig/glm_pages_host.zig index 2563fac518..5ab24dac46 100644 --- a/zig/glm_pages_host.zig +++ b/zig/glm_pages_host.zig @@ -205,15 +205,15 @@ test "GLM pages: a follower refuses a duplicate owner and rolls back a failed ta defer f.deinit(); f.p.leader = false; F.fail_put = true; - try t.expectError(error.NotReady, f.p.map(0, &.{ 3, 2 })); + try t.expectError(error.NotReady, f.p.map(0, &.{ 3, 2 }, 0)); try t.expectEqual(@as(usize, 0), f.p.leases[0].ids.items.len); try t.expect(!f.p.used[0]); F.fail_put = false; - try f.p.map(0, &.{ 3, 2 }); - try t.expectError(error.BadStep, f.p.map(0, &.{ 1, 0 })); + try f.p.map(0, &.{ 3, 2 }, 0); + try t.expectError(error.BadStep, f.p.map(0, &.{ 1, 0 }, 0)); try t.expectEqualSlices(u32, &.{ 3, 2 }, f.e.seqs.items[0].?.st.pages); f.p.close(0); - try f.p.map(0, &.{ 1, 0 }); + try f.p.map(0, &.{ 1, 0 }, 0); f.p.close(0); try t.expect(f.p.pool.idle()); } diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index cdf90bbfb5..c6cce49f23 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -19,8 +19,8 @@ pub const Header = extern struct { kind: u32, n: u32 = 0, // ids that follow depth: u32 = 0, // draft: drafts to chain; chunk: where it starts; shared: its parts - flags: u32 = 0, // begin: sampled; draft: from the prompt (1), chances (2); prefill, chunk: prompt_logprobs (1); chunk: indexer rows split (2); verify, shared: whole rows - a: u64 = 0, // keep: rows; prefill, chunk: the kept state it resumes (0: none); redraw: the prompt's length + flags: u32 = 0, // begin: sampled; draft: from the prompt (1), chances (2); prefill, chunk: prompt_logprobs (1); chunk: indexer rows split (2); verify, shared: whole rows; snapshots: pages shared + a: u64 = 0, // begin: the prompt cache's pages first in its table; keep: rows; prefill, chunk: the kept state it resumes (0: none); redraw: the prompt's length seed: u64 = 0, // begin: the sampling seed; prefill, chunk: the kept state's position temperature: f64 = 0, top_p: f64 = 0, @@ -172,14 +172,14 @@ pub const Leader = struct { } pub fn begin(x: Leader, s: ?sampler.Sampling) !void { - return x.beginOn(s, &.{}); + return x.beginOn(s, &.{}, 0); } - /// `begin` on a paged sequence: the page ids it maps (rank 0's lease) go with the step, to every rank's table. - pub fn beginOn(x: Leader, s: ?sampler.Sampling, pages: []const u32) !void { + /// `begin` on a paged sequence: its page ids (rank 0's lease, the cache's `kept` first) go to every rank's table. + pub fn beginOn(x: Leader, s: ?sampler.Sampling, pages: []const u32, kept: usize) !void { try x.flush(); const c = sampler.check(s); - var h: Header = .{ .kind = @backingInt(Kind.begin) }; + var h: Header = .{ .kind = @backingInt(Kind.begin), .a = kept }; if (c) |v| { h.flags = 1; h.seed = v.seed; @@ -192,9 +192,9 @@ pub const Leader = struct { try x.e.begin(c); } - /// Every rank's prompt cache at rank 0's budget (snapshots.agree): the others make theirs on this step. + /// Every rank's prompt cache at rank 0's budget (snapshots.agree), its states past the pool's pages as rank 0's. pub fn snapshots(x: Leader, budget: u64) !void { - try x.send(.{ .kind = @backingInt(Kind.snapshots), .a = budget }, &.{}); + try x.send(.{ .kind = @backingInt(Kind.snapshots), .a = budget, .flags = @intFromBool(x.e.share_pages) }, &.{}); try @import("cuda_snapshots.zig").agree(x.e, budget); } @@ -331,7 +331,7 @@ pub fn follow(e: *Engine, link: *Link) !void { e.structured = structuredStep(h); // rank 0 alone masks them switch (kind) { .begin => { - if (step.ids.len > 0) try (e.pages orelse return error.BadStep).map(h.slot, step.ids); + if (step.ids.len > 0) try (e.pages orelse return error.BadStep).map(h.slot, step.ids, std.math.cast(usize, h.a) orelse return error.BadStep); try e.beginFollowing(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null); }, .pool => { @@ -399,7 +399,10 @@ pub fn follow(e: *Engine, link: *Link) !void { try v.take(h.slot, step.ids); try link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); }, - .snapshots => try @import("cuda_snapshots.zig").agree(e, h.a), + .snapshots => { + e.share_pages = h.flags & 1 != 0; // TF_GLM_KV_SHARE on rank 0: a kept state leaves the prompt cache's pages out + try @import("cuda_snapshots.zig").agree(e, h.a); + }, .drop => for (step.ids) |id| e.dropState(id), .release => e.release(h.slot), .idle => cuda.tp_watch.sleep(e.io), diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index f676ddcc57..be2b5c7d68 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -295,7 +295,7 @@ pub const Lanes = struct { _ = self.lanes.remove(s); } if (self.x.e.images) |v| v.clear(); // no earlier request's rows splice into this one - try self.x.beginOn(s.sampling, pages); + try self.x.beginOn(s.sampling, pages, adopt.items.len); s.cached = 0; s.prompt_rows.clearRetainingCapacity(); const l = self.lanes.getPtr(s).?; diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 28feff0a4f..02c794f42e 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -219,8 +219,6 @@ fn openRank(a: StartArgs, set: ranks.Settings) !Loaded { const roce_bytes = a.roce_bytes; // TF_GLM_KV_SHARE keeps prompts as the pool's pages: refused before the weights load, every rank told (startup) if (a.share > 0 and !set.pages) return error.GlmShareNeedsPages; - // cc-defer: one rank; a follower keeps its fixed parts and maps adopted pages when the TP stage lands (V9F_RADIX) - if (a.share > 0 and p.world > 1) return error.GlmShareOneRank; cuda.memory_receipt.emit(@intCast(p.rank), .load_begin, true); // the four phases tools/memory_receipts.py reads const own = try gpa.create(Owned); errdefer gpa.destroy(own); @@ -452,8 +450,7 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { error.RankMismatch => "the ranks' settings differ (each rank's log names the one): start them alike", error.Fp8OnPages => "--kv-cache fp8 keeps a window per stream: unset TF_KV_PAGES, or set it to auto or off", error.BadGlmKvShare => "TF_GLM_KV_SHARE: give how many states the prompt cache over the KV pages keeps, or 0 or unset for off", - error.GlmShareNeedsPages => "TF_GLM_KV_SHARE keeps prompts as KV pool pages: on one rank set TF_KV_PAGES=on as well", - error.GlmShareOneRank => "TF_GLM_KV_SHARE runs on one rank for now: serve TF_GLM_WORLD=1, or unset it", + error.GlmShareNeedsPages => "TF_GLM_KV_SHARE keeps prompts as KV pool pages: set TF_KV_PAGES=on as well", error.TooManyRanks => std.fmt.comptimePrint( "GLM-5.3-Flash on CUDA runs on at most {d} ranks: set TF_GLM_WORLD to {d} or fewer", .{ ranks.max_world, ranks.max_world }, diff --git a/zig/src/families/glm/cuda_pages.zig b/zig/src/families/glm/cuda_pages.zig index a6ac11642d..4bb2f5a5f5 100644 --- a/zig/src/families/glm/cuda_pages.zig +++ b/zig/src/families/glm/cuda_pages.zig @@ -162,9 +162,9 @@ pub const Pages = struct { return l.ids.items; } - /// Another rank: paged slot `slot` maps rank 0's page ids (sent with the stream's begin step). - pub fn map(p: *Pages, slot: u32, ids: []const u32) !void { - if (!p.paged(slot) or ids.len == 0 or ids.len > tableLen(p.e.capacity)) return error.BadStep; + /// Another rank: paged slot `slot` maps rank 0's page ids (its begin step's), the prompt cache's `kept` first. + pub fn map(p: *Pages, slot: u32, ids: []const u32, kept: usize) !void { + if (!p.paged(slot) or ids.len == 0 or ids.len > tableLen(p.e.capacity) or kept > ids.len) return error.BadStep; if (p.used[slot]) return error.BadStep; for (ids) |id| if (id >= p.pool.total()) return error.BadStep; const l = &p.leases[slot]; @@ -173,7 +173,7 @@ pub const Pages = struct { errdefer l.ids.clearRetainingCapacity(); var up: Upload = .{ .p = p, .slot = slot }; try Upload.write(&up, l.ids.items, 0); - try p.mapped(slot, l.ids.items, 0); + try p.mapped(slot, l.ids.items, kept); } /// The slot's sequence reaches its pages' rows (past the `kept` first, NaN first under TF_KV_POISON). diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index f49f22afda..ee0199839a 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -174,16 +174,47 @@ test "a follower maps rank 0's ids, and refuses ids past its pool or a slot that const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); const p = try f.pool(l, 3, 8, false); defer f.deinit(p); - try p.map(2, &.{ 5, 3, 7 }); + try p.map(2, &.{ 5, 3, 7 }, 0); try std.testing.expectEqualSlices(u32, &.{ 5, 3, 7 }, F.table(f.e.seqs.items[2].?.st.table, 3)); try expectEqual(@as(usize, 1536), f.e.seqs.items[2].?.st.capacity); - try std.testing.expectError(error.BadStep, p.map(1, &.{8})); - try std.testing.expectError(error.BadStep, p.map(3, &.{1})); + try std.testing.expectError(error.BadStep, p.map(1, &.{8}, 0)); + try std.testing.expectError(error.BadStep, p.map(3, &.{1}, 0)); p.close(2); try expectEqual(@as(usize, 0), f.e.seqs.items[2].?.st.capacity); try expect(p.pool.idle()); // a follower's pool lends nothing } +test "TF_GLM_KV_SHARE at TP: two followers map rank 0's resumed table, NaN past the cache's pages, more kept refused" { + var ids: [4]u32 = undefined; + { + var f: Fake = undefined; // rank 0: the prompt cache keeps the first two pages, and a resume holds them again + try f.init(std.testing.allocator); + const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); + const p = try f.pool(l, 3, 8, true); + defer f.deinit(p); + _ = try p.open(1, 1500, 100); + for ([_]u32{ 0, 1, 0, 1 }) |id| p.pool.retain(id); + p.close(1); + @memcpy(&ids, try p.openWith(2, 1500, 100, &.{ 0, 1 })); + p.close(2); + for ([_]u32{ 0, 1 }) |id| p.pool.release(id); + try expect(p.pool.idle()); + } + for (1..3) |_| { // ranks 1 and 2, one after the other on the fake driver + var f: Fake = undefined; + try f.init(std.testing.allocator); + const l = try pages_mod.layout(std.testing.allocator, &f.e, 0); + const p = try f.pool(l, 3, 8, false); + defer f.deinit(p); + p.poison = true; + try p.map(2, &ids, 2); + try std.testing.expectEqualSlices(u32, &ids, F.table(f.e.seqs.items[2].?.st.table, 4)); + try expectEqual(@as(usize, 2), F.fills); // the stream's two fresh pages alone + p.close(2); + try std.testing.expectError(error.BadStep, p.map(2, ids[0..2], 3)); + } +} + test "a follower refuses rank 0's pool when its KV page has other bytes, and makes no pool" { var f: Fake = undefined; try f.init(std.testing.allocator); diff --git a/zig/tests/cuda/glm_paged_tests.zig b/zig/tests/cuda/glm_paged_tests.zig index 5a993a82bb..bfdc8b02ae 100644 --- a/zig/tests/cuda/glm_paged_tests.zig +++ b/zig/tests/cuda/glm_paged_tests.zig @@ -66,7 +66,7 @@ pub fn pagedCheck(gpu: Gpu, checkpoint: []const u8, rank: c_int, world: c_int, i var drafts: [2][glm.engine.max_drafts]u32 = undefined; for ([_]u32{ paged_slot, twin_slot }, 0..) |slot, j| { try x.bind(slot); - if (slot == paged_slot) try x.beginOn(null, held) else try x.begin(null); + if (slot == paged_slot) try x.beginOn(null, held, 0) else try x.begin(null); first[j] = try x.prefill(ids, null, null, null); try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(if (j == 0) mine.ptr else twin.ptr, e.bp.logits, V * 2), "cuMemcpyDtoH"); try x.draft(true, &.{first[j]}, 2, &drafts[j], null); From af80027e1cb5834f838652b6cbb61a69abe83043 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Fri, 30 Oct 2026 23:48:10 +0900 Subject: [PATCH 3/4] glm cuda: TF_GLM_KV_SHARE warns at startup below 4 states a stream Six 8k conversations sent twice at --parallel 4 resumed none at n=4 and all at n=16, with pages to spare: the store limits states as well as pages, and each new conversation evicted the state the next turn wanted. The server now logs a warning below 4 x --parallel and names the value to set, as #1357 does for TF_QWEN_KV_SHARE. --- zig/src/families/glm/cuda.zig | 1 + zig/src/families/glm/cuda_native.zig | 8 ++++++-- zig/src/families/glm/cuda_prefix.zig | 12 ++++++++++++ 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 3992fabf83..c55f7bfd9d 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -79,6 +79,7 @@ test { _ = lanes; _ = native; _ = @import("cuda_qpack.zig"); + _ = @import("cuda_prefix.zig"); _ = @import("cuda_mtp_chance.zig"); _ = @import("cuda_tile.zig"); _ = @import("heat"); // upstream's heat bands, run where the GLM ranks use them diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 02c794f42e..3c662933c9 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -76,16 +76,20 @@ pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { const l = own.layout orelse return error.NoPages; try own.lanes.x.pool(streams + 1, pages, l.pageBytes()); own.e.pages = try pages_mod.Pages.init(own.gpa, own.e, l, streams + 1, pages, true); - if (own.share > 0) try share(own); + if (own.share > 0) try share(own, streams + 1); } /// TF_GLM_KV_SHARE: the pool's prompt cache, its `own.share` snapshots sized at the window's last mark (rings whole). -fn share(own: *Owned) !void { +fn share(own: *Owned, parallel: u32) !void { const e = own.e; e.share_pages = true; // cc-defer: the snapshots take the room set aside for --prompt-cache-gib; from the plan once it plans them try readySnapshots(own, own.share * e.stateBytes(e.capacity - e.capacity % snap_grid)); own.lanes.prefix = try prefix_mod.Prefix.init(e.gpa, e, e.pages.?, snap_grid, own.share); + if (prefix_mod.fewStates(own.share, parallel)) std.log.warn( + "TF_GLM_KV_SHARE={d} keeps under 4 states a stream at --parallel {d}: live streams evict each other's before a resume; set it to {d} or more", + .{ own.share, parallel, 4 * parallel }, + ); } /// Whether a request's caches fit now (the lane host's admission, on its engine thread). diff --git a/zig/src/families/glm/cuda_prefix.zig b/zig/src/families/glm/cuda_prefix.zig index 33bd4591f9..9bc698df81 100644 --- a/zig/src/families/glm/cuda_prefix.zig +++ b/zig/src/families/glm/cuda_prefix.zig @@ -14,6 +14,18 @@ pub fn sharing() !usize { return std.fmt.parseInt(usize, std.mem.span(v), 10) catch error.BadGlmKvShare; } +/// Below 4 states a stream, live streams evict each other's before a resume reaches them. +pub fn fewStates(states: usize, parallel: usize) bool { + return states > 0 and states < 4 * parallel; +} + +test "TF_GLM_KV_SHARE below 4 states a stream is too few; off is not" { + try std.testing.expect(fewStates(4, 4)); + try std.testing.expect(fewStates(15, 4)); + try std.testing.expect(!fewStates(16, 4)); + try std.testing.expect(!fewStates(0, 4)); +} + /// The store's rules: 512-row pages, states on the 2,048-row prompt grid, the MTP head's row at-1 keyed by token at. pub fn rules(grid: u32) radix.Rules { return .{ .page = page_rows, .planned = true, .grid = grid, .lookahead = 1, .min_prompt = 0, .tail = true }; From 682cbd29ac486131223cabf839fa434c72b8d095 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 31 Oct 2026 17:05:18 +0000 Subject: [PATCH 4/4] changelog: #1732 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 41d8319875..7c2addd590 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 +- GLM-5.3-Flash on CUDA: `TF_GLM_KV_SHARE=n` with `TF_KV_PAGES=on` keeps prompts as the KV pool's pages in the shared prompt store (the one `TF_QWEN_KV_SHARE` uses), with at most n states past them, on one GPU and at tensor parallel. Where main's flat store has no room for a long state (GB10s at TP=3 with the 1M window), a 64k message sent again resumes in about 1.7 s instead of prefilling for 36 s; token-id prompts resume at every width (#1510). Unset, nothing changes. + - MOSS-Image-Base's transformer runs on CUDA too, through the same test tool (`tf-moss-image`) and cases as on Metal (#1528, for #1518). It runs on an RTX 5090 (sm_120) and a GB10 (sm_121) and refuses other GPUs; the prompt encoder and image decoder stay in torch, and nothing is served yet. Against the publisher's torch bf16 pipeline, every step's cosine is 0.99967 or better, the same step run twice gives the same bytes, and a 1024x1024 picture is 44.0 dB (RTX 5090) and 45.9 dB (GB10) PSNR from torch's. Its 50 guided steps at 1024x1024 take 24.6 s on an RTX 5090 and 69.6 s on a GB10, 1.26x and 2.03x faster than torch. Its four kernels and full attention at head dim 128, a new instance of the vision towers' kernel, pass the kernel contract with five planted controls red; the towers' existing attention kernels compile to the same SASS. See [the recipe](docs/recipes/moss-image-cuda.md). Thanks to @BobClawblaw. - DeepSeek-V4.x replies carry the Spark recipe's per-response stats (`prompt`, `cached`, `ttft_s`, `finish`, `completion`, `prefill_s`, `decode_s`, `rounds`, `drafted`, `accepted`) in their `tensorfold` block, after its own keys; `TENSORFOLD_WIRE=spark` turns them on for any family and `native` turns them off.