glm cuda: TF_GLM_HOST_CACHE keeps kept prompt states in pinned host memory

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_HOST_CACHE=1` on every rank keeps its prompt states in pinned host memory mapped into the device, sized by `--prompt-cache-gib`, so on discrete cards where the KV pool leaves the device arena a fraction of a GiB a long conversation's state is kept instead of refused and the next turn resumes rather than prefilling the whole conversation (#1735, #1510). Unset, nothing changes.
+
 - EXL3 mixture-of-experts models on CUDA (GLM-5.3-Flash and the Qwen3.5-MoE EXL3 packs): on sm_120 and sm_121 (RTX PRO and RTX 50 Blackwell, GB10), the grouped experts kernel reads a warp's next tile of weights while it decodes the current one (#1505). The decode and the sums are unchanged, so outputs are the same bits; the kernel contract checks them against the previous loop at one to three tiles a warp and at the served models' widths, with two planted load faults that must fail. On a GB10 the kernel ran 11% faster at MiMo-V2.6-Flash's expert shapes (109 to 97 µs a launch), and GLM-5.3-Flash at two GB10s decoded the same tokens, within 1% of its previous speed. Ampere and Ada GPUs build the same code as before. Thanks to @ajensenwaud.
 
 - CUDA graphs: a capture whose end fails because it was invalidated (`CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`) no longer leaves its stream marked as capturing, which had refused every later capture on that stream and every temporary host upload in the process with `UnownedCaptureSource`. A failed end that leaves the stream capturing, such as one from the wrong thread, still keeps the guard. Thanks to @jayleaton.

docs/recipes/glm-5.3-flash-cuda.md

@@ -43,6 +43,10 @@ left after a 12.2 GiB reserve` and a 5.8 GiB prompt cache
   on, and required and named tool calls use the model's tool markup ([1.1.1](../../CHANGELOG.md#111): #1016 to #1018).
 - `--kv-cache fp8` is an opt-in FP8 latent cache. BF16 remains the default, and its native quality and device
   qualification are pending ([1.1.2](../../CHANGELOG.md#112)).
+- `TF_GLM_HOST_CACHE=1` on every rank keeps kept prompt states in pinned host memory instead of on the card, so a
+  long conversation's next turn resumes when the card has no room left for its state (#1510). `--prompt-cache-gib`
+  sizes the host cache on each rank (8 GiB unasked, at most half of the available memory), and the KV pool sets no
+  device cache aside. Checked on two RTX PRO 6000 Blackwell cards at TP=2 with the EXL3 4bpw checkpoint.
 
 ## More ranks
 

zig/glm_snapshot_host.zig

@@ -58,6 +58,10 @@ test "GLM snapshots: a short prefix fits 4 GiB even with a 524288-token configur
 
 const rec = @import("test_support/glm_snap_driver.zig").Driver;
 const Fixture = @import("test_support/glm_snap_fixture.zig").Fixture;
+
+test {
+    _ = @import("glm_snapshot_host_ram.zig");
+}
 const pc = @import("engine_api").prompt_cache;
 const adapter = @import("copies");
 
@@ -91,7 +95,7 @@ test "GLM snapshots: grid and committed off-grid read sets survive copies on bot
             try t.expectEqual(@as(u8, 0xE9), rec.bytes(tail, 1)[0]);
             const charge = f.e.snaps.get(id).?.lease.ids.items.len * glm.snapshots.pages.page;
             try t.expect(charge <= f.e.stateBytes(at));
-            try t.expectEqual(f.e.snap_pages.?.arena.?.len, charge + f.e.snap_pages.?.spare());
+            try t.expectEqual(f.e.snap_pages.?.arena.?.len(), charge + f.e.snap_pages.?.spare());
             f.e.dropState(id);
             // In production the following gate or the next Leader step drains these notes.
             f.e.dropped.clearRetainingCapacity();
@@ -169,7 +173,7 @@ test "GLM snapshots: copy faults drain before pages return and a restore can ret
 test "GLM snapshots: fragmented pages join a larger state and cannot exceed the admitted budget" {
     var d = rec.init();
     const p = glm.snapshots.pages;
-    var pool = try p.Pool.init(t.allocator, &d, 5 * p.page + p.page - 1);
+    var pool = try p.Pool.init(t.allocator, &d, 5 * p.page + p.page - 1, false);
     defer pool.deinit();
     var a = try pool.take(1);
     var b = try pool.take(2 * p.page);
@@ -363,13 +367,13 @@ test "GLM snapshots: pool construction failures refund device accounting and fre
     const before = cuda.usage(false).device;
     for (0..3) |index| {
         var failing = t.FailingAllocator.init(t.allocator, .{ .fail_index = index });
-        try t.expectError(error.OutOfMemory, glm.snapshots.pages.Pool.init(failing.allocator(), &d, 2 << 20));
+        try t.expectError(error.OutOfMemory, glm.snapshots.pages.Pool.init(failing.allocator(), &d, 2 << 20, false));
     }
     rec.fail_alloc = true;
-    try t.expectError(error.NotReady, glm.snapshots.pages.Pool.init(t.allocator, &d, 2 << 20));
+    try t.expectError(error.NotReady, glm.snapshots.pages.Pool.init(t.allocator, &d, 2 << 20, false));
     rec.fail_alloc = false;
     try t.expectEqual(before, cuda.usage(false).device);
-    var empty = try glm.snapshots.pages.Pool.init(t.allocator, &d, (2 << 20) - 1);
+    var empty = try glm.snapshots.pages.Pool.init(t.allocator, &d, (2 << 20) - 1, false);
     defer empty.deinit();
     try t.expectEqual(@as(u64, 0), empty.spare());
     try t.expectError(error.NoSnapshotRoom, empty.take(1));

zig/glm_snapshot_host_ram.zig

@@ -0,0 +1,42 @@
+//! GLM kept states in pinned host memory (TF_GLM_HOST_CACHE) through the actual rank owner with a recording driver.
+const std = @import("std");
+const glm = @import("glm");
+const cuda = @import("cuda");
+const t = std.testing;
+const rec = @import("test_support/glm_snap_driver.zig").Driver;
+const Fixture = @import("test_support/glm_snap_fixture.zig").Fixture;
+
+extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int;
+extern "c" fn unsetenv(name: [*:0]const u8) c_int;
+
+test "GLM snapshots: TF_GLM_HOST_CACHE keeps states in pinned host memory and restores them byte for byte" {
+    var d = rec.init();
+    _ = setenv("TF_GLM_HOST_CACHE", "1", 1);
+    defer _ = unsetenv("TF_GLM_HOST_CACHE");
+    const before = cuda.usage(false).device;
+    for (0..2) |rank| {
+        var f: Fixture = undefined;
+        try f.init(&d, rank, 2, true);
+        defer f.deinit();
+        const device_allocs = rec.allocs;
+        const host_allocs = rec.host_allocs;
+        try f.e.prepareSnapshots(2 * f.e.stateBytes(8192));
+        try t.expectEqual(device_allocs, rec.allocs); // the arena is host memory: no device allocation
+        try t.expectEqual(host_allocs + 1, rec.host_allocs);
+        try t.expect(f.e.snap_pages.?.arena.? == .host);
+        for ([_]usize{ 2048, 4096, 8192 }) |at| {
+            try f.mark(at, at, @intCast(at % 251));
+            const saved = try t.allocator.dupe(u8, rec.bytes(f.seq.st.arena.buf.ptr, f.seq.st.arena.buf.len));
+            defer t.allocator.free(saved);
+            const id = try f.e.keepState();
+            @memset(rec.bytes(f.seq.st.arena.buf.ptr, f.seq.st.arena.buf.len), 0x5A);
+            try f.seq.st.restoreHost(f.e.stream, 17, 18, &.{0}, .{ 0, 0 });
+            try f.e.restoreState(id);
+            try f.expectPrefix(saved, at, at);
+            try t.expectEqual(at, f.seq.st.pos);
+            f.e.dropState(id);
+            f.e.dropped.clearRetainingCapacity();
+        }
+    }
+    try t.expectEqual(before, cuda.usage(false).device);
+}

zig/src/families/glm/cuda_native.zig

@@ -65,6 +65,7 @@ pub const Loaded = struct {
     vision: ?*images_mod.Images = null, // rank 0's tower for the server's image offer
     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
 };
 
 /// The pool behind every stream: the engine's own sequence and `streams` more on `pages` KV pages, on every rank.
@@ -330,6 +331,7 @@ fn finishRank(a: OpenArgs) !Loaded {
         .vision = own.e.images,
         .layout = own.layout,
         .engine_paged = set.pages,
+        .host_snaps = snapshots.pages.onHost(),
     };
 }
 

zig/src/families/glm/cuda_snapshot_pages.zig

@@ -11,22 +11,62 @@ pub fn rounded(bytes: usize) !usize {
     return (try std.math.add(usize, bytes, page - 1)) & ~(page - 1);
 }
 
+/// TF_GLM_HOST_CACHE=1: kept states in pinned host memory, not on the device; every rank must be given the same value.
+pub fn onHost() bool {
+    const v = std.c.getenv("TF_GLM_HOST_CACHE") orelse return false;
+    return std.mem.eql(u8, std.mem.span(v), "1");
+}
+
+/// The pages' memory: a device buffer, or mapped pinned host memory that the same device-to-device copies reach.
+const Arena = union(enum) {
+    device: cuda.DeviceBuffer,
+    host: struct { buf: cuda.HostBuffer, ptr: u64 },
+
+    fn ptr(a: Arena) u64 {
+        return switch (a) {
+            .device => |b| b.ptr,
+            .host => |h| h.ptr,
+        };
+    }
+
+    pub fn len(a: Arena) usize {
+        return switch (a) {
+            .device => |b| b.len,
+            .host => |h| h.buf.bytes.len,
+        };
+    }
+
+    fn driver(a: Arena) *const cuda.Driver {
+        return switch (a) {
+            .device => |b| b.d,
+            .host => |h| h.buf.d,
+        };
+    }
+};
+
 pub const Pool = struct {
     pages: kv.Pool,
-    arena: ?cuda.DeviceBuffer,
+    arena: ?Arena,
 
-    pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, budget: u64) !Pool {
+    pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, budget: u64, host: bool) !Pool {
         const count = std.math.cast(usize, budget / page) orelse return error.SnapshotBudget;
         if (count > std.math.maxInt(u32)) return error.SnapshotBudget;
         var pages = try kv.Pool.init(gpa, count, page);
         errdefer pages.deinit();
-        const arena = if (count > 0) try cuda.DeviceBuffer.alloc(d, count * page) else null;
+        const arena: ?Arena = if (count == 0) null else if (host) blk: {
+            var buf = try cuda.HostBuffer.allocMapped(d, count * page);
+            errdefer buf.free();
+            break :blk .{ .host = .{ .buf = buf, .ptr = try buf.device() } };
+        } else .{ .device = try cuda.DeviceBuffer.alloc(d, count * page) };
         return .{ .pages = pages, .arena = arena };
     }
 
     pub fn deinit(p: *Pool) void {
         std.debug.assert(p.pages.available() == p.pages.total());
-        if (p.arena) |*a| a.free();
+        if (p.arena) |*a| switch (a.*) {
+            .device => |*b| b.free(),
+            .host => |*h| h.buf.free(),
+        };
         p.pages.deinit();
     }
 
@@ -51,15 +91,17 @@ pub const Pool = struct {
     pub fn move(p: *const Pool, stream: cuda.Stream, lease: *const Lease, plan: *const Plan, into: bool) !void {
         if (plan.len > lease.ids.items.len * page) return error.SnapshotSize;
         const arena = p.arena orelse return error.NoSnapshotRoom;
+        const base = arena.ptr();
+        const d = arena.driver();
         var off: usize = 0;
         for (plan.parts[0..plan.count]) |part| {
             var done: usize = 0;
             while (done < part.len) {
                 const at = off + done;
-                const target = arena.ptr + @as(u64, lease.ids.items[at / page]) * page + at % page;
+                const target = base + @as(u64, lease.ids.items[at / page]) * page + at % page;
                 const live = plan.piece(part, done);
                 const n = @min(live.len, page - at % page);
-                try arena.d.check(arena.d.api.cuMemcpyDtoDAsync_v2(if (into) target else live.live, if (into) live.live else target, n, stream.handle), "cuMemcpyDtoDAsync");
+                try d.check(d.api.cuMemcpyDtoDAsync_v2(if (into) target else live.live, if (into) live.live else target, n, stream.handle), "cuMemcpyDtoDAsync");
                 done += n;
             }
             off += part.len;

zig/src/families/glm/cuda_snapshots.zig

@@ -12,7 +12,7 @@ pub fn prepare(e: anytype, budget: u64) !void {
     const count = std.math.cast(u32, budget / pages.page) orelse return error.SnapshotBudget;
     try e.snaps.ensureTotalCapacity(e.gpa, count);
     try e.dropped.ensureTotalCapacity(e.gpa, @as(usize, count) + 1);
-    e.snap_pages = try pages.Pool.init(e.gpa, e.d, budget);
+    e.snap_pages = try pages.Pool.init(e.gpa, e.d, budget, pages.onHost());
 }
 
 /// A failed allocation still joins the startup agreement, so a peer never waits for a missing participant.

zig/test_support/glm_snap_driver.zig

@@ -11,6 +11,7 @@ pub const Driver = struct {
 
     pub fn init() cuda.Driver {
         allocs = 0;
+        host_allocs = 0;
         copies = 0;
         syncs = 0;
         fail_copy = null;
@@ -20,6 +21,9 @@ pub const Driver = struct {
         d.api.cuMemFree_v2 = free;
         d.api.cuMemsetD8_v2 = zero;
         d.api.cuMemcpyDtoDAsync_v2 = copy;
+        d.api.cuMemHostAlloc = hostAlloc;
+        d.api.cuMemHostGetDevicePointer_v2 = hostDevice;
+        d.api.cuMemFreeHost = hostFree;
         d.api.cuMemsetD32Async = set;
         d.api.cuStreamSynchronize = sync;
         d.api.cuCtxSynchronize = contextSync;
@@ -51,6 +55,24 @@ pub const Driver = struct {
         };
         return 1;
     }
+    pub var host_allocs: usize = 0;
+    // pinned host memory: the same recorded blocks, its device address the host address (as under UVA)
+    fn hostAlloc(out: *?*anyopaque, n: usize, _: c_uint) callconv(.c) cuda.abi.Result {
+        var p: u64 = 0;
+        const r = alloc(&p, n);
+        if (r != 0) return r;
+        allocs -= 1;
+        host_allocs += 1;
+        out.* = @ptrFromInt(p);
+        return 0;
+    }
+    fn hostDevice(out: *u64, p: ?*anyopaque, _: c_uint) callconv(.c) cuda.abi.Result {
+        out.* = @intFromPtr(p);
+        return 0;
+    }
+    fn hostFree(p: ?*anyopaque) callconv(.c) cuda.abi.Result {
+        return free(@intFromPtr(p));
+    }
     fn zero(ptr: u64, v: u8, n: usize) callconv(.c) cuda.abi.Result {
         @memset(bytes(ptr, n), v);
         return 0;