glm cuda: TF_GLM_KV_HOST keeps the KV pages' latent rows in pinned host memory

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- GLM 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.

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.

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

zig/kernels/cuda/glm_ops.cu

@@ -53,3 +53,105 @@ 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,
+                                           volatile int* over) {
+    const long long i = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i >= total) return;
+    if (static_cast<int>(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;
+    } else {
+        slot_of[t] = -1;  // past the slots: left unseen, and the call flagged in host memory (Gather.check reads it)
+        *over = 1;
+    }
+}
+
+// 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<long long>(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<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i >= total) return;
+    if (static_cast<int>(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<long long>(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<long long>(blockIdx.x) * (blockDim.x / 32) + threadIdx.x / 32;
+    if (s >= min(static_cast<long long>(*pos) + rows, static_cast<long long>(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<long long>(*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<long long>(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<long long>(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<long long>(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<long long>(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;
+}

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.

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

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

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