diff --git a/CHANGELOG.md b/CHANGELOG.md index 1240b394ff..adce9aae0e 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_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. diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 353c51fba3..12d83c424e 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/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 diff --git a/zig/glm_snapshot_host.zig b/zig/glm_snapshot_host.zig index 70cc555f3e..2a83b3f43a 100644 --- a/zig/glm_snapshot_host.zig +++ b/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)); diff --git a/zig/glm_snapshot_host_ram.zig b/zig/glm_snapshot_host_ram.zig new file mode 100644 index 0000000000..fdcb3974fc --- /dev/null +++ b/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); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 345a487473..c985ead7a3 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/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(), }; } diff --git a/zig/src/families/glm/cuda_snapshot_pages.zig b/zig/src/families/glm/cuda_snapshot_pages.zig index 43753bb1a9..70ee337c92 100644 --- a/zig/src/families/glm/cuda_snapshot_pages.zig +++ b/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; diff --git a/zig/src/families/glm/cuda_snapshots.zig b/zig/src/families/glm/cuda_snapshots.zig index 08101ab6c1..d3a9f6a5a6 100644 --- a/zig/src/families/glm/cuda_snapshots.zig +++ b/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. diff --git a/zig/test_support/glm_snap_driver.zig b/zig/test_support/glm_snap_driver.zig index 68fc426486..3afb37502b 100644 --- a/zig/test_support/glm_snap_driver.zig +++ b/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;