From 063ce375040f3a1987ea0e57094ed6ee15d858af Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sat, 31 Oct 2026 15:09:26 -0700 Subject: [PATCH 1/7] glm cuda: TF_GLM_KV_HOST keeps the KV pages' latent rows in pinned host memory With TF_GLM_KV_HOST=1 (BF16 KV on pages) the pool's pages live in pinned host memory and only the pooled index keys stay on the device, on pages of their own. Each sparse DSA call gathers the rows it picked, once each, into device slots and attends there with the existing kernels; dense calls gather their visible rows the same way. Decode and verify calls keep each layer's last gathered rows on the device and fetch only the rest. The plan counts a page's device cost as its pooled keys; TF_GLM_KV_HOST_GIB sizes the host pages (one window unasked), and TF_GLM_KV_HOST_SHARED=1 gives the ranks of one host a single copy in /dev/shm. --- docs/recipes/glm-5.3-flash-cuda.md | 5 + zig/kernels/cuda/glm_ops.cu | 98 ++++++++ zig/src/families/glm/cuda_dsa_block.zig | 18 +- zig/src/families/glm/cuda_kv_host.zig | 228 ++++++++++++++++++ zig/src/families/glm/cuda_native.zig | 9 +- zig/src/families/glm/cuda_pages.zig | 51 +++- zig/src/families/glm/cuda_snapshot_layout.zig | 8 +- zig/src/families/glm/cuda_state.zig | 14 +- 8 files changed, 413 insertions(+), 18 deletions(-) create mode 100644 zig/src/families/glm/cuda_kv_host.zig diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 6fb85e59bf..460c8bb143 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -49,6 +49,11 @@ left after a 12.2 GiB reserve` and a 5.8 GiB prompt cache 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. +- `TF_GLM_KV_HOST=1` on every rank (with KV pages, BF16) keeps the KV pages' latent rows in pinned host memory: the + device keeps the pooled index keys, and each sparse call copies only the rows it picked into device slots. The KV + pool then costs the card 768 bytes a token instead of 13,056, so a 1,048,576-token window fits beside the weights. + `TF_GLM_KV_HOST_GIB` sizes the host pages (one full window unasked); `TF_GLM_KV_HOST_SHARED=1` gives the ranks of + one host a single copy in `/dev/shm`. Checked on two RTX PRO 6000 Blackwell cards at TP=2. - `TF_GLM_CHUNK_SOLO=4096` on every rank (opt-in) lets a prompt pass take 4,096 rows while no other stream is decoding. Unset or `2048`, every pass takes 2,048 rows and nothing extra is held. See [Prompt chunk size](#prompt-chunk-size). - `TF_GLM_DENSE_BITS=4` (EXL3 at TP=2, off by default) repacks the BF16 dense layers to 4 bits at load. Decode is diff --git a/zig/kernels/cuda/glm_ops.cu b/zig/kernels/cuda/glm_ops.cu index adbf1abf14..a6a9fd202c 100644 --- a/zig/kernels/cuda/glm_ops.cu +++ b/zig/kernels/cuda/glm_ops.cu @@ -53,3 +53,101 @@ extern "C" __global__ void tf_glm_copy_rows(const unsigned char* __restrict__ sr const long long r = i / row_bytes, c = i % row_bytes; dst[r * dst_pitch + c] = src[r * src_pitch + c]; } + +// TF_GLM_KV_HOST mark: each picked row (first counts[r] of tokens [rows][width]) gets a slot; slot_of -1 if unseen. +extern "C" __global__ void tf_glm_kvh_mark(const int* __restrict__ tokens, const int* __restrict__ counts, int width, + long long total, int* slot_of, int* list, int* n, int cap) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + if (static_cast(i % width) >= counts[i / width]) return; + const int t = tokens[i]; + if (t < 0 || atomicCAS(slot_of + t, -1, -2) != -1) return; + const int s = atomicAdd(n, 1); + if (s < cap) { + list[s] = t; + slot_of[t] = s; + } +} + +// copy: slot s gets row list[s] (row_vec 16-byte words a row), one warp a row. +extern "C" __global__ void tf_glm_kvh_copy(const uint4* __restrict__ lc, const int* __restrict__ list, + const int* __restrict__ n, uint4* __restrict__ slots, int row_vec, int cap) { + const long long s = static_cast(blockIdx.x) * (blockDim.x / 32) + threadIdx.x / 32; + if (s >= min(*n, cap)) return; + const long long t = list[s]; + for (int j = threadIdx.x % 32; j < row_vec; j += 32) slots[s * row_vec + j] = lc[t * row_vec + j]; +} + +// remap: the call's token lists name slots instead of rows (same order, so the attention sums are unchanged). +extern "C" __global__ void tf_glm_kvh_remap(int* tokens, const int* __restrict__ counts, int width, long long total, + const int* __restrict__ slot_of) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + if (static_cast(i % width) >= counts[i / width]) return; + const int t = tokens[i]; + if (t >= 0) tokens[i] = slot_of[t]; +} + +// reset: the rows this call marked are unseen again. +extern "C" __global__ void tf_glm_kvh_reset(const int* __restrict__ list, const int* __restrict__ n, int* slot_of, int cap) { + const long long s = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (s < min(*n, cap)) slot_of[list[s]] = -1; +} + +// rows: a sequence's rows 0 .. *pos + rows (at most max_rows) into slots in position order, through its page table. +extern "C" __global__ void tf_glm_kvh_rows(const uint4* __restrict__ lc, const int* __restrict__ pt, long long page_rows, + const int* __restrict__ pos, int rows, uint4* __restrict__ slots, int row_vec, + int max_rows) { + const long long s = static_cast(blockIdx.x) * (blockDim.x / 32) + threadIdx.x / 32; + if (s >= min(static_cast(*pos) + rows, static_cast(max_rows))) return; + const long long t = pt ? pt[s / PAGE] * page_rows + s % PAGE : s; + for (int j = threadIdx.x % 32; j < row_vec; j += 32) slots[s * row_vec + j] = lc[t * row_vec + j]; +} + +// prev_mark: prev_of[row] = its slot in the layer's row cache, except rows this call writes (rejected drafts' rows). +extern "C" __global__ void tf_glm_kvh_prev_mark(const int* __restrict__ prev_rows, const int* __restrict__ prev_n, + int* prev_of, const int* __restrict__ pt, long long page_rows, + const int* __restrict__ pos, int rows) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= *prev_n) return; + const int t = prev_rows[i]; + for (int r = 0; r < rows; ++r) { + const long long q = static_cast(*pos) + r; + if (pt[q / PAGE] * page_rows + q % PAGE == t) return; + } + prev_of[t] = i; +} + +// copy_cached: as copy, from the layer's cache where the row is there, else from host memory. +extern "C" __global__ void tf_glm_kvh_copy_cached(const uint4* __restrict__ lc, const int* __restrict__ list, + const int* __restrict__ n, uint4* __restrict__ slots, int row_vec, + int cap, const int* __restrict__ prev_of, + const uint4* __restrict__ cache) { + const long long s = static_cast(blockIdx.x) * (blockDim.x / 32) + threadIdx.x / 32; + if (s >= min(*n, cap)) return; + const long long t = list[s]; + const int p = prev_of[t]; + const uint4* src = p >= 0 ? cache + static_cast(p) * row_vec : lc + t * row_vec; + for (int j = threadIdx.x % 32; j < row_vec; j += 32) slots[s * row_vec + j] = src[j]; +} + +// keep: the call's gathered rows become the layer's cache (the first cache_rows of them), one warp a row. +extern "C" __global__ void tf_glm_kvh_keep(const uint4* __restrict__ slots, const int* __restrict__ list, + const int* __restrict__ n, uint4* __restrict__ cache, int* prev_rows, + int* prev_n, int row_vec, int cache_rows) { + const long long s = static_cast(blockIdx.x) * (blockDim.x / 32) + threadIdx.x / 32; + const int m = min(*n, cache_rows); + if (s == 0 && threadIdx.x == 0) *prev_n = m; + if (s >= m) return; + if (threadIdx.x % 32 == 0) prev_rows[s] = list[s]; + for (int j = threadIdx.x % 32; j < row_vec; j += 32) cache[s * row_vec + j] = slots[s * row_vec + j]; +} + +// unmark: the previous call's rows (old_rows, old_n) and this call's (list, n) are unseen again. +extern "C" __global__ void tf_glm_kvh_unmark(const int* __restrict__ old_rows, const int* __restrict__ old_n, + int* prev_of, const int* __restrict__ list, const int* __restrict__ n, + int* slot_of, int cap) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < *old_n) prev_of[old_rows[i]] = -1; + if (i < min(*n, cap)) slot_of[list[i]] = -1; +} diff --git a/zig/src/families/glm/cuda_dsa_block.zig b/zig/src/families/glm/cuda_dsa_block.zig index 515470e85e..788157b131 100644 --- a/zig/src/families/glm/cuda_dsa_block.zig +++ b/zig/src/families/glm/cuda_dsa_block.zig @@ -7,6 +7,7 @@ const projection = @import("cuda_proj.zig"); const latent = @import("cuda_latent.zig"); const sparse = @import("cuda_sparse.zig"); const prof = @import("cuda_prof.zig"); +const kv_host = @import("cuda_kv_host.zig"); /// A prompt chunk's rows one dense latent attention call takes (geometry.MLA_PROMPT_ATT_ROWS). pub const prompt_att_rows = 512; @@ -44,6 +45,7 @@ pub const Caches = struct { pk_rows: usize, pt: u64 = 0, // on pages: the stream's table (i32 page ids), lc and pk then this cache's place in page 0; 0: rows ps: u64 = 0, // on pages: bytes from one page to the next + pks: u64 = 0, // TF_GLM_KV_HOST: the pooled keys' own page stride (their pages on the device); 0: ps }; /// The block's scratch for up to `rows` rows. @@ -135,16 +137,25 @@ pub fn attend(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, v.qi += r0 * ix.heads * ix.d * 2; try writeKeys(k, stream, l, c, v, rows, pos); const t: latent.Table = .{ .pt = c.pt, .ps = c.ps, .page_rows = c.ps / (l.lw * 2) }; + const tk: latent.Table = .{ .pt = c.pt, .ps = if (c.pks != 0) c.pks else c.ps, .page_rows = t.page_rows }; const scale: f32 = @floatCast(std.math.pow(f64, @floatFromInt(l.qk), -0.5)); if (host_pos < l.dense_limit) { const stamp = prof.span(k.prof, stream, .mla); - try denseAttention(k, stream, l, c, v, pos, rows, nch, part_rows, scale); + var dc = c; + if (c.pks != 0) if (kv_host.active) |*g| { // host rows: the visible ones, in order, into the slots + dc = .{ .lc = try g.visible(k.ops, stream, c.lc, c.pt, t.page_rows, pos, rows, l.dense_limit + rows), .ik = c.ik, .ig = c.ig, .pk = c.pk, .pk_rows = c.pk_rows }; + }; + try denseAttention(k, stream, l, dc, v, pos, rows, nch, part_rows, scale); stamp.end(); } if (sparseRows(l, host_pos, rows)) { - try sparse.selectTokensProfile(set, k.ops, stream, v.qi, v.ikr + ix.d * 2, kw_rows, c.pk, t, pos, host_pos, rows, ix.heads, ix.d, c.pk_rows - 2, v.select, v.tokens, v.counts, k.split, k.prof); + try sparse.selectTokensProfile(set, k.ops, stream, v.qi, v.ikr + ix.d * 2, kw_rows, c.pk, tk, pos, host_pos, rows, ix.heads, ix.d, c.pk_rows - 2, v.select, v.tokens, v.counts, k.split, k.prof); const stamp = prof.span(k.prof, stream, .mla); defer stamp.end(); + if (c.pks != 0) if (kv_host.active) |*g| { // host rows: gather the picks into device slots, attend there + const slots = try g.gather(k.ops, stream, ix.gate, c.lc, c.pt, t.page_rows, pos, v.tokens, v.counts, rows, sparse.width); + return latent.sparseAttentionFormat(set, stream, v.qa, slots, 0, .{}, v.tokens, v.counts, sparse.width, v.partials, v.ol, rows, l.heads, l.lw, scale); + }; try latent.sparseAttentionFormat(set, stream, v.qa, c.lc, c.ls, t, v.tokens, v.counts, sparse.width, v.partials, v.ol, rows, l.heads, l.lw, scale); } } @@ -175,11 +186,12 @@ fn writeKeys(k: Kernels, stream: cuda.Stream, l: Layer, c: Caches, s: Scratch, r const kw_rows = ix.d + ix.heads; if (c.ps % (l.lw * 2) != 0) return error.LatentShape; // pages whole latent rows apart (the token lists' unit) const t: latent.Table = .{ .pt = c.pt, .ps = c.ps, .page_rows = c.ps / (l.lw * 2) }; + const tk: latent.Table = .{ .pt = c.pt, .ps = if (c.pks != 0) c.pks else c.ps, .page_rows = t.page_rows }; var stamp = prof.span(k.prof, stream, .cache_write); try latent.latentWriteFormat(set, stream, s.lat, l.lw, c.lc, c.ls, t, pos, rows, l.lw); stamp.end(); stamp = prof.span(k.prof, stream, .index_update); - try sparse.indexUpdate(set, stream, s.ikr, kw_rows, s.igr, ix.ln_w, ix.ln_b, ix.ape, c.ik, c.ig, c.pk, t, pos, rows); + try sparse.indexUpdate(set, stream, s.ikr, kw_rows, s.igr, ix.ln_w, ix.ln_b, ix.ape, c.ik, c.ig, c.pk, tk, pos, rows); stamp.end(); } diff --git a/zig/src/families/glm/cuda_kv_host.zig b/zig/src/families/glm/cuda_kv_host.zig new file mode 100644 index 0000000000..abbece1aa1 --- /dev/null +++ b/zig/src/families/glm/cuda_kv_host.zig @@ -0,0 +1,228 @@ +//! TF_GLM_KV_HOST=1: the DSA layers' latent rows on pinned host pages, each call's picked rows gathered to the device. +const std = @import("std"); +const cuda = @import("cuda"); +const posix = std.posix; + +pub fn onHost() bool { + return set("TF_GLM_KV_HOST"); +} + +/// TF_GLM_KV_HOST_SHARED=1: the ranks on this host share one copy of the pages (their latent rows are identical). +pub fn shared() bool { + return set("TF_GLM_KV_HOST_SHARED"); +} + +fn set(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// Host pages of `page` bytes for a `capacity`-row window: `gib` GiB of them, else one full window and 16. +fn pagesOf(gib: ?f64, page: u64, capacity: usize) usize { + const g = gib orelse return cdiv(capacity, page_rows) + 16; + return @intFromFloat(g * (1 << 30) / @as(f64, @floatFromInt(page))); +} + +/// TF_GLM_KV_HOST_GIB's host pages, else one full window's. +pub fn hostPages(page: u64, capacity: usize) usize { + const v = std.c.getenv("TF_GLM_KV_HOST_GIB") orelse return pagesOf(null, page, capacity); + return pagesOf(std.fmt.parseFloat(f64, std.mem.span(v)) catch null, page, capacity); +} + +/// The plan's `pages` (counted by their device keys alone), held to the host pages when the pages are in host memory. +pub fn pagesFor(page: u64, capacity: usize, pages: usize) usize { + if (!onHost()) return pages; + return @min(pages, hostPages(page, capacity)); +} + +/// Before the ranks measure their room: the gather's device buffers, so the budget counts them. +pub fn prepare(d: *const cuda.Driver, page: u64, capacity: usize, row_bytes: usize) !void { + if (!onHost() or active != null) return; + active = try Gather.init(d, hostPages(page, capacity) * page / row_bytes, capacity, row_bytes); +} + +/// The engine's gather, set when its pages go to host memory. +pub var active: ?Gather = null; + +/// Rows a KV page holds; the row cache's layers (11 DSA and MTP, with room), rows a layer, and its largest call. +const page_rows = 512; +const layers = 16; +const cache_rows = 8192; +const cached_call = 16; + +/// The gather's device buffers (one engine stream runs every call, so one set serves all layers and sequences). +pub const Gather = struct { + slot_of: cuda.DeviceBuffer, // i32 per arena row: its slot in this call, -1 unless marked + list: cuda.DeviceBuffer, // i32 [cap]: the call's rows by slot + slots: cuda.DeviceBuffer, // bf16 [cap][lw] + n: cuda.DeviceBuffer, // i32: the call's rows + prev_of: cuda.DeviceBuffer, // i32 per arena row: its slot in the layer's cache, -1 unless marked + prev_rows: cuda.DeviceBuffer, // i32 [layers][cache_rows]: each layer's last decode call's rows + prev_n: cuda.DeviceBuffer, // i32 [layers] + cache: cuda.DeviceBuffer, // bf16 [layers][cache_rows][lw] + cap: usize, + row_bytes: usize, + tags: [layers]u64 = @splat(0), + + pub fn init(d: *const cuda.Driver, arena_rows: usize, cap: usize, row_bytes: usize) !Gather { + var slot_of = try cuda.DeviceBuffer.alloc(d, arena_rows * 4); + errdefer slot_of.free(); + var list = try cuda.DeviceBuffer.alloc(d, cap * 4); + errdefer list.free(); + var slots = try cuda.DeviceBuffer.alloc(d, cap * row_bytes); + errdefer slots.free(); + var n = try cuda.DeviceBuffer.alloc(d, 16); + errdefer n.free(); + var prev_of = try cuda.DeviceBuffer.alloc(d, arena_rows * 4); + errdefer prev_of.free(); + var prev_rows = try cuda.DeviceBuffer.alloc(d, layers * cache_rows * 4); + errdefer prev_rows.free(); + var prev_n = try cuda.DeviceBuffer.alloc(d, layers * 4); + errdefer prev_n.free(); + const cache = try cuda.DeviceBuffer.alloc(d, layers * cache_rows * row_bytes); + try d.check(d.api.cuMemsetD32_v2(slot_of.ptr, 0xFFFF_FFFF, arena_rows), "cuMemsetD32"); + try d.check(d.api.cuMemsetD32_v2(prev_of.ptr, 0xFFFF_FFFF, arena_rows), "cuMemsetD32"); + try d.check(d.api.cuMemsetD32_v2(prev_n.ptr, 0, layers), "cuMemsetD32"); + std.log.info("kv host: {d} gather slots ({d} MiB), {d} MiB of row marks and a {d} MiB row cache on the device", .{ cap, cap * row_bytes >> 20, 2 * arena_rows * 4 >> 20, layers * cache_rows * row_bytes >> 20 }); + return .{ .slot_of = slot_of, .list = list, .slots = slots, .n = n, .prev_of = prev_of, .prev_rows = prev_rows, .prev_n = prev_n, .cache = cache, .cap = cap, .row_bytes = row_bytes }; + } + + pub fn deinit(g: *Gather) void { + inline for (.{ &g.slot_of, &g.list, &g.slots, &g.n, &g.prev_of, &g.prev_rows, &g.prev_n, &g.cache }) |b| b.free(); + } + + /// Pages changed hands: no layer's cached rows may be used again. + pub fn invalidate(g: *const Gather, stream: cuda.Stream) !void { + try stream.d.check(stream.d.api.cuMemsetD32Async(g.prev_n.ptr, 0, layers, stream.handle), "cuMemsetD32Async"); + } + + /// The cache index of the layer `tag` names (numbered on first sight); null past `layers`. + fn layerOf(g: *Gather, tag: u64) ?usize { + for (&g.tags, 0..) |*t, i| { + if (t.* == tag) return i; + if (t.* == 0) { + t.* = tag; + return i; + } + } + return null; + } + + /// A sparse call's picked rows of `lc` (tokens [rows][width], counts) into the slots; the lists then name slots. + pub fn gather(g: *Gather, ops: cuda.Module, stream: cuda.Stream, tag: u64, lc: u64, pt: u64, page_rows_at: u64, pos: u64, tokens: u64, counts: u64, rows: usize, width: usize) !u64 { + const layer = if (rows <= cached_call) g.layerOf(tag) else null; // decode and verify calls use the row cache + const total = rows * width; + const upper = @min(total, g.cap); + const cap: c_int = @intCast(g.cap); + const rv: c_int = @intCast(g.row_bytes / 16); + try stream.d.check(stream.d.api.cuMemsetD32Async(g.n.ptr, 0, 1, stream.handle), "cuMemsetD32Async"); + try launch(ops, stream, "tf_glm_kvh_mark", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr, g.list.ptr, g.n.ptr, cap }); + const l = layer orelse { + try launch(ops, stream, "tf_glm_kvh_copy", cdiv(upper, 8), .{ lc, g.list.ptr, g.n.ptr, g.slots.ptr, rv, cap }); + try launch(ops, stream, "tf_glm_kvh_remap", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr }); + try launch(ops, stream, "tf_glm_kvh_reset", cdiv(upper, 256), .{ g.list.ptr, g.n.ptr, g.slot_of.ptr, cap }); + return g.slots.ptr; + }; + const prev_rows = g.prev_rows.ptr + l * cache_rows * 4; + const prev_n = g.prev_n.ptr + l * 4; + const cache = g.cache.ptr + l * cache_rows * g.row_bytes; + try launch(ops, stream, "tf_glm_kvh_prev_mark", cdiv(cache_rows, 256), .{ prev_rows, prev_n, g.prev_of.ptr, pt, @as(i64, @intCast(page_rows_at)), pos, @as(c_int, @intCast(rows)) }); + try launch(ops, stream, "tf_glm_kvh_copy_cached", cdiv(upper, 8), .{ lc, g.list.ptr, g.n.ptr, g.slots.ptr, rv, cap, g.prev_of.ptr, cache }); + try launch(ops, stream, "tf_glm_kvh_remap", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr }); + try launch(ops, stream, "tf_glm_kvh_unmark", cdiv(@max(upper, cache_rows), 256), .{ prev_rows, prev_n, g.prev_of.ptr, g.list.ptr, g.n.ptr, g.slot_of.ptr, cap }); + try launch(ops, stream, "tf_glm_kvh_keep", cdiv(cache_rows, 8), .{ g.slots.ptr, g.list.ptr, g.n.ptr, cache, prev_rows, prev_n, rv, @as(c_int, cache_rows) }); + return g.slots.ptr; + } + + /// A dense call's visible rows 0 .. pos + rows (at most `max_rows`) into the slots, in position order. + pub fn visible(g: *const Gather, ops: cuda.Module, stream: cuda.Stream, lc: u64, pt: u64, page_rows_at: u64, pos: u64, rows: usize, max_rows: usize) !u64 { + const n = @min(max_rows, g.cap); + try launch(ops, stream, "tf_glm_kvh_rows", cdiv(n, 8), .{ lc, pt, @as(i64, @intCast(page_rows_at)), pos, @as(c_int, @intCast(rows)), g.slots.ptr, @as(c_int, @intCast(g.row_bytes / 16)), @as(c_int, @intCast(n)) }); + return g.slots.ptr; + } +}; + +/// One glm_ops kernel of 256 threads a block over `blocks` blocks. +fn launch(ops: cuda.Module, stream: cuda.Stream, name: [:0]const u8, blocks: usize, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(try ops.function(name), .{ .grid = .{ .x = @intCast(blocks) }, .block = .{ .x = 256 } }, stream, &a); +} + +/// The shared pages' file for an NCCL id (`host:port`), the same name on every rank. +fn shmName(buf: *[96:0]u8, id: []const u8) ![]u8 { + buf.* = @splat(0); + const name = std.fmt.bufPrint(buf, "/dev/shm/tf-glm-kv-{s}", .{id}) catch return error.SharedMemoryUnavailable; + for (name["/dev/shm/".len..]) |*c| if (c.* == ':' or c.* == '/') { + c.* = '-'; + }; + return name; +} + +/// The pages as one /dev/shm file every rank maps: rank 0 makes it, the others wait for it at its full length. +pub const SharedPages = struct { + region: []align(std.heap.page_size_min) u8, + dev: u64, + lib: std.DynLib, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + leader: bool, + path: [96:0]u8, + + pub fn open(d: *const cuda.Driver, len: usize, leader: bool) !SharedPages { + var path: [96:0]u8 = undefined; + const name = try shmName(&path, if (std.c.getenv("NCCL_COMM_ID")) |v| std.mem.span(v) else "local"); + const flags: std.c.O = if (leader) .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true, .CLOEXEC = true } else .{ .ACCMODE = .RDWR, .CLOEXEC = true }; + var fd: c_int = -1; + var tries: usize = 0; + while (true) : (tries += 1) { // the others may get rank 0's pool before it has made the file + fd = std.c.open(&path, flags, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and leader) break; + if (fd >= 0 and std.c.lseek(fd, 0, std.c.SEEK.END) == @as(std.c.off_t, @intCast(len))) break; + if (fd >= 0) _ = std.c.close(fd); + if (tries == 600) return error.SharedMemoryUnavailable; + const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000_000 }; + _ = std.c.nanosleep(&ts, null); + } + defer _ = std.c.close(fd); + if (leader and std.c.ftruncate(fd, @intCast(len)) != 0) return error.SharedMemoryUnavailable; + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + const register = lib.lookup(*const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, "cuMemHostRegister_v2") orelse return error.MissingSymbol; + const unregister = lib.lookup(*const fn (?*anyopaque) callconv(.c) cuda.abi.Result, "cuMemHostUnregister") orelse return error.MissingSymbol; + try d.check(register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = unregister(region.ptr); + var dev: u64 = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + std.log.info("kv host: pages shared through {s} ({d} MiB, {s})", .{ name, len >> 20, if (leader) "made" else "mapped" }); + return .{ .region = region, .dev = dev, .lib = lib, .unregister = unregister, .leader = leader, .path = path }; + } + + pub fn close(s: *SharedPages) void { + _ = s.unregister(s.region.ptr); + posix.munmap(s.region); + if (s.leader) _ = std.c.unlink(&s.path); + s.lib.close(); + } +}; + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +test "kv host pages: one full window and 16 by default, TF_GLM_KV_HOST_GIB's worth when set" { + const page: u64 = 13_056 * page_rows; // a GLM-5.3-Flash page: 12 DSA and MTP caches' rows + try std.testing.expectEqual(@as(usize, 2048 + 16), pagesOf(null, page, 1 << 20)); + try std.testing.expectEqual(@as(usize, 1 + 16), pagesOf(null, page, 300)); + try std.testing.expectEqual(@as(usize, 2088), pagesOf(13, page, 1 << 20)); + try std.testing.expectEqual(@as(usize, 0), pagesOf(0, page, 1 << 20)); +} + +test "kv host shared pages: one file name per NCCL id, with no separators past /dev/shm" { + var buf: [96:0]u8 = undefined; + try std.testing.expectEqualStrings("/dev/shm/tf-glm-kv-127.0.0.1-29583", try shmName(&buf, "127.0.0.1:29583")); + try std.testing.expectEqualStrings("/dev/shm/tf-glm-kv-a-b-1", try shmName(&buf, "a/b:1")); + const long: [90]u8 = @splat('x'); + try std.testing.expectError(error.SharedMemoryUnavailable, shmName(&buf, &long)); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 06591615c7..bcded6032c 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -74,8 +74,9 @@ pub const Loaded = struct { pub fn attach(p: *anyopaque, streams: u32, pages: usize) !void { const own: *Owned = @ptrCast(@alignCast(p)); 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); + const kept = @import("cuda_kv_host.zig").pagesFor(@import("cuda_state.zig").pageBytes(&own.e.w), own.e.capacity, pages); + try own.lanes.x.pool(streams + 1, kept, l.pageBytes()); + own.e.pages = try pages_mod.Pages.init(own.gpa, own.e, l, streams + 1, kept, true); if (own.share > 0) try share(own, streams + 1); } @@ -224,6 +225,7 @@ fn openRank(a: StartArgs, set: ranks.Settings) !Loaded { // 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; try @import("cuda_prompt_solo.zig").refuse(set.solo, o.drafter); + if (@import("cuda_kv_host.zig").onHost() and !set.pages) return error.GlmKvHostNeedsPages; 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); @@ -305,6 +307,7 @@ fn finishRank(a: OpenArgs) !Loaded { own.link = try follow.Link.init(gpa, ctx.d, &own.group.lib, own.group.comm, own.e.stream, try linkCapacity(cap, o.vision)); errdefer own.link.deinit(); own.e.snap_setup = .{ .ptr = own, .run = readySnapshots }; + try @import("cuda_kv_host.zig").prepare(ctx.d, @import("cuda_state.zig").pageBytes(&own.e.w), own.e.capacity, own.e.w.config.kv_lora * 2); // Rank 0 admits by the smallest room of the ranks; its sizes are the largest (the first ranks take the remainder). const rooms = try gpa.alloc(u64, p.world); defer gpa.free(rooms); @@ -457,6 +460,7 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { 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: set TF_KV_PAGES=on as well", error.GlmSoloWithDrafter => "TF_GLM_CHUNK_SOLO=4096 holds buffers no pass uses beside --drafter: unset it, or leave out --drafter", + error.GlmKvHostNeedsPages => "TF_GLM_KV_HOST keeps the KV pool's pages in host memory: set TF_KV_PAGES=on (and --kv-cache bf16) 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 }, @@ -503,6 +507,7 @@ test "TF_KV_PAGES with --kv-cache fp8: auto and off keep windows at every rank c for ([_][]const u8{ "on", "1", "all" }) |v| try std.testing.expectError(error.Fp8OnPages, paging(v, w, true)); } try std.testing.expect(explain(null, error.Fp8OnPages) != null); + try std.testing.expect(explain(null, error.GlmKvHostNeedsPages) != null); } test "TF_GLM_CHUNK_SOLO=4096 beside --drafter is refused by name" { diff --git a/zig/src/families/glm/cuda_pages.zig b/zig/src/families/glm/cuda_pages.zig index 4bb2f5a5f5..0276a2e0a6 100644 --- a/zig/src/families/glm/cuda_pages.zig +++ b/zig/src/families/glm/cuda_pages.zig @@ -6,6 +6,7 @@ const fw = @import("cuda_forward.zig"); const state = @import("cuda_state.zig"); const types = @import("cuda_engine_types.zig"); const Engine = @import("cuda_engine.zig").Engine; +const kv_host = @import("cuda_kv_host.zig"); /// Table entries a sequence of `capacity` rows takes. pub fn tableLen(capacity: usize) usize { @@ -15,7 +16,8 @@ pub fn tableLen(capacity: usize) usize { /// What a stream of engine `e` holds: its fixed part and table (and `extra`, its other buffers), 13,056 bytes a token. pub fn layout(gpa: std.mem.Allocator, e: *const Engine, extra: u64) !kv.Layout { const fixed = try state.State.pagedBytes(gpa, &e.w, e.capacity, tableLen(e.capacity)); - const page = state.pageBytes(&e.w); + // TF_GLM_KV_HOST: a page costs the device only its pooled keys (the pages themselves are in host memory) + const page = if (kv_host.onHost()) state.pkPageBytes(&e.w) else state.pageBytes(&e.w); return .{ .fixed_bytes = fixed + extra, .token_bytes = page / state.page_rows, .page_tokens = state.page_rows, .snap_bytes = 0 }; } @@ -45,8 +47,13 @@ pub const Pages = struct { gpa: std.mem.Allocator, e: *Engine, layout: kv.Layout, + page: u64, // a page's bytes in its arena (TF_GLM_KV_HOST: the whole page, though the layout counts its keys) pool: kv.Pool, // rank 0 takes and gives back pages; the other ranks map the ids it sends arena: ?cuda.DeviceBuffer, + host: ?cuda.HostBuffer = null, // TF_GLM_KV_HOST: the pages in pinned host memory instead of `arena` + shared: ?kv_host.SharedPages = null, // TF_GLM_KV_HOST_SHARED: the same, one copy for the ranks on this host + host_ptr: u64 = 0, + pk_arena: ?cuda.DeviceBuffer = null, // TF_GLM_KV_HOST: the pooled keys' device pages leases: []kv.Lease, // slot i's: rank 0's lease, or the ids another rank mapped used: []bool, leader: bool, @@ -58,11 +65,27 @@ pub const Pages = struct { if (own == null or own.?.st.table == 0) return error.OwnNotPaged; // slot 0 holds a window, or none was made const p = try gpa.create(Pages); errdefer gpa.destroy(p); - const page = l.pageBytes(); - p.* = .{ .gpa = gpa, .e = e, .layout = l, .pool = try kv.Pool.init(gpa, pages, page), .arena = null, .leases = &.{}, .used = &.{}, .leader = leader, .poison = poisoned() }; + const page = if (kv_host.onHost()) state.pageBytes(&e.w) else l.pageBytes(); + p.* = .{ .gpa = gpa, .e = e, .layout = l, .page = page, .pool = try kv.Pool.init(gpa, pages, page), .arena = null, .leases = &.{}, .used = &.{}, .leader = leader, .poison = poisoned() }; errdefer p.pool.deinit(); - if (pages > 0) p.arena = try cuda.DeviceBuffer.alloc(e.d, pages * page); + const host = kv_host.onHost() and pages > 0; + if (pages > 0 and !host) p.arena = try cuda.DeviceBuffer.alloc(e.d, pages * page); errdefer if (p.arena) |*a| a.free(); + if (host) { + if (kv_host.shared()) { + p.shared = try kv_host.SharedPages.open(e.d, pages * page, leader); + p.host_ptr = p.shared.?.dev; + } else { + p.host = try cuda.HostBuffer.allocMapped(e.d, pages * page); + p.host_ptr = try p.host.?.device(); + } + p.pk_arena = try cuda.DeviceBuffer.alloc(e.d, pages * state.pkPageBytes(&e.w)); + try kv_host.prepare(e.d, page, e.capacity, e.w.config.kv_lora * 2); // made before the room was read + std.log.info("kv host: {d} pages ({d} MiB) in pinned host memory, pooled keys {d} MiB on the device", .{ pages, pages * page >> 20, pages * state.pkPageBytes(&e.w) >> 20 }); + } + errdefer if (p.host) |*h| h.free(); + errdefer if (p.shared) |*h| h.close(); + errdefer if (p.pk_arena) |*a| a.free(); p.leases = try gpa.alloc(kv.Lease, streams); errdefer gpa.free(p.leases); @memset(p.leases, .{}); @@ -78,7 +101,11 @@ pub const Pages = struct { q.* = null; }; errdefer cuda.completion.drain(e.d, e.stream, "GLM page construction did not drain"); - const at = on(e, if (p.arena) |a| a.ptr else 0, page); + var at = on(e, p.base(), page); + if (p.pk_arena) |a| { + at.pk = a.ptr; + at.pk_page = state.pkPageBytes(&e.w); + } for (seqs) |*q| { if (q.* != null) return error.SlotTaken; // attach comes before any stream const s = try e.gpa.create(types.Seq); @@ -102,6 +129,11 @@ pub const Pages = struct { return init(gpa, e, l, streams, pages, false); } + /// The pages' first byte as kernels address it (device arena or mapped host pages). + fn base(p: *const Pages) u64 { + return if (p.arena) |a| a.ptr else p.host_ptr; + } + fn poisoned() bool { const v = std.c.getenv("TF_KV_POISON") orelse return false; return std.mem.eql(u8, std.mem.span(v), "1"); @@ -116,6 +148,9 @@ pub const Pages = struct { p.gpa.free(p.leases); p.gpa.free(p.used); if (p.arena) |*a| a.free(); + if (p.host) |*h| h.free(); + if (p.shared) |*h| h.close(); + if (p.pk_arena) |*a| a.free(); p.pool.deinit(); p.gpa.destroy(p); } @@ -179,11 +214,12 @@ pub const Pages = struct { /// 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[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"); + const page = p.page; + if (p.poison) for (ids[kept..]) |id| try p.e.d.check(p.e.d.api.cuMemsetD32Async(p.base() + 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; + if (kv_host.active) |*g| try g.invalidate(p.e.stream); // the row cache may hold these pages' old rows } /// The stream on paged slot `slot` ended: its sequence reaches no rows, and rank 0's lease gives its pages back. @@ -193,6 +229,7 @@ pub const Pages = struct { st.pages = &.{}; st.capacity = 0; p.used[slot] = false; + if (kv_host.active) |*g| g.invalidate(p.e.stream) catch {}; if (!p.leader) return p.leases[slot].ids.clearRetainingCapacity(); p.pool.finish(&p.leases[slot]); for (p.used) |u| if (u) return; diff --git a/zig/src/families/glm/cuda_snapshot_layout.zig b/zig/src/families/glm/cuda_snapshot_layout.zig index 5454cd38df..8a2531324f 100644 --- a/zig/src/families/glm/cuda_snapshot_layout.zig +++ b/zig/src/families/glm/cuda_snapshot_layout.zig @@ -54,21 +54,21 @@ pub const Plan = struct { const latent = try std.math.mul(usize, at, @as(usize, c.kv_lora) * @as(usize, if (x.ls == 0) 2 else 1)); const ring = try std.math.mul(usize, @min(at, st.ring_rows), st.ring_row); const pooled = try std.math.mul(usize, @min(at / sparse.pool + 2, x.pk_rows), @as(usize, c.i_dim) * 2); - try p.addCache(st, x, x.lc, latent, page_rows * @as(usize, c.kv_lora) * 2); + try p.addCache(st, x, x.lc, latent, page_rows * @as(usize, c.kv_lora) * 2, x.ps); if (x.ls != 0) try p.add(st, x.ls, try std.math.mul(usize, at, @sizeOf(f32))); try p.add(st, st.rings[2 * j], ring); try p.add(st, st.rings[2 * j + 1], ring); - try p.addCache(st, x, x.pk, pooled, page_rows / sparse.pool * @as(usize, c.i_dim) * 2); + try p.addCache(st, x, x.pk, pooled, page_rows / sparse.pool * @as(usize, c.i_dim) * 2, if (x.pks != 0) x.pks else x.ps); } /// 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 { + fn addCache(p: *Plan, st: *const State, x: cache.Caches, live: u64, len: usize, span: usize, page: u64) !void { if (x.pt == 0) return p.add(st, live, len); 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 - from); - p.parts[p.count] = .{ .live = live, .len = len - from, .span = span, .page = x.ps, .from = from }; + p.parts[p.count] = .{ .live = live, .len = len - from, .span = span, .page = page, .from = from }; p.count += 1; } diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 72196feeed..e4ee5f4348 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -56,7 +56,7 @@ fn arena(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { pub const page_rows = latent.chunk; /// Where a paged sequence's caches are: the pool's pages at `kv`, `page_bytes` apart, through a table of `pages`. -pub const Paged = struct { kv: u64, pages: usize, page_bytes: u64 }; +pub const Paged = struct { kv: u64, pages: usize, page_bytes: u64, pk: u64 = 0, pk_page: u64 = 0 }; // pk: TF_GLM_KV_HOST /// The DSA (MLA) layers among the loaded ones; the MTP layer's caches come after theirs. pub fn mlaLayers(w: *const weights.Weights) usize { @@ -70,6 +70,15 @@ fn pageCacheBytes(c: anytype) usize { return page_rows * c.kv_lora * 2 + page_rows / sparse.pool * c.i_dim * 2; } +/// TF_GLM_KV_HOST: one cache's pooled keys on a device page of their own, and that page's bytes. +fn pkCacheBytes(c: anytype) usize { + return page_rows / sparse.pool * c.i_dim * 2; +} + +pub fn pkPageBytes(w: *const weights.Weights) u64 { + return (mlaLayers(w) + 1) * pkCacheBytes(&w.config); +} + /// A page's bytes: every DSA layer's part, then the MTP layer's (13,056 a row on GLM-5.3-Flash). pub fn pageBytes(w: *const weights.Weights) u64 { return (mlaLayers(w) + 1) * pageCacheBytes(&w.config); @@ -78,7 +87,8 @@ pub fn pageBytes(w: *const weights.Weights) u64 { /// Cache `j`'s rows on pages: its part of page 0 at `p.kv`, the others through the table. fn onPages(x: *dsa_block.Caches, c: anytype, p: Paged, table: u64, j: usize) void { x.lc = p.kv + j * pageCacheBytes(c); - x.pk = x.lc + page_rows * c.kv_lora * 2; + x.pk = if (p.pk != 0) p.pk + j * pkCacheBytes(c) else x.lc + page_rows * c.kv_lora * 2; + x.pks = p.pk_page; x.pt = table; x.ps = p.page_bytes; } From 4ae9f099bfe6f0188eb6a1957ee04be89f89154a Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:05:35 -0700 Subject: [PATCH 2/7] glm cuda: the last rank to map TF_GLM_KV_HOST_SHARED's page file unlinks it The shared pages' file kept its name until rank 0 closed it, so a rank stopped by a signal or a crash left the whole window's pages in /dev/shm. A count past the pages now records each rank's mapping, and the rank that completes it unlinks the file: the mappings stay, and nothing outlives the ranks. Rank 0 still unlinks at close if a rank never mapped it. --- zig/src/families/glm/cuda_kv_host.zig | 39 ++++++++++++++++++++++----- zig/src/families/glm/cuda_pages.zig | 2 +- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/zig/src/families/glm/cuda_kv_host.zig b/zig/src/families/glm/cuda_kv_host.zig index abbece1aa1..616d81f752 100644 --- a/zig/src/families/glm/cuda_kv_host.zig +++ b/zig/src/families/glm/cuda_kv_host.zig @@ -159,6 +159,14 @@ fn shmName(buf: *[96:0]u8, id: []const u8) ![]u8 { return name; } +/// The shared file's tail past the pages: a count of the ranks that have mapped it. +const mark_bytes = 4096; + +/// Counts this rank's mapping; true for the last of `world` ranks, which takes the file's name away. +fn lastToMap(count: *u32, world: usize) bool { + return @atomicRmw(u32, count, .Add, 1, .seq_cst) + 1 == world; +} + /// The pages as one /dev/shm file every rank maps: rank 0 makes it, the others wait for it at its full length. pub const SharedPages = struct { region: []align(std.heap.page_size_min) u8, @@ -166,9 +174,11 @@ pub const SharedPages = struct { lib: std.DynLib, unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, leader: bool, + world: usize, // ranks that map the file; the last to map it unlinks it, so nothing outlives the ranks path: [96:0]u8, - pub fn open(d: *const cuda.Driver, len: usize, leader: bool) !SharedPages { + pub fn open(d: *const cuda.Driver, len: usize, leader: bool, world: usize) !SharedPages { + const file_len = len + mark_bytes; var path: [96:0]u8 = undefined; const name = try shmName(&path, if (std.c.getenv("NCCL_COMM_ID")) |v| std.mem.span(v) else "local"); const flags: std.c.O = if (leader) .{ .ACCMODE = .RDWR, .CREAT = true, .TRUNC = true, .CLOEXEC = true } else .{ .ACCMODE = .RDWR, .CLOEXEC = true }; @@ -177,15 +187,15 @@ pub const SharedPages = struct { while (true) : (tries += 1) { // the others may get rank 0's pool before it has made the file fd = std.c.open(&path, flags, @as(std.c.mode_t, 0o600)); if (fd >= 0 and leader) break; - if (fd >= 0 and std.c.lseek(fd, 0, std.c.SEEK.END) == @as(std.c.off_t, @intCast(len))) break; + if (fd >= 0 and std.c.lseek(fd, 0, std.c.SEEK.END) == @as(std.c.off_t, @intCast(file_len))) break; if (fd >= 0) _ = std.c.close(fd); if (tries == 600) return error.SharedMemoryUnavailable; const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000_000 }; _ = std.c.nanosleep(&ts, null); } defer _ = std.c.close(fd); - if (leader and std.c.ftruncate(fd, @intCast(len)) != 0) return error.SharedMemoryUnavailable; - const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + if (leader and std.c.ftruncate(fd, @intCast(file_len)) != 0) return error.SharedMemoryUnavailable; + const region = try posix.mmap(null, file_len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); errdefer posix.munmap(region); var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; errdefer lib.close(); @@ -195,14 +205,18 @@ pub const SharedPages = struct { errdefer _ = unregister(region.ptr); var dev: u64 = 0; try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); - std.log.info("kv host: pages shared through {s} ({d} MiB, {s})", .{ name, len >> 20, if (leader) "made" else "mapped" }); - return .{ .region = region, .dev = dev, .lib = lib, .unregister = unregister, .leader = leader, .path = path }; + const last = lastToMap(@ptrCast(@alignCast(region.ptr + len)), world); + if (last) _ = std.c.unlink(&path); + std.log.info("kv host: pages shared through {s} ({d} MiB, {s}{s})", .{ name, len >> 20, if (leader) "made" else "mapped", if (last) ", every rank mapped: unlinked" else "" }); + return .{ .region = region, .dev = dev, .lib = lib, .unregister = unregister, .leader = leader, .world = world, .path = path }; } pub fn close(s: *SharedPages) void { + const count: *u32 = @ptrCast(@alignCast(s.region.ptr + s.region.len - mark_bytes)); + const named = s.leader and @atomicLoad(u32, count, .seq_cst) < s.world; // a rank never mapped it: still named _ = s.unregister(s.region.ptr); posix.munmap(s.region); - if (s.leader) _ = std.c.unlink(&s.path); + if (named) _ = std.c.unlink(&s.path); s.lib.close(); } }; @@ -219,6 +233,17 @@ test "kv host pages: one full window and 16 by default, TF_GLM_KV_HOST_GIB's wor try std.testing.expectEqual(@as(usize, 0), pagesOf(0, page, 1 << 20)); } +test "kv host shared pages: the last of the ranks to map the file takes its name away" { + var count: u32 = 0; + try std.testing.expect(!lastToMap(&count, 2)); + try std.testing.expect(lastToMap(&count, 2)); + count = 0; + try std.testing.expect(lastToMap(&count, 1)); + count = 0; + for (0..2) |_| try std.testing.expect(!lastToMap(&count, 3)); + try std.testing.expect(lastToMap(&count, 3)); +} + test "kv host shared pages: one file name per NCCL id, with no separators past /dev/shm" { var buf: [96:0]u8 = undefined; try std.testing.expectEqualStrings("/dev/shm/tf-glm-kv-127.0.0.1-29583", try shmName(&buf, "127.0.0.1:29583")); diff --git a/zig/src/families/glm/cuda_pages.zig b/zig/src/families/glm/cuda_pages.zig index 0276a2e0a6..509368c5ba 100644 --- a/zig/src/families/glm/cuda_pages.zig +++ b/zig/src/families/glm/cuda_pages.zig @@ -73,7 +73,7 @@ pub const Pages = struct { errdefer if (p.arena) |*a| a.free(); if (host) { if (kv_host.shared()) { - p.shared = try kv_host.SharedPages.open(e.d, pages * page, leader); + p.shared = try kv_host.SharedPages.open(e.d, pages * page, leader, e.share.world); p.host_ptr = p.shared.?.dev; } else { p.host = try cuda.HostBuffer.allocMapped(e.d, pages * page); From 76612bb0240e0f340771713b63c3a981993a2d72 Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:27:09 -0700 Subject: [PATCH 3/7] glm cuda: TF_GLM_KV_HOST_SHARED's /dev/shm pages build on Linux only The shared page file used POSIX open flags and DynLib.close, which the Windows test builds can't compile. Off Linux, SharedPages is now a stub that refuses with SharedMemoryUnavailable, so each rank keeps its own pages; Linux is unchanged. --- zig/src/families/glm/cuda_kv_host.zig | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/zig/src/families/glm/cuda_kv_host.zig b/zig/src/families/glm/cuda_kv_host.zig index 616d81f752..6a369eea9f 100644 --- a/zig/src/families/glm/cuda_kv_host.zig +++ b/zig/src/families/glm/cuda_kv_host.zig @@ -167,8 +167,22 @@ fn lastToMap(count: *u32, world: usize) bool { return @atomicRmw(u32, count, .Add, 1, .seq_cst) + 1 == world; } -/// The pages as one /dev/shm file every rank maps: rank 0 makes it, the others wait for it at its full length. -pub const SharedPages = struct { +/// The pages as one /dev/shm file every rank maps (Linux; elsewhere TF_GLM_KV_HOST_SHARED is refused). +pub const SharedPages = if (@import("builtin").os.tag == .linux) ShmPages else NoSharedPages; + +/// Off Linux: no /dev/shm, so the shared copy is refused and each rank keeps its own pages. +const NoSharedPages = struct { + dev: u64, + + pub fn open(_: *const cuda.Driver, _: usize, _: bool, _: usize) !NoSharedPages { + return error.SharedMemoryUnavailable; + } + + pub fn close(_: *NoSharedPages) void {} +}; + +/// Rank 0 makes the file, the others wait for it at its full length; the last rank to map it unlinks it. +const ShmPages = struct { region: []align(std.heap.page_size_min) u8, dev: u64, lib: std.DynLib, @@ -177,7 +191,7 @@ pub const SharedPages = struct { world: usize, // ranks that map the file; the last to map it unlinks it, so nothing outlives the ranks path: [96:0]u8, - pub fn open(d: *const cuda.Driver, len: usize, leader: bool, world: usize) !SharedPages { + pub fn open(d: *const cuda.Driver, len: usize, leader: bool, world: usize) !ShmPages { const file_len = len + mark_bytes; var path: [96:0]u8 = undefined; const name = try shmName(&path, if (std.c.getenv("NCCL_COMM_ID")) |v| std.mem.span(v) else "local"); @@ -211,7 +225,7 @@ pub const SharedPages = struct { return .{ .region = region, .dev = dev, .lib = lib, .unregister = unregister, .leader = leader, .world = world, .path = path }; } - pub fn close(s: *SharedPages) void { + pub fn close(s: *ShmPages) void { const count: *u32 = @ptrCast(@alignCast(s.region.ptr + s.region.len - mark_bytes)); const named = s.leader and @atomicLoad(u32, count, .seq_cst) < s.world; // a rank never mapped it: still named _ = s.unregister(s.region.ptr); From 3a132d63c9696a366822936113151ae4e6c7a693 Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sat, 31 Oct 2026 21:05:28 -0700 Subject: [PATCH 4/7] glm cuda: TF_GLM_KV_HOST, _SHARED and _GIB join the ranks' fingerprint A rank started without TF_GLM_KV_HOST would size a different pool from its peers, _SHARED needs every rank to map the same file, and each rank sizes its row marks from its own _GIB while rank 0 counts the pages. Settings now reads the three (refusing a bad value, _SHARED without the base, or a _GIB that isn't positive), and the fingerprint carries kv_host (0, 1 or 2) and kv_host_gib, so ranks that disagree stop at startup naming the setting. --- zig/src/families/glm/cuda_ranks.zig | 33 +++++++++++++++++++- zig/src/families/glm/cuda_ranks_test.zig | 39 +++++++++++++++++++++++- 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index 5b6bd58c61..d6b128f75c 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -22,6 +22,8 @@ pub const Settings = struct { dense_bits: u8 = 16, // TF_GLM_DENSE_BITS: the dense linears packed at load to 4, 6 or 8 bits (16: bf16) head_bits: u8 = 16, // TF_GLM_HEAD_BITS: the head, likewise roce_multi: bool = false, // TF_GLM_ROCE_MULTI: 1 (three ranks or more: small all-gathers on fabric/tp's multirank RoCE beside NCCL) or 0 + kv_host: u8 = 0, // TF_GLM_KV_HOST=1: KV pages in host memory (1); with TF_GLM_KV_HOST_SHARED=1, one copy a host (2) + kv_host_gib: f64 = 0, // TF_GLM_KV_HOST_GIB: the host pages' GiB (0: one full window) /// The settings `get` (an environment lookup) gives; `problem` names a value that is not one of the choices. pub fn read(get: *const fn ([:0]const u8) ?[]const u8, problem: *[]const u8) !Settings { @@ -107,6 +109,25 @@ pub const Settings = struct { problem.* = "TF_GLM_PROMPT_K=shape requires TF_GLM_MM=1"; return error.BadSetting; } + if (get("TF_GLM_KV_HOST")) |v| { + if (std.mem.eql(u8, v, "1")) s.kv_host = 1 else if (!std.mem.eql(u8, v, "0")) { + problem.* = "TF_GLM_KV_HOST: 0 or 1"; + return error.BadSetting; + } + } + if (get("TF_GLM_KV_HOST_SHARED")) |v| { + if (std.mem.eql(u8, v, "1") and s.kv_host == 1) s.kv_host = 2 else if (!std.mem.eql(u8, v, "0")) { + problem.* = if (std.mem.eql(u8, v, "1")) "TF_GLM_KV_HOST_SHARED=1 requires TF_GLM_KV_HOST=1" else "TF_GLM_KV_HOST_SHARED: 0 or 1"; + return error.BadSetting; + } + } + if (get("TF_GLM_KV_HOST_GIB")) |v| { + s.kv_host_gib = std.fmt.parseFloat(f64, v) catch 0; + if (!(s.kv_host_gib > 0) or s.kv_host == 0) { + problem.* = "TF_GLM_KV_HOST_GIB: the host pages' GiB, more than 0, with TF_GLM_KV_HOST=1"; + return error.BadSetting; + } + } // refused here, before the weights load: the wire's bootstrap meets at rank 0, its kernel comes from the image if (s.roce_multi and (get("TF_TP_MASTER") == null or get("TF_TP_KERNEL_IMAGE") == null)) { problem.* = "TF_GLM_ROCE_MULTI=1 requires TF_TP_MASTER and TF_TP_KERNEL_IMAGE"; @@ -126,7 +147,7 @@ pub fn environment(name: [:0]const u8) ?[]const u8 { /// What every rank must agree on, as bytes the communicator carries. pub const Fingerprint = extern struct { - layout: u32 = 14, // this struct's version + layout: u32 = 15, // this struct's version world: u32, rank: u32, kv8: u32 = 0, @@ -144,11 +165,13 @@ pub const Fingerprint = extern struct { max_videos: u32 = 0, max_frames: u32 = 0, roce_multi: u32 = 0, // TF_GLM_ROCE_MULTI: every rank opens the borrowed wire or none does (a lone rank would wait on it) + kv_host: u32 = 0, // TF_GLM_KV_HOST(_SHARED): every rank sizes the pool alike and maps the shared file, or none does solo: u32 = 0, // TF_GLM_CHUNK_SOLO: rank 0's solo passes need the wider prompt buffers on every rank bits: u32 = 16 << 8 | 16, // TF_GLM_DENSE_BITS << 8 | TF_GLM_HEAD_BITS: every rank's matmuls read the same codes capacity: u64, // tokens a sequence holds heat_high: f64, // the heat bands in degrees (0 and 0: off): every rank must wait alike heat_low: f64, + kv_host_gib: f64 = 0, // TF_GLM_KV_HOST_GIB: each rank's row marks span the host pages rank 0 counts config: [32]u8, // sha256 of config.json program: [32]u8, // sha256 of this executable checkpoint: [64]u8, // hex digest of the checkpoint's metadata and every shard byte, through the digest cache @@ -172,11 +195,13 @@ pub const Fingerprint = extern struct { max_videos, max_frames, roce_multi, + kv_host, solo, bits, capacity, heat_high, heat_low, + kv_host_gib, config, program, checkpoint, @@ -205,6 +230,8 @@ pub const Fingerprint = extern struct { .layers = s.layers, .mtp = @intFromBool(mtp), .roce_multi = @intFromBool(s.roce_multi), + .kv_host = s.kv_host, + .kv_host_gib = s.kv_host_gib, .solo = @intFromBool(s.solo), .bits = @as(u32, s.dense_bits) << 8 | s.head_bits, .capacity = capacity, @@ -242,11 +269,13 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "max_videos", "max_frames", "roce_multi", + "kv_host", "solo", "bits", "capacity", "heat_high", "heat_low", + "kv_host_gib", "config", "program", "checkpoint", @@ -352,6 +381,8 @@ fn explain(f: Fingerprint.Field) []const u8 { .max_videos => "the video count limit", .max_frames => "the frame limit", .roce_multi => "TF_GLM_ROCE_MULTI", + .kv_host => "TF_GLM_KV_HOST or TF_GLM_KV_HOST_SHARED", + .kv_host_gib => "TF_GLM_KV_HOST_GIB", .solo => "TF_GLM_CHUNK_SOLO", .bits => "TF_GLM_DENSE_BITS or TF_GLM_HEAD_BITS", .capacity => "the sequence capacity (the context)", diff --git a/zig/src/families/glm/cuda_ranks_test.zig b/zig/src/families/glm/cuda_ranks_test.zig index 09ecd38fd4..38006176ec 100644 --- a/zig/src/families/glm/cuda_ranks_test.zig +++ b/zig/src/families/glm/cuda_ranks_test.zig @@ -186,7 +186,7 @@ test "every media limit is part of the ranks' agreement" { test "every scheduling and identity field that differs stops the ranks" { const zero = Fingerprint.of(2, 0, .{}, 8192, true, @splat(0), @splat(0), hashBytes("/same")); const one = Fingerprint.of(2, 1, .{}, 8192, true, @splat(0), @splat(0), hashBytes("/same")); - inline for (.{ "layout", "world", "rank", "pages", "split", "kda_wide", "pieces", "solo", "layers", "mtp", "vision", "video_tokens", "max_images", "max_videos", "max_frames", "roce_multi", "capacity", "heat_high", "heat_low" }) |name| { + inline for (.{ "layout", "world", "rank", "pages", "split", "kda_wide", "pieces", "solo", "layers", "mtp", "vision", "video_tokens", "max_images", "max_videos", "max_frames", "roce_multi", "kv_host", "capacity", "heat_high", "heat_low", "kv_host_gib" }) |name| { var all = [_]Fingerprint{ zero, one }; @field(all[1], name) += 1; try std.testing.expectEqual(@field(Fingerprint.Field, name), mismatch(&all).?.field); @@ -347,3 +347,40 @@ test "TF_GLM_ROCE_MULTI: off by default, 1 and 0, and refusals naming a bad valu const h: [32]u8 = @splat(0); try std.testing.expectEqual(@as(u32, 1), Fingerprint.of(3, 0, .{ .roce_multi = true }, 8192, true, h, h, h).roce_multi); } + +test "TF_GLM_KV_HOST, _SHARED and _GIB: off by default, the choices, and refusals naming a bad value or a missing base" { + const Reader = struct { + var vars: []const [2][]const u8 = &.{}; + fn get(key: [:0]const u8) ?[]const u8 { + for (vars) |kv| if (std.mem.eql(u8, kv[0], key)) return kv[1]; + return null; + } + }; + var problem: []const u8 = ""; + Reader.vars = &.{}; + try std.testing.expectEqual(@as(u8, 0), (try Settings.read(Reader.get, &problem)).kv_host); + Reader.vars = &.{.{ "TF_GLM_KV_HOST", "1" }}; + try std.testing.expectEqual(@as(u8, 1), (try Settings.read(Reader.get, &problem)).kv_host); + Reader.vars = &.{ .{ "TF_GLM_KV_HOST", "1" }, .{ "TF_GLM_KV_HOST_SHARED", "1" }, .{ "TF_GLM_KV_HOST_GIB", "13" } }; + const shared = try Settings.read(Reader.get, &problem); + try std.testing.expectEqual(@as(u8, 2), shared.kv_host); + try std.testing.expectEqual(@as(f64, 13), shared.kv_host_gib); + const refusals = .{ + .{ &[_][2][]const u8{.{ "TF_GLM_KV_HOST", "yes" }}, "TF_GLM_KV_HOST: 0 or 1" }, + .{ &[_][2][]const u8{.{ "TF_GLM_KV_HOST_SHARED", "1" }}, "TF_GLM_KV_HOST_SHARED=1 requires TF_GLM_KV_HOST=1" }, + .{ &[_][2][]const u8{ .{ "TF_GLM_KV_HOST", "1" }, .{ "TF_GLM_KV_HOST_SHARED", "2" } }, "TF_GLM_KV_HOST_SHARED: 0 or 1" }, + .{ &[_][2][]const u8{.{ "TF_GLM_KV_HOST_GIB", "13" }}, "TF_GLM_KV_HOST_GIB: the host pages' GiB, more than 0, with TF_GLM_KV_HOST=1" }, + .{ &[_][2][]const u8{ .{ "TF_GLM_KV_HOST", "1" }, .{ "TF_GLM_KV_HOST_GIB", "0" } }, "TF_GLM_KV_HOST_GIB: the host pages' GiB, more than 0, with TF_GLM_KV_HOST=1" }, + }; + inline for (refusals) |r| { + Reader.vars = r[0]; + try std.testing.expectError(error.BadSetting, Settings.read(Reader.get, &problem)); + try std.testing.expectEqualStrings(r[1], problem); + } + const h: [32]u8 = @splat(0); + const f = Fingerprint.of(2, 0, shared, 8192, true, h, h, h); + try std.testing.expectEqual(@as(u32, 2), f.kv_host); + try std.testing.expectEqual(@as(f64, 13), f.kv_host_gib); + var all = [_]Fingerprint{ f, Fingerprint.of(2, 1, .{ .kv_host = 1 }, 8192, true, h, h, h) }; + try std.testing.expectEqual(Fingerprint.Field.kv_host, mismatch(&all).?.field); +} From 1ae7acb76077a921972f9ecf3e7380199544cbe2 Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sun, 1 Nov 2026 12:52:47 -0800 Subject: [PATCH 5/7] glm cuda: TF_GLM_KV_HOST's gather kernels join the contract, and a call past its slots is refused contract_glm_kv_host.zig checks the nine tf_glm_kvh_* kernels exactly against a host replay, at decode, verify and 2,048-row chunk shapes with pages out of order, each with a planted control (a stale mark, a dropped row or swapped slots); glm_cuda_records.zig records their launches from nvcc 13.0. A call's distinct rows are bounded by the sequence's capacity, which Gather.cap equals, but mark no longer leaves a row past the slots at -2: it sets it unseen and raises a flag in mapped host memory, which Engine.checkGathers reads after its synchronize and refuses with error.GlmKvHostSlots. --- docs/kernel-contract.md | 5 + zig/kernels/cuda/glm_ops.cu | 6 +- zig/src/core/contract/glm_cuda_records.zig | 9 + zig/src/families/glm/cuda_engine.zig | 3 +- zig/src/families/glm/cuda_kv_host.zig | 24 +- zig/tests/cuda/contract_glm_controls.zig | 13 +- zig/tests/cuda/contract_glm_controls_host.zig | 2 +- zig/tests/cuda/contract_glm_kv_host.zig | 554 ++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 2 +- 9 files changed, 608 insertions(+), 10 deletions(-) create mode 100644 zig/tests/cuda/contract_glm_kv_host.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 55abdb5d2c..cd66bcbf6d 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -266,6 +266,11 @@ sparse). Its prompt indexer (`tf_glm_index_scores`, `tf_glm_index_pools`) is bit `tf_glm_select_rows`, and its decode select at a 1,024-thread block a row (`tf_glm_select_rows_wide`) to `tf_glm_select_rows`; both stated rules rank any NaN above +inf, as the GPU's `s + 0` makes every NaN one positive NaN. +CUDA also checks `TF_GLM_KV_HOST`'s gather (`tf_glm_kvh_*` in `glm_ops.cu`, launched by `cuda_kv_host.zig`): the +marks (a slot for each distinct pick; past its slots, the call flagged and no row left marked), the picked rows into +their slots from host pages or the layer's row cache, the token lists' slots, the marks cleared, the dense rows through +pages out of order, and the layer's cache kept, each exact against the host's replay (nine controls, one a kernel: a +stale mark, a dropped row or swapped slots). GLM's Metal window projections register narrow and wide rows at 2, 3, 4, 5, 6 and 8 bits, plus both gathered paths. Independent FP64 linear cases accompany byte replacements against frozen pre-window-optimization kernels. diff --git a/zig/kernels/cuda/glm_ops.cu b/zig/kernels/cuda/glm_ops.cu index a6a9fd202c..b2fbee106c 100644 --- a/zig/kernels/cuda/glm_ops.cu +++ b/zig/kernels/cuda/glm_ops.cu @@ -56,7 +56,8 @@ extern "C" __global__ void tf_glm_copy_rows(const unsigned char* __restrict__ sr // TF_GLM_KV_HOST mark: each picked row (first counts[r] of tokens [rows][width]) gets a slot; slot_of -1 if unseen. extern "C" __global__ void tf_glm_kvh_mark(const int* __restrict__ tokens, const int* __restrict__ counts, int width, - long long total, int* slot_of, int* list, int* n, int cap) { + long long total, int* slot_of, int* list, int* n, int cap, + volatile int* over) { const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (i >= total) return; if (static_cast(i % width) >= counts[i / width]) return; @@ -66,6 +67,9 @@ extern "C" __global__ void tf_glm_kvh_mark(const int* __restrict__ tokens, const if (s < cap) { list[s] = t; slot_of[t] = s; + } else { + slot_of[t] = -1; // past the slots: left unseen, and the call flagged in host memory (Gather.check reads it) + *over = 1; } } diff --git a/zig/src/core/contract/glm_cuda_records.zig b/zig/src/core/contract/glm_cuda_records.zig index 6d23ed6650..578debeddd 100644 --- a/zig/src/core/contract/glm_cuda_records.zig +++ b/zig/src/core/contract/glm_cuda_records.zig @@ -48,6 +48,15 @@ pub const cuda = [_]registry.CudaRecord{ .{ .launch = .{ .name = "glm_glue tf_glm_chunks_16", .symbol = "tf_glm_chunks_16", .image = "glm_glue", .threads = 256, .dynamic = 83328, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 123, .static = 0, .bound = 256 }, .{ .regs = 123, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "glm_glue tf_glm_chunks_32", .symbol = "tf_glm_chunks_32", .image = "glm_glue", .threads = 256, .dynamic = 101120, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 130, .static = 0, .bound = 256 }, .{ .regs = 130, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "glm_ops tf_glm_dsa_tokens", .symbol = "tf_glm_dsa_tokens", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_mark", .symbol = "tf_glm_kvh_mark", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_copy", .symbol = "tf_glm_kvh_copy", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_remap", .symbol = "tf_glm_kvh_remap", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_reset", .symbol = "tf_glm_kvh_reset", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_rows", .symbol = "tf_glm_kvh_rows", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_prev_mark", .symbol = "tf_glm_kvh_prev_mark", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 15, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_copy_cached", .symbol = "tf_glm_kvh_copy_cached", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_keep", .symbol = "tf_glm_kvh_keep", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 46, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 }, .{ .regs = 42, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_ops tf_glm_kvh_unmark", .symbol = "tf_glm_kvh_unmark", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 } } }, }; /// nvcc 13.2's usage where it differs from 13.0's. diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 87a4b07a06..774acc78fc 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -573,9 +573,10 @@ pub const Engine = struct { const greedy = @import("cuda_engine_pick.zig").greedy; const greedyChance = @import("cuda_engine_pick.zig").greedyChance; - /// A failed RoCE gather's error, after a synchronize whose result read gathered rows (they hold nothing then). + /// A failed RoCE gather (its rows hold nothing) or TF_GLM_KV_HOST gather past its slots, after a synchronize. pub fn checkGathers(e: *const Engine) !void { if (e.group) |g| try g.check(); + if (@import("cuda_kv_host.zig").active) |*g| try g.check(); // a TF_GLM_KV_HOST gather past its slots } pub fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void { diff --git a/zig/src/families/glm/cuda_kv_host.zig b/zig/src/families/glm/cuda_kv_host.zig index 6a369eea9f..d06b6cd28a 100644 --- a/zig/src/families/glm/cuda_kv_host.zig +++ b/zig/src/families/glm/cuda_kv_host.zig @@ -60,6 +60,8 @@ pub const Gather = struct { prev_rows: cuda.DeviceBuffer, // i32 [layers][cache_rows]: each layer's last decode call's rows prev_n: cuda.DeviceBuffer, // i32 [layers] cache: cuda.DeviceBuffer, // bf16 [layers][cache_rows][lw] + over: cuda.HostBuffer, // i32 in mapped host memory: set when a call picked more distinct rows than `cap` + over_dev: u64, // its device address cap: usize, row_bytes: usize, tags: [layers]u64 = @splat(0), @@ -79,16 +81,30 @@ pub const Gather = struct { errdefer prev_rows.free(); var prev_n = try cuda.DeviceBuffer.alloc(d, layers * 4); errdefer prev_n.free(); - const cache = try cuda.DeviceBuffer.alloc(d, layers * cache_rows * row_bytes); + var cache = try cuda.DeviceBuffer.alloc(d, layers * cache_rows * row_bytes); + errdefer cache.free(); + var over = try cuda.HostBuffer.allocMapped(d, 16); + errdefer over.free(); + @memset(over.bytes, 0); + const over_dev = try over.device(); try d.check(d.api.cuMemsetD32_v2(slot_of.ptr, 0xFFFF_FFFF, arena_rows), "cuMemsetD32"); try d.check(d.api.cuMemsetD32_v2(prev_of.ptr, 0xFFFF_FFFF, arena_rows), "cuMemsetD32"); try d.check(d.api.cuMemsetD32_v2(prev_n.ptr, 0, layers), "cuMemsetD32"); std.log.info("kv host: {d} gather slots ({d} MiB), {d} MiB of row marks and a {d} MiB row cache on the device", .{ cap, cap * row_bytes >> 20, 2 * arena_rows * 4 >> 20, layers * cache_rows * row_bytes >> 20 }); - return .{ .slot_of = slot_of, .list = list, .slots = slots, .n = n, .prev_of = prev_of, .prev_rows = prev_rows, .prev_n = prev_n, .cache = cache, .cap = cap, .row_bytes = row_bytes }; + return .{ .slot_of = slot_of, .list = list, .slots = slots, .n = n, .prev_of = prev_of, .prev_rows = prev_rows, .prev_n = prev_n, .cache = cache, .over = over, .over_dev = over_dev, .cap = cap, .row_bytes = row_bytes }; } pub fn deinit(g: *Gather) void { - inline for (.{ &g.slot_of, &g.list, &g.slots, &g.n, &g.prev_of, &g.prev_rows, &g.prev_n, &g.cache }) |b| b.free(); + inline for (.{ &g.slot_of, &g.list, &g.slots, &g.n, &g.prev_of, &g.prev_rows, &g.prev_n, &g.cache, &g.over }) |b| b.free(); + } + + /// After a synchronize: refuses (and clears) a call since the last check whose distinct picks passed the slots. + pub fn check(g: *const Gather) !void { + const flag: *volatile i32 = @ptrCast(g.over.slice(i32).ptr); + if (flag.* == 0) return; + flag.* = 0; + std.log.err("kv host: a call picked more distinct rows than the gather's {d} slots", .{g.cap}); + return error.GlmKvHostSlots; } /// Pages changed hands: no layer's cached rows may be used again. @@ -116,7 +132,7 @@ pub const Gather = struct { const cap: c_int = @intCast(g.cap); const rv: c_int = @intCast(g.row_bytes / 16); try stream.d.check(stream.d.api.cuMemsetD32Async(g.n.ptr, 0, 1, stream.handle), "cuMemsetD32Async"); - try launch(ops, stream, "tf_glm_kvh_mark", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr, g.list.ptr, g.n.ptr, cap }); + try launch(ops, stream, "tf_glm_kvh_mark", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr, g.list.ptr, g.n.ptr, cap, g.over_dev }); const l = layer orelse { try launch(ops, stream, "tf_glm_kvh_copy", cdiv(upper, 8), .{ lc, g.list.ptr, g.n.ptr, g.slots.ptr, rv, cap }); try launch(ops, stream, "tf_glm_kvh_remap", cdiv(total, 256), .{ tokens, counts, @as(c_int, @intCast(width)), @as(i64, @intCast(total)), g.slot_of.ptr }); diff --git a/zig/tests/cuda/contract_glm_controls.zig b/zig/tests/cuda/contract_glm_controls.zig index 74a559dd3c..83868a1f7b 100644 --- a/zig/tests/cuda/contract_glm_controls.zig +++ b/zig/tests/cuda/contract_glm_controls.zig @@ -1,4 +1,4 @@ -//! Real GLM launches with planted page-table or sparse physical-row mismatches, one per paged contract case. +//! Real GLM launches with planted page-table or sparse physical-row mismatches, and TF_GLM_KV_HOST's gather kernels. const std = @import("std"); const cuda = @import("cuda"); const glm = @import("glm_cuda"); @@ -7,6 +7,7 @@ const g = @import("core").contract.gates; const pages = @import("contract_glm_pages.zig"); const index = @import("contract_glm_index.zig"); const attn = @import("contract_glm_attn.zig"); +const kv_host = @import("contract_glm_kv_host.zig"); pub const Controls = struct { lat: pages.LatWrite, @@ -14,14 +15,18 @@ pub const Controls = struct { scores: index.Scores, tokens: index.Tokens, chunks: [4]attn.Chunks, + ops: ?cuda.Module, // glm_ops: TF_GLM_KV_HOST's gather kernels and their controls (null: the page cases alone) + kvh: ?*kv_host.Set = null, + a: ?std.mem.Allocator = null, - pub fn init(gpu: check.Gpu, s: cuda.Stream, f: glm.glue.Kernels, tokens: cuda.Function) Controls { + pub fn init(gpu: check.Gpu, s: cuda.Stream, f: glm.glue.Kernels, tokens: cuda.Function, ops: ?cuda.Module) Controls { var out: Controls = .{ .lat = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f.lat_write, .wrong_table = true }, .pool = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f.pool_keys, .wrong_table = true }, .scores = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f.scores, .wrong_table = true }, .tokens = .{ .rig = .{ .gpu = gpu, .s = s }, .f = tokens, .wrong_table = true }, .chunks = undefined, + .ops = ops, }; for (&out.chunks, [_]u32{ 16, 16, 32, 32 }, [_]bool{ false, true, false, true }) |*x, hb, sparse| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = if (hb == 16) f.chunks_16 else f.chunks_32, .merge = f.merge, .form = .{ .hb = hb, .sparse = sparse }, .wrong_table = true }; @@ -34,6 +39,7 @@ pub const Controls = struct { self.scores.deinit(); self.tokens.deinit(); for (&self.chunks) |*x| x.deinit(); + if (self.kvh) |k| k.destroy(self.a.?); } pub fn append(self: *Controls, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !void { @@ -42,5 +48,8 @@ pub const Controls = struct { try kernels.append(a, self.scores.kernel()); try kernels.append(a, self.tokens.kernel()); for (&self.chunks) |*x| try kernels.append(a, x.kernel()); + const ops = self.ops orelse return; + self.kvh = try kv_host.Set.create(a, self.lat.rig.gpu, self.lat.rig.s, ops, kernels); + self.a = a; } }; diff --git a/zig/tests/cuda/contract_glm_controls_host.zig b/zig/tests/cuda/contract_glm_controls_host.zig index 4edba66b3e..81b8a0c6bc 100644 --- a/zig/tests/cuda/contract_glm_controls_host.zig +++ b/zig/tests/cuda/contract_glm_controls_host.zig @@ -24,7 +24,7 @@ fn registered(fault: bool) !void { @field(functions, name) = .{ .d = &d, .handle = null }; const gpu: check.Gpu = .{ .d = &d, .ctx = &ctx, .gpa = t.allocator, .io = t.io }; const stream: cuda.Stream = .{ .d = &d, .handle = null }; - var controls = Controls.init(gpu, stream, functions, .{ .d = &d, .handle = null }); + var controls = Controls.init(gpu, stream, functions, .{ .d = &d, .handle = null }, null); // Only descriptors are read: empty rigs own no allocations and no CUDA callback is invoked. if (fault) controls.lat.wrong_table = false; var kernels: std.ArrayList(g.Kernel) = .empty; diff --git a/zig/tests/cuda/contract_glm_kv_host.zig b/zig/tests/cuda/contract_glm_kv_host.zig new file mode 100644 index 0000000000..9c685bcdfe --- /dev/null +++ b/zig/tests/cuda/contract_glm_kv_host.zig @@ -0,0 +1,554 @@ +//! TF_GLM_KV_HOST's gather kernels (glm_ops.cu) under the contract on CUDA: ids and rows moved, checked exactly. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const pages = @import("contract_glm_pages.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const sparse = glm.sparse; + +const row_bytes = pages.lw * 2; // a latent row, as the gather moves it +const rv = row_bytes / 16; // its 16-byte words +const page = pages.rows_a_page; +const cache_rows = 8192; // a layer's row cache, as cuda_kv_host.zig keeps it +const at = 3 * page - 7; // the call's first position: its rows cross a page edge + +pub const Which = enum { mark, copy, remap, reset, rows, prev_mark, copy_cached, keep, unmark }; + +/// What a planted control breaks, and the clean case each one shadows. +const planted = std.enums.EnumArray(Which, []const u8).init(.{ + .mark = "control: GLM kv host mark meets a stale mark", + .copy = "control: GLM kv host copy drops its last row", + .remap = "control: GLM kv host remap reads swapped slots", + .reset = "control: GLM kv host reset drops its last row", + .rows = "control: GLM kv host dense rows use a swapped table", + .prev_mark = "control: GLM kv host prev_mark keeps the rows this call rewrites (a stale mark)", + .copy_cached = "control: GLM kv host copy_cached reads swapped cache slots", + .keep = "control: GLM kv host keep drops its last row", + .unmark = "control: GLM kv host unmark drops a previous row (a stale mark)", +}); + +fn noSums(_: c.Shape) usize { + return 0; +} + +const ids_rule = "integers only: arena rows, slots and counts"; +const rows_rule = "row copies: every byte as stored, NaN and infinities included"; + +/// Rows a call x picks a row: decode and verify calls (16 rows or fewer) use the row cache, prompt chunks do not. +pub fn spec(which: Which) c.Spec { + const w = sparse.width; + const both: []const c.Shape = &.{ .{ .rows = 1, .n = w, .k = 1 }, .{ .rows = 8, .n = w, .k = 1 }, .{ .rows = 2048, .n = w, .k = 1 } }; + const chunk: []const c.Shape = &.{ .{ .rows = 64, .n = w, .k = 1 }, .{ .rows = 2048, .n = w, .k = 1 } }; + const cached: []const c.Shape = &.{ .{ .rows = 1, .n = w, .k = 1 }, .{ .rows = 8, .n = w, .k = 1 }, .{ .rows = 16, .n = w, .k = 1 } }; + const s: struct { []const u8, []const c.Shape, []const u8 } = switch (which) { + .mark => .{ "cuda glm tf_glm_kvh_mark: a slot for each distinct pick", both, "integers only: slots follow atomicAdd's order, so each pick is checked through its slot to its row" }, + .copy => .{ "cuda glm tf_glm_kvh_copy: marked rows into their slots", chunk, rows_rule }, + .remap => .{ "cuda glm tf_glm_kvh_remap: token lists name slots", both, ids_rule }, + .reset => .{ "cuda glm tf_glm_kvh_reset: marked rows unseen again", chunk, ids_rule }, + .rows => .{ "cuda glm tf_glm_kvh_rows: visible rows in order through pages out of order", both, rows_rule }, + .prev_mark => .{ "cuda glm tf_glm_kvh_prev_mark: the layer cache's rows, less this call's", cached, ids_rule }, + .copy_cached => .{ "cuda glm tf_glm_kvh_copy_cached: rows from the layer cache or host pages", cached, rows_rule }, + .keep => .{ "cuda glm tf_glm_kvh_keep: the call's rows become the layer cache", cached, rows_rule }, + .unmark => .{ "cuda glm tf_glm_kvh_unmark: both calls' marks cleared", cached, ids_rule }, + }; + return .{ .name = s[0], .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = s[2] }, .served = s[1] }; +} + +/// One gather kernel's case: a sequence at `at` with `rows` rows, its picks on pages out of order. +pub const Kvh = struct { + rig: rig_.Rig, + f: cuda.Function, + which: Which, + plant: bool = false, + short: bool = false, // mark only: half the slots its distinct picks need + shape: c.Shape = undefined, + ids: []i32 = &.{}, // the page table + tokens: []i32 = &.{}, // [rows][width] arena rows + counts: []i32 = &.{}, + distinct: []i32 = &.{}, // the picks' distinct rows, in slot order + slot_of: []i32 = &.{}, // per arena row: its slot, -1 unpicked + prev: []i32 = &.{}, // the layer cache's rows (prev_mark, unmark) or a cache slot per arena row (copy_cached) + lc: []u8 = &.{}, // the arena's rows + cache: []u8 = &.{}, // the layer cache's rows (copy_cached), the gathered slots (keep) + arena_rows: usize = 0, + unused: i32 = 0, // a row on a page the table skips: what the lists hold past their counts + in: [8]u64 = @splat(0), + state: [4]u64 = @splat(0), + out: [3]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Kvh) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Kvh) void { + const a = self.rig.gpu.gpa; + inline for (.{ &self.ids, &self.tokens, &self.counts, &self.distinct, &self.slot_of, &self.prev }) |x| { + a.free(x.*); + x.* = &.{}; + } + a.free(self.lc); + a.free(self.cache); + self.lc = &.{}; + self.cache = &.{}; + } + + pub fn kernel(self: *Kvh) g.Kernel { + var declared = spec(self.which); + if (self.short) declared.name = "cuda glm tf_glm_kvh_mark past its slots: flags the call, leaves no row marked"; + if (self.plant) declared.name = planted.get(self.which); + return .{ .spec = declared, .ctx = self, .vt = &vt, .control = if (self.plant) .shapes else null }; + } + + const vt = g.hooks(Kvh); + + fn width(self: *const Kvh) usize { + return self.shape.n; + } + + fn total(self: *const Kvh) usize { + return self.shape.rows * self.width(); + } + + fn visible(self: *const Kvh) usize { + return at + self.shape.rows; + } + + /// Position q's row in the arena: row q % 512 of page ids[q / 512]. + fn arena(self: *const Kvh, q: usize) i32 { + return self.ids[q / page] * page + @as(i32, @intCast(q % page)); + } + + fn picked(self: *const Kvh, i: usize) bool { + return i % self.width() < self.counts[i / self.width()]; + } + + fn cap(self: *const Kvh) usize { + return if (self.short) self.distinct.len / 2 else self.distinct.len; + } + + fn i32s(self: *Kvh, x: []const i32) !u64 { + return self.rig.input(std.mem.sliceAsBytes(x)); + } + + fn one(self: *Kvh, v: usize) !u64 { + const x: i32 = @intCast(v); + return self.rig.input(std.mem.asBytes(&x)); + } + + /// Per arena row: -1, or its index in `rows` where given. + fn marks(self: *Kvh, rows: []const i32) ![]i32 { + const m = try self.rig.gpu.gpa.alloc(i32, self.arena_rows); + @memset(m, -1); + for (rows, 0..) |t, i| m[@intCast(t)] = @intCast(i); + return m; + } + + pub fn prepare(self: *Kvh, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + self.free(); + self.rig.reset(); + self.shape = shape; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const n_pages = (self.visible() + page - 1) / page; + self.ids = try pages.shuffled(a, rng, n_pages); + self.arena_rows = (n_pages + 3) * page; + const skipped = for (0..n_pages + 3) |p| { + if (std.mem.indexOfScalar(i32, self.ids, @intCast(p)) == null) break p; + } else unreachable; + self.unused = @intCast(skipped * page + 7); + // the picks: each row's first counts[r] are rows it sees (row 1 sees none), the rest a row no call picks + self.counts = try a.alloc(i32, shape.rows); + for (self.counts, 0..) |*n, r| n.* = if (r == 1) 0 else @intCast(self.width() - r % 4); + self.tokens = try a.alloc(i32, self.total()); + for (self.tokens, 0..) |*t, i| t.* = if (self.picked(i)) self.arena(rng.uintAtMost(usize, at + i / self.width())) else self.unused; + var seen = try a.alloc(bool, self.arena_rows); + defer a.free(seen); + @memset(seen, false); + var list: std.ArrayList(i32) = .empty; + defer list.deinit(a); + for (self.tokens, 0..) |t, i| if (self.picked(i) and !seen[@intCast(t)]) { + seen[@intCast(t)] = true; + try list.append(a, t); + }; + rng.shuffle(i32, list.items); + self.distinct = try list.toOwnedSlice(a); + self.slot_of = try self.marks(self.distinct); + self.lc = try a.alloc(u8, self.arena_rows * row_bytes); + rng.bytes(self.lc); + if (fill == .nonfinite) for (self.distinct[0..@min(3, self.distinct.len)], [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }) |t, v| { + const at_row = @as(usize, @intCast(t)) * row_bytes; + std.mem.writeInt(u16, self.lc[at_row..][0..2], c.toBf16(v), .little); + std.mem.writeInt(u16, self.lc[at_row + row_bytes - 2 ..][0..2], c.toBf16(v), .little); + }; + try self.bind(rng); + return true; + } + + /// The kernel's inputs, state (reset from inputs each run) and outputs. + fn bind(self: *Kvh, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const d = self.distinct.len; + const n_in = if (self.plant and (self.which == .copy or self.which == .reset or self.which == .keep)) d - 1 else d; // a dropped row + switch (self.which) { + .mark => { + const clean = try self.marks(&.{}); + defer a.free(clean); + if (self.plant) clean[@intCast(self.distinct[0])] = -2; // a stale mark from an earlier call + self.in[0] = try self.i32s(self.tokens); + self.in[1] = try self.i32s(self.counts); + self.in[2] = try self.i32s(clean); + self.state = .{ try self.rig.scratch(self.arena_rows * 4), try self.rig.scratch(self.cap() * 4), try self.rig.scratch(16), try self.rig.scratch(16) }; + self.out[0] = try self.rig.output(self.viewLen() * 4); + }, + .copy => { + self.in[0] = try self.rig.input(self.lc); + self.in[1] = try self.i32s(self.distinct); + self.in[2] = try self.one(n_in); + self.out[0] = try self.rig.output(d * row_bytes); + }, + .remap => { + const marks_in = try a.dupe(i32, self.slot_of); + defer a.free(marks_in); + if (self.plant and d > 1) std.mem.swap(i32, &marks_in[@intCast(self.distinct[0])], &marks_in[@intCast(self.distinct[1])]); + self.in[0] = try self.i32s(self.tokens); + self.in[1] = try self.i32s(self.counts); + self.in[2] = try self.i32s(marks_in); + self.out[0] = try self.rig.output(self.total() * 4); + }, + .reset => { + self.in[0] = try self.i32s(self.distinct); + self.in[1] = try self.one(n_in); + self.in[2] = try self.i32s(self.slot_of); + self.out[0] = try self.rig.output(self.arena_rows * 4); + }, + .rows => { + self.in[0] = try self.rig.input(self.lc); + self.in[1] = try pages.uploadTable(&self.rig, self.ids, self.plant); + self.in[2] = try self.one(at); + self.out[0] = try self.rig.output(self.visible() * row_bytes); + }, + .prev_mark => { + // the cache: half the picks' rows and every row this call writes (a rejected draft's), shuffled + var rows: std.ArrayList(i32) = .empty; + defer rows.deinit(a); + for (self.distinct[0 .. d / 2]) |t| if (!self.written(t)) try rows.append(a, t); + for (at..self.visible()) |q| try rows.append(a, self.arena(q)); + rng.shuffle(i32, rows.items); + self.prev = try a.dupe(i32, rows.items[0..@min(rows.items.len, cache_rows)]); + const clean = try self.marks(&.{}); + defer a.free(clean); + self.in[0] = try self.i32s(self.prev); + self.in[1] = try self.one(self.prev.len); + self.in[2] = try pages.uploadTable(&self.rig, self.ids, false); + self.in[3] = try self.one(if (self.plant) at - self.shape.rows else at); // planted: this call's rows unseen + self.in[4] = try self.i32s(clean); + self.out[0] = try self.rig.output(self.arena_rows * 4); + }, + .copy_cached => { + // half the picks sit in the layer cache, at shuffled slots, with rows unlike their host rows + const held = d / 2; + const slots = try a.alloc(i32, held); + defer a.free(slots); + for (slots, 0..) |*s, i| s.* = @intCast(i); + rng.shuffle(i32, slots); + self.prev = try self.marks(&.{}); + for (self.distinct[0..held], slots) |t, s| self.prev[@intCast(t)] = s; + const prev_in = try a.dupe(i32, self.prev); + defer a.free(prev_in); + if (self.plant and held > 1) std.mem.swap(i32, &prev_in[@intCast(self.distinct[0])], &prev_in[@intCast(self.distinct[1])]); + self.cache = try a.alloc(u8, @max(held, 1) * row_bytes); + rng.bytes(self.cache); + self.in[0] = try self.rig.input(self.lc); + self.in[1] = try self.i32s(self.distinct); + self.in[2] = try self.one(d); + self.in[3] = try self.i32s(prev_in); + self.in[4] = try self.rig.input(self.cache); + self.out[0] = try self.rig.output(d * row_bytes); + }, + .keep => { + self.cache = try a.alloc(u8, d * row_bytes); // the call's gathered slots + rng.bytes(self.cache); + self.in[0] = try self.rig.input(self.cache); + self.in[1] = try self.i32s(self.distinct); + self.in[2] = try self.one(n_in); + const m: usize = @min(d, cache_rows); + self.out = .{ try self.rig.output(m * row_bytes), try self.rig.output(m * 4), try self.rig.output(4) }; + }, + .unmark => { + // the previous call's rows: a third of these picks and as many others + var rows: std.ArrayList(i32) = .empty; + defer rows.deinit(a); + for (self.distinct[0 .. d / 3]) |t| try rows.append(a, t); + for (0..d / 3) |i| try rows.append(a, self.arena(at - 1 - i % at)); + std.mem.sortUnstable(i32, rows.items, {}, std.sort.asc(i32)); + var k: usize = 0; + for (rows.items) |t| if (k == 0 or t != rows.items[k - 1]) { + rows.items[k] = t; + k += 1; + }; + rng.shuffle(i32, rows.items[0..k]); + self.prev = try a.dupe(i32, rows.items[0..@min(k, cache_rows)]); + const prev_marks = try self.marks(self.prev); + defer a.free(prev_marks); + self.in[0] = try self.i32s(self.prev); + self.in[1] = try self.one(if (self.plant) self.prev.len - 1 else self.prev.len); + self.in[2] = try self.i32s(prev_marks); + self.in[3] = try self.i32s(self.distinct); + self.in[4] = try self.one(d); + self.in[5] = try self.i32s(self.slot_of); + self.out = .{ try self.rig.output(self.arena_rows * 4), try self.rig.output(self.arena_rows * 4), 0 }; + }, + } + } + + /// Whether arena row t is one this call writes (positions at .. at + rows). + fn written(self: *const Kvh, t: i32) bool { + for (at..self.visible()) |q| if (self.arena(q) == t) return true; + return false; + } + + /// mark's checked view: each pick's row through its slot, n, the flag, -2 rows; past its slots, counts instead. + fn viewLen(self: *const Kvh) usize { + return if (self.short) 5 else self.total() + 3; + } + + fn launch(self: *Kvh, blocks: usize, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(@max(blocks, 1)) }, .block = .{ .x = 256 } }, self.rig.s, &a); + } + + fn copy(self: *Kvh, dst: u64, src: u64, len: usize) !void { + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = dst, .len = len }).copyFrom(0, src, len, self.rig.s.handle); + } + + /// As cuda_kv_host.zig launches each: 256 threads a block, a warp a row for the row copies. + pub fn run(self: *Kvh) !void { + const w: c_int = @intCast(self.width()); + const all: i64 = @intCast(self.total()); + const cap_: c_int = @intCast(self.cap()); + const upper = @min(self.total(), self.cap()); + const r: c_int = @intCast(self.shape.rows); + switch (self.which) { + .mark => { + try self.copy(self.state[0], self.in[2], self.arena_rows * 4); + for (self.state[2..4]) |p| try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = p, .len = 16 }).fill8(0, self.rig.s.handle); + try self.launch(cdiv(self.total(), 256), .{ self.in[0], self.in[1], w, all, self.state[0], self.state[1], self.state[2], cap_, self.state[3] }); + try self.view(); + }, + .copy => try self.launch(cdiv(upper, 8), .{ self.in[0], self.in[1], self.in[2], self.out[0], @as(c_int, rv), cap_ }), + .remap => { + try self.copy(self.out[0], self.in[0], self.total() * 4); + try self.launch(cdiv(self.total(), 256), .{ self.out[0], self.in[1], w, all, self.in[2] }); + }, + .reset => { + try self.copy(self.out[0], self.in[2], self.arena_rows * 4); + try self.launch(cdiv(upper, 256), .{ self.in[0], self.in[1], self.out[0], cap_ }); + }, + .rows => try self.launch(cdiv(self.visible(), 8), .{ self.in[0], self.in[1], @as(i64, page), self.in[2], r, self.out[0], @as(c_int, rv), @as(c_int, @intCast(self.visible())) }), + .prev_mark => { + try self.copy(self.out[0], self.in[4], self.arena_rows * 4); + try self.launch(cdiv(cache_rows, 256), .{ self.in[0], self.in[1], self.out[0], self.in[2], @as(i64, page), self.in[3], r }); + }, + .copy_cached => try self.launch(cdiv(upper, 8), .{ self.in[0], self.in[1], self.in[2], self.out[0], @as(c_int, rv), cap_, self.in[3], self.in[4] }), + .keep => try self.launch(cdiv(cache_rows, 8), .{ self.in[0], self.in[1], self.in[2], self.out[0], self.out[1], self.out[2], @as(c_int, rv), @as(c_int, cache_rows) }), + .unmark => { + try self.copy(self.out[0], self.in[2], self.arena_rows * 4); + try self.copy(self.out[1], self.in[5], self.arena_rows * 4); + try self.launch(cdiv(@max(upper, cache_rows), 256), .{ self.in[0], self.in[1], self.out[0], self.in[3], self.in[4], self.out[1], cap_ }); + }, + } + try self.rig.s.synchronize(); + } + + /// mark's state read back and reduced to the view its check compares (slot numbers themselves vary run to run). + fn view(self: *Kvh) !void { + const a = self.rig.gpu.gpa; + const marks_now = try a.dupe(u8, try self.rig.bytes(self.state[0], self.arena_rows * 4)); + defer a.free(marks_now); + const list = try a.dupe(u8, try self.rig.bytes(self.state[1], self.cap() * 4)); + defer a.free(list); + const head = try self.rig.bytes(self.state[2], 4); + const n = std.mem.readInt(i32, head[0..4], .little); + const flag_b = try self.rig.bytes(self.state[3], 4); + const flag = std.mem.readInt(i32, flag_b[0..4], .little); + const slot = std.mem.bytesAsSlice(i32, marks_now); + const rows = std.mem.bytesAsSlice(i32, list); + var stale: i32 = 0; + for (slot) |s| stale += @intFromBool(s == -2); + const out = try a.alloc(i32, self.viewLen()); + defer a.free(out); + if (self.short) { + var resolved: i32 = 0; + var wrong: i32 = 0; + for (slot, 0..) |s, t| if (s >= 0) { + resolved += 1; + wrong += @intFromBool(s >= rows.len or rows[@intCast(s)] != @as(i32, @intCast(t))); + }; + const most: i32 = @intCast(self.distinct.len); // a row past the slots is unseen again: its later picks count again + out[0..5].* = .{ resolved, wrong, @min(n, most), flag, stale }; + } else { + for (self.tokens, 0..) |t, i| { + const s = if (self.picked(i)) slot[@intCast(t)] else -1; + out[i] = if (s >= 0 and s < rows.len) rows[@intCast(s)] else -1; + } + out[self.total()..][0..3].* = .{ n, flag, stale }; + } + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = self.out[0], .len = out.len * 4 }).upload(0, std.mem.sliceAsBytes(out)); + } + + pub fn count(self: *Kvh) usize { + return switch (self.which) { + .mark => self.viewLen() * 4, + .copy, .copy_cached => self.distinct.len * row_bytes, + .remap => self.total() * 4, + .reset, .prev_mark => self.arena_rows * 4, + .rows => self.visible() * row_bytes, + .keep => @as(usize, @min(self.distinct.len, cache_rows)) * (row_bytes + 4) + 4, + .unmark => self.arena_rows * 8, + }; + } + + pub fn reference(_: *Kvh, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + fn got(self: *Kvh, ptr: u64, len: usize) ![]u8 { + return self.rig.gpu.gpa.dupe(u8, try self.rig.bytes(ptr, len)); + } + + /// Ids `v` against `want`: null when equal, else the first difference named. + fn ids_(self: *Kvh, what: []const u8, v: []const u8, want: []const i32) !?[]const u8 { + for (want, 0..) |w, i| { + const x = std.mem.readInt(i32, v[i * 4 ..][0..4], .little); + if (x != w) return try std.fmt.bufPrint(&self.msg, "{s} {d}: {d} where {d} is due", .{ what, i, x, w }); + } + return null; + } + + /// Row `s` of `v` against `src`'s row: null when the bytes match. + fn row(self: *Kvh, what: []const u8, v: []const u8, s: usize, src: []const u8, from: usize) !?[]const u8 { + if (std.mem.eql(u8, v[s * row_bytes ..][0..row_bytes], src[from * row_bytes ..][0..row_bytes])) return null; + return try std.fmt.bufPrint(&self.msg, "{s} slot {d}: not the bytes of row {d}", .{ what, s, from }); + } + + pub fn check(self: *Kvh, _: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const d = self.distinct.len; + switch (self.which) { + .mark => { + const v = try self.got(self.out[0], self.viewLen() * 4); + defer a.free(v); + const want = try a.alloc(i32, self.viewLen()); + defer a.free(want); + if (self.short) { + want[0..5].* = .{ @intCast(self.cap()), 0, @intCast(d), 1, 0 }; + return self.ids_("past its slots: (resolved, misresolved, n, flag, rows left at -2)", v, want); + } + for (self.tokens, 0..) |t, i| want[i] = if (self.picked(i)) t else -1; + want[self.total()..][0..3].* = .{ @intCast(d), 0, 0 }; + return self.ids_("pick (then n, flag, rows left at -2)", v, want); + }, + .copy, .copy_cached => { + const v = try self.got(self.out[0], d * row_bytes); + defer a.free(v); + for (self.distinct, 0..) |t, s| { + const p = if (self.which == .copy_cached) self.prev[@intCast(t)] else -1; + const why = if (p >= 0) try self.row("from the cache", v, s, self.cache, @intCast(p)) else try self.row("from host pages", v, s, self.lc, @intCast(t)); + if (why != null) return why; + } + return null; + }, + .remap => { + const v = try self.got(self.out[0], self.total() * 4); + defer a.free(v); + const want = try a.alloc(i32, self.total()); + defer a.free(want); + for (self.tokens, want, 0..) |t, *w, i| w.* = if (self.picked(i)) self.slot_of[@intCast(t)] else t; + return self.ids_("token", v, want); + }, + .reset, .prev_mark => { + const v = try self.got(self.out[0], self.arena_rows * 4); + defer a.free(v); + const want = try self.marks(&.{}); + defer a.free(want); + if (self.which == .prev_mark) for (self.prev, 0..) |t, i| if (!self.written(t)) { + want[@intCast(t)] = @intCast(i); + }; + return self.ids_("arena row", v, want); + }, + .rows => { + const v = try self.got(self.out[0], self.visible() * row_bytes); + defer a.free(v); + for (0..self.visible()) |q| if (try self.row("visible", v, q, self.lc, @intCast(self.arena(q)))) |why| return why; + return null; + }, + .keep => { + const m: usize = @min(d, cache_rows); + const v = try self.got(self.out[0], m * row_bytes); + defer a.free(v); + for (0..m) |s| if (try self.row("cache", v, s, self.cache, s)) |why| return why; + const kept = try self.got(self.out[1], m * 4); + defer a.free(kept); + if (try self.ids_("cache row", kept, self.distinct[0..m])) |why| return why; + const n = try self.got(self.out[2], 4); + defer a.free(n); + return self.ids_("cache count", n, &.{@intCast(m)}); + }, + .unmark => { + const want = try self.marks(&.{}); + defer a.free(want); + for (self.out[0..2], [_][]const u8{ "cache mark of arena row", "slot mark of arena row" }) |p, what| { + const v = try self.got(p, self.arena_rows * 4); + defer a.free(v); + if (try self.ids_(what, v, want)) |why| return why; + } + return null; + }, + } + } +}; + +fn cdiv(x: usize, y: usize) usize { + return (x + y - 1) / y; +} + +/// The nine kernels, mark past its slots, and a planted control each, at stable addresses. +pub const Set = struct { + cases: [2 * std.enums.values(Which).len + 1]Kvh, + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, ops: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !*Set { + const self = try a.create(Set); + errdefer a.destroy(self); + const all = std.enums.values(Which); + for (all, 0..) |w, i| { + self.cases[i] = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try ops.function(symbol(w)), .which = w }; + self.cases[all.len + i] = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try ops.function(symbol(w)), .which = w, .plant = true }; + } + self.cases[2 * all.len] = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try ops.function(symbol(.mark)), .which = .mark, .short = true }; + for (&self.cases) |*k| try kernels.append(a, k.kernel()); + return self; + } + + pub fn destroy(self: *Set, a: std.mem.Allocator) void { + for (&self.cases) |*k| k.deinit(); + a.destroy(self); + } +}; + +fn symbol(w: Which) [:0]const u8 { + return switch (w) { + inline else => |x| "tf_glm_kvh_" ++ @tagName(x), + }; +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index ca9742e59f..92c1d691a6 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -554,7 +554,7 @@ pub fn run(gpu: check.Gpu) !void { x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = if (hb == 16) glm_fns.chunks_16 else glm_fns.chunks_32, .merge = glm_fns.merge, .form = .{ .hb = hb, .sparse = sparse } }; defer for (&glm_chunks) |*x| x.deinit(); for (&glm_chunks) |*x| try kernels.append(a, x.kernel()); - var glm_controls = @import("contract_glm_controls.zig").Controls.init(gpu, s, glm_fns, glm_tokens.f); + var glm_controls = @import("contract_glm_controls.zig").Controls.init(gpu, s, glm_fns, glm_tokens.f, glm_ops); defer glm_controls.deinit(); try glm_controls.append(a, &kernels); var q35_flat_mod = try cuda.Module.load(gpu.d, kk.qwen35); From 6da7ba6506af75e6afb8031d17923f9741344c8d Mon Sep 17 00:00:00 2001 From: Russell Marr <35648773+russellmm@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:00:59 -0800 Subject: [PATCH 6/7] glm cuda: a TF_GLM_KV_HOST call past its slots is refused by name, with a test Commit 5's evidence, here since the branch is under review and not force-pushed. The nine tf_glm_kvh_* kernels and their planted controls in the CUDA kernel contract (RTX PRO 6000, sm_120, nvcc 13.4; the kernels and limit records narrowed to tf_glm_kvh_* by a local name filter): SUMMARY 10 kernels: 10 pass gates 1-4, 0 fail, 0 known findings; 9 controls: 9 red as planted, 0 missed; 0 limit problems; gates 405 ms, limits 0 ms The refusal: Engine.checkGathers now calls kv_host.checkActive(), and `tf-cuda-test glm-kv-host-slots` (one GPU) runs the real Gather.gather with 256 distinct rows for 64 slots: it is refused as error.GlmKvHostSlots, no row is left marked, and the next call, 48 distinct rows, is accepted with each token's slot holding its row. The flag's read-and-clear is a host test, so CI runs that part. --- zig/src/families/glm/cuda.zig | 1 + zig/src/families/glm/cuda_engine.zig | 2 +- zig/src/families/glm/cuda_kv_host.zig | 25 +++++++- zig/tests/cuda/glm_kv_host_slots.zig | 86 +++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 5 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 zig/tests/cuda/glm_kv_host_slots.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 877c38e354..52859f5346 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -38,6 +38,7 @@ pub const follow = @import("cuda_follow.zig"); pub const lanes = @import("cuda_lanes.zig"); pub const native = @import("cuda_native.zig"); pub const snapshots = @import("cuda_snapshots.zig"); +pub const kv_host = @import("cuda_kv_host.zig"); test { _ = @import("cuda_tile.zig"); diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 774acc78fc..88e1e31a71 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -576,7 +576,7 @@ pub const Engine = struct { /// A failed RoCE gather (its rows hold nothing) or TF_GLM_KV_HOST gather past its slots, after a synchronize. pub fn checkGathers(e: *const Engine) !void { if (e.group) |g| try g.check(); - if (@import("cuda_kv_host.zig").active) |*g| try g.check(); // a TF_GLM_KV_HOST gather past its slots + try @import("cuda_kv_host.zig").checkActive(); // a TF_GLM_KV_HOST gather past its slots } pub fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void { diff --git a/zig/src/families/glm/cuda_kv_host.zig b/zig/src/families/glm/cuda_kv_host.zig index d06b6cd28a..464405dc85 100644 --- a/zig/src/families/glm/cuda_kv_host.zig +++ b/zig/src/families/glm/cuda_kv_host.zig @@ -44,6 +44,18 @@ pub fn prepare(d: *const cuda.Driver, page: u64, capacity: usize, row_bytes: usi /// The engine's gather, set when its pages go to host memory. pub var active: ?Gather = null; +/// Engine.checkGathers' part: the engine's gather refuses a call past its slots by name (none: nothing to check). +pub fn checkActive() !void { + if (active) |*g| try g.check(); +} + +/// Whether mark flagged a call past the slots since the last look; the flag is cleared for the next. +fn taken(flag: *volatile i32) bool { + if (flag.* == 0) return false; + flag.* = 0; + return true; +} + /// Rows a KV page holds; the row cache's layers (11 DSA and MTP, with room), rows a layer, and its largest call. const page_rows = 512; const layers = 16; @@ -100,9 +112,7 @@ pub const Gather = struct { /// After a synchronize: refuses (and clears) a call since the last check whose distinct picks passed the slots. pub fn check(g: *const Gather) !void { - const flag: *volatile i32 = @ptrCast(g.over.slice(i32).ptr); - if (flag.* == 0) return; - flag.* = 0; + if (!taken(@ptrCast(g.over.slice(i32).ptr))) return; std.log.err("kv host: a call picked more distinct rows than the gather's {d} slots", .{g.cap}); return error.GlmKvHostSlots; } @@ -263,6 +273,15 @@ test "kv host pages: one full window and 16 by default, TF_GLM_KV_HOST_GIB's wor try std.testing.expectEqual(@as(usize, 0), pagesOf(0, page, 1 << 20)); } +test "kv host slots: mark's flag is seen once, then cleared for the next call" { + var flag: i32 = 0; + try std.testing.expect(!taken(&flag)); + flag = 1; // mark: a call picked more distinct rows than the slots + try std.testing.expect(taken(&flag)); + try std.testing.expectEqual(@as(i32, 0), flag); + try std.testing.expect(!taken(&flag)); +} + test "kv host shared pages: the last of the ranks to map the file takes its name away" { var count: u32 = 0; try std.testing.expect(!lastToMap(&count, 2)); diff --git a/zig/tests/cuda/glm_kv_host_slots.zig b/zig/tests/cuda/glm_kv_host_slots.zig new file mode 100644 index 0000000000..8f78411071 --- /dev/null +++ b/zig/tests/cuda/glm_kv_host_slots.zig @@ -0,0 +1,86 @@ +//! glm-kv-host-slots (GPU): TF_GLM_KV_HOST's gather, past its slots and within them, through the engine's check. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const kv_host = @import("glm_cuda").kv_host; + +const arena_rows = 1024; +const row_bytes = 64; +const cap = 64; // slots: the first call picks 256 distinct rows, the second 48 +const rows = 32; // a call past the row cache's 16, so mark, copy, remap and reset +const width = 8; + +/// Arena row `r`'s bytes: words naming the row and the word, so a row in the wrong slot shows. +fn word(r: usize, i: usize) u32 { + return @intCast(r * 1000 + i); +} + +/// One gather of `picks` (rows x width, every count full); returns the token lists it remapped to slots. +fn call(gpu: check.Gpu, ops: cuda.Module, s: cuda.Stream, lc: u64, picks: []const i32) ![]i32 { + var tokens = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer tokens.free(); + const full: [rows]i32 = @splat(width); + var counts = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&full)); + defer counts.free(); + try gpu.d.synchronize(); // the uploads, before the non-blocking stream + _ = try kv_host.active.?.gather(ops, s, 1, lc, 0, 0, 0, tokens.ptr, counts.ptr, rows, width); + try s.synchronize(); + const out = try gpu.gpa.alloc(i32, picks.len); + try tokens.download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// No arena row left marked: every slot_of entry back at -1. +fn unmarked(gpu: check.Gpu) !bool { + const bytes = try check.download(gpu, kv_host.active.?.slot_of); + defer gpu.gpa.free(bytes); + return std.mem.allEqual(u8, bytes, 0xff); +} + +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var ops = try cuda.Module.load(gpu.d, cuda.kernels.glm_ops); + defer ops.unload(); + var host: [arena_rows * row_bytes / 4]u32 = undefined; + for (&host, 0..) |*w, i| w.* = word(i / (row_bytes / 4), i % (row_bytes / 4)); + var lc = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&host)); + defer lc.free(); + kv_host.active = try kv_host.Gather.init(gpu.d, arena_rows, cap, row_bytes); + defer { + kv_host.active.?.deinit(); + kv_host.active = null; + } + + var past: [rows * width]i32 = undefined; + for (&past, 0..) |*t, i| t.* = @intCast(i * 37 % arena_rows); // 256 distinct rows + gpu.gpa.free(try call(gpu, ops, s, lc.ptr, &past)); + kv_host.checkActive() catch |err| { + try check.expect(err == error.GlmKvHostSlots, "kv host slots: past the slots gave {s}", .{@errorName(err)}); + try check.expect(try unmarked(gpu), "kv host slots: a row past the slots was left marked", .{}); + try kv_host.checkActive(); // the flag is cleared: the next call is judged alone + check.pass("kv host slots: 256 distinct rows for {d} slots refused as error.{s}, no row left marked", .{ cap, @errorName(err) }); + return within(gpu, ops, s, lc.ptr, &host); + }; + try check.expect(false, "kv host slots: 256 distinct rows for {d} slots were not refused", .{cap}); +} + +/// A call within the slots after the refused one: accepted, and each token's slot holds its row's bytes. +fn within(gpu: check.Gpu, ops: cuda.Module, s: cuda.Stream, lc: u64, host: []const u32) !void { + var picks: [rows * width]i32 = undefined; + for (&picks, 0..) |*t, i| t.* = @intCast(i % 48 * 5); // 48 distinct rows, each picked more than once + const slots_of = try call(gpu, ops, s, lc, &picks); + defer gpu.gpa.free(slots_of); + try kv_host.checkActive(); + try check.expect(try unmarked(gpu), "kv host slots: a call within the slots left a row marked", .{}); + const slots = try check.download(gpu, kv_host.active.?.slots); + defer gpu.gpa.free(slots); + const arena = std.mem.sliceAsBytes(host); + for (picks, slots_of) |r, slot| { + try check.expect(slot >= 0 and slot < 48, "kv host slots: row {d} got slot {d}", .{ r, slot }); + const at: usize = @intCast(slot); + const row: usize = @intCast(r); + try check.sameBytes("kv host slots: a slot's row", slots[at * row_bytes ..][0..row_bytes], arena[row * row_bytes ..][0..row_bytes]); + } + check.pass("kv host slots: then 48 distinct rows for {d} slots accepted, each token's slot holding its row", .{cap}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 17c095a35c..fd2376f18a 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -69,6 +69,7 @@ const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_split_tests = @import("glm_split_tests.zig"); const glm_solo_tests = @import("glm_solo_tests.zig"); const glm_paged_tests = @import("glm_paged_tests.zig"); +const glm_kv_host_slots = @import("glm_kv_host_slots.zig"); const glm_synth = @import("glm_synth.zig"); const tree_round_tests = @import("tree_round_tests.zig"); const tree_tests = @import("tree_tests.zig"); @@ -252,6 +253,7 @@ const usage = \\ glm-mm-fixed Prompt K policies: FP64, every row alone/batched, offsets and guarded tails \\ glm-mm-bench GLM's bf16 linears: b16.matmul against glm_mm, by rows (one GPU) \\ glm-index-check GLM's indexer with glm_index against tf_glm_scores and tf_glm_select_rows, byte for byte (one GPU) + \\ glm-kv-host-slots TF_GLM_KV_HOST's gather past its slots is refused as error.GlmKvHostSlots, then one within (one GPU) \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run @@ -394,6 +396,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-mm-fixed")) return glm_mm_fixed.run(gpu); if (std.mem.eql(u8, cmd, "glm-mm-bench")) return glm_mm_bench.run(gpu); if (std.mem.eql(u8, cmd, "glm-index-check")) return glm_index_check.run(gpu); + if (std.mem.eql(u8, cmd, "glm-kv-host-slots")) return glm_kv_host_slots.run(gpu); if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "moe")) return moe_tests.run(gpu); From adc001b4bdf9fcdde38647ae5a46fbafd3e829f3 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:42:33 +0000 Subject: [PATCH 7/7] changelog: #1795 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3eb3c2ea39..4bd5a3bc51 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 on CUDA can keep the KV pages' latent rows in pinned host memory with `TF_GLM_KV_HOST`, for a 1M-token window on two 96 GB cards; unset, the engine is as before, with identical tokens at TP=2 on two GB10s (#1795, part of #1794). Thanks to @russellmm. + - On HIP, the planned decode score takes four keys a warp on long spans and two on short ones, with the same lane order and folds, so decode runs 7.5-18.9% faster on Qwen3.6-35B-A3B and Qwen3.8-27B at 32K and 64K with identical tokens on a gfx1151 (#1932). Thanks to @AbrahamCete. - Ling's grouped expert routing runs on Metal, checked against a float64 reference: sigmoid scores plus the routing bias, groups picked by their top scores, empty picks handled and both infinities in a group refused (#1836, part of #1247). Thanks to @ThinkOffApp.