glm cuda: TF_GLM_CHUNK_SOLO takes two grid chunks in one pass while no other stream decodes

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- GLM-5.3-Flash on CUDA: `TF_GLM_CHUNK_SOLO=4096` (opt-in, on every rank) lets a prompt pass take 4,096 rows while no other stream is decoding, with the same bytes. On GB10s a cold prompt's first token comes 6% to 9% sooner from 8K to 128K tokens. It holds 1.6 to 1.9 GiB more for prompt buffers, and a request that arrives during a long pass waits for it to end. Unset, nothing changes (#1823, option C of #1554). Thanks to @Bizuayeu.
+
 - Qwen3.8-27B on M5: checkpoints that mix bit widths in a fused projection (such as the oQ4 and oQ2 quantizations) answer prompts longer than 128 tokens, instead of failing with `BadCoreProjectionRows`; each 128-row piece is the call main already made, so replies keep main's tokens (#1822). Thanks to @grantoverton.
 
 - Windows: the cost cache reads a table through a shared file handle, so a reader that meets a table being replaced sees a miss quietly instead of logging an unexpected `FILE_DELETED` status, which failed a host test (fixes #1805).

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

@@ -49,6 +49,8 @@ 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_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).
 
 ## More ranks
 
@@ -58,6 +60,36 @@ left after a 12.2 GiB reserve` and a 5.8 GiB prompt cache
 (`TF_GLM_PREFILL_REDUCE`) and the two-rank RoCE gathers (`TF_ROCE`) are in
 [placement](../multi-rank-receipts.md#the-settings-the-ranks-agreed-on).
 
+## Prompt chunk size
+
+A prompt runs in passes of 2,048 rows. Kept prompt states sit on that grid, and a stream decoding beside a prompt waits
+for each pass to end. A pass that takes more rows reads each routed expert's weights fewer times, so a lone prompt
+comes sooner. `TF_GLM_CHUNK_SOLO=4096` lets a pass take two of those chunks when no other stream is decoding as it
+starts. With one decoding, or with the DFlash2 drafter, `prompt_logprobs` or image rows, the pass takes 2,048.
+
+The grid does not move, so warm turns resume where they do unset, and streams decoding beside a prompt see the same
+passes. Replies are the same bytes either way. Two costs remain:
+- A request that arrives during a solo pass waits for that pass to end before its own prompt runs. After that, the
+  passes are short again.
+- The prompt buffers are sized for 4,096 rows for the whole run: 1.92 GiB more at TP=2 and 1.56 GiB at TP=3, which
+  comes out of the room for KV pages and the prompt cache. In the runs below, TP=2's default cache goes from 4.7 to
+  2.8 GiB.
+
+| | TP=2 unset | TP=2 `=4096` | TP=3 unset | TP=3 `=4096` |
+| --- | --- | --- | --- | --- |
+| First token of a cold prompt, 8K / 32K / 128K (s) | 6.08 / 23.94 / 98.16 | 5.72 / 22.22 / 91.33 | 5.00 / 19.61 / 80.26 | 4.64 / 17.99 / 73.45 |
+| First token of a warm turn, 10K history + 500 / 1,500 / 3,000 tokens (s) | 1.40 / 2.06 / 2.87 | 1.40 / 2.05 / 2.83 | 1.14 / 1.67 / 2.35 | 1.14 / 1.68 / 2.31 |
+| First token of a warm turn, 40K history + 500 / 1,500 / 3,000 tokens (s) | 1.25 / 2.03 / 3.05 | 1.26 / 2.04 / 3.03 | 1.02 / 1.67 / 2.51 | 1.02 / 1.66 / 2.48 |
+| Gap between a decoding stream's tokens beside a 32K prompt, median / longest (ms) | 1,474 / 1,749 | 1,478 / 1,751 | 1,219 / 1,444 | 1,219 / 1,459 |
+| That stream's tokens during the prompt | 17.0 | 16.5 | 16.0 | 16.5 |
+| First token of a short request sent 6 to 8 s into a 32K prompt, mean (lowest to highest) (s) | 1.27 (0.88 to 1.65) | 2.48 (1.35 to 3.47) | 1.24 (0.70 to 1.67) | 1.73 (1.07 to 2.67) |
+| That prompt's first token (s) | 23.54 | 23.27 | 19.32 | 18.99 |
+
+Each setting is two fresh servers (one at TP=2 for the cold row), `--parallel 4 --context 524288`, `TF_KV_PAGES=on`,
+`Bizuayeu/GLM-5.3-Flash-XSFP4` with its MTP drafts. Cold prompts are Python source at `--prompt-cache-gib 0`, and a
+repeated filler line gains 2 to 5% at the same sizes. The other rows ran with the prompt cache on (the default at TP=2,
+2 GiB at TP=3). Three rows at TP=2 that waited on the heat guard are left out.
+
 ## Long runs on GB10
 
 Long prompts and long runs of replies heat a GB10. The runbook's [heat bands](../../RUNBOOK.md#linux-and-cuda) (with a

zig/src/families/glm/cuda.zig

@@ -28,6 +28,7 @@ pub const q4 = @import("cuda_q4.zig");
 pub const ranks = @import("cuda_ranks.zig");
 pub const sampler = @import("cuda_sampler.zig");
 pub const prompt_sink = @import("cuda_prompt_sink.zig");
+pub const prompt_solo = @import("cuda_prompt_solo.zig");
 pub const prompt_logprobs = @import("cuda_prompt_logprobs.zig");
 pub const engine = @import("cuda_engine.zig");
 pub const pages = @import("cuda_pages.zig");
@@ -49,6 +50,7 @@ test {
     _ = @import("cuda_ranks_test.zig");
     _ = @import("cuda_roce_multi.zig");
     _ = @import("cuda_lanes_test.zig");
+    _ = prompt_solo;
     _ = experts;
     _ = exl3;
     _ = vision;

zig/src/families/glm/cuda_engine.zig

@@ -30,6 +30,7 @@ const pair_bytes = 16;
 
 /// A prompt pass's chunk rows (decode.prefill's).
 pub const prompt_rows = fw.prefill_rows;
+pub const solo_rows = fw.solo_rows;
 pub const max_together = @import("cuda_state.zig").prompt_heads;
 pub const Prompt = struct { slot: u32, ids: []const u32 };
 
@@ -39,6 +40,7 @@ pub const max_shared = 16;
 
 const types = @import("cuda_engine_types.zig");
 const snapshots = @import("cuda_snapshots.zig");
+const prompt_solo = @import("cuda_prompt_solo.zig");
 pub const Part = types.Part;
 pub const Kept = types.Kept;
 const Seq = types.Seq;
@@ -175,8 +177,9 @@ pub const Engine = struct {
             e.free(0);
             e.seqs.deinit(gpa);
         }
-        e.bp = try fw.Buffers.init(d, &e.w, fw.prefill_rows, true, o.capacity);
+        e.bp = try fw.Buffers.init(d, &e.w, prompt_solo.bufferRows(o.settings.solo), true, o.capacity);
         errdefer e.bp.deinit();
+        if (o.settings.solo and rank == 0) std.log.info("glm: prompt passes of {d} rows, {d} while another stream decodes", .{ fw.solo_rows, fw.prefill_rows });
         e.bd = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity);
         errdefer e.bd.deinit();
         e.bm = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity);
@@ -338,8 +341,9 @@ pub const Engine = struct {
         e.index_split = false;
         if ((from != null or start == 0) and try e.prefillStart(len, from) != start) return error.BadChunk;
         if (start >= len or start % fw.prefill_rows != 0) return error.BadChunk;
-        const rows = @min(len - start, fw.prefill_rows);
+        const rows = @min(len - start, @as(usize, if (part.len > fw.prefill_rows + 1) fw.solo_rows else fw.prefill_rows));
         if (part.len != @min(len - start, rows + 1)) return error.BadChunk;
+        if (rows > fw.prefill_rows and !prompt_solo.allowed(e, start, rows)) return error.BadChunk;
         var one: [1]fw.Rank = undefined;
         var f = e.view(&e.bp, &one);
         f.index_split = split;
@@ -383,17 +387,7 @@ pub const Engine = struct {
         if (f.prof) |p| p.mark(e.stream, .draw);
         if (last) try e.holdPrompt(e.bp.fnormed + (rows - 1) * e.c.hidden * 2);
         stamp = prof_mod.span(f.prof, e.stream, .mtp);
-        if (e.mtp) {
-            const nxt = as_i32[1..];
-            if (nxt.len > 0) { // decode._absorb_rows
-                const mp = f.prof;
-                defer f.prof = mp;
-                f.prof = null; // the head's pass as one section
-                try f.mtp(nxt, &.{e.bp.fnormed});
-                f.prof = mp;
-                try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len);
-            }
-        }
+        if (e.mtp) try prompt_solo.mtpRows(e, &f, as_i32);
         stamp.end();
         if (f.prof) |p| p.mark(e.stream, .mtp);
         stamp = prof_mod.span(f.prof, e.stream, .commit);

zig/src/families/glm/cuda_follow.zig

@@ -3,6 +3,7 @@
 const std = @import("std");
 const cuda = @import("cuda");
 const engine = @import("cuda_engine.zig");
+const prompt_solo = @import("cuda_prompt_solo.zig");
 const images_mod = @import("cuda_images.zig");
 const vcfg = @import("vision_config.zig");
 const sampler = @import("cuda_sampler.zig");
@@ -120,6 +121,9 @@ pub const Link = struct {
     }
 };
 
+/// A prompt step's outcome (Leader.chunk): the first token once it ends the prompt, and the rows it took.
+pub const Chunk = struct { first: ?u32, rows: usize };
+
 /// Rank 0's engine: every step goes to the other ranks first.
 pub const Leader = struct {
     e: *Engine,
@@ -219,15 +223,17 @@ pub const Leader = struct {
         return x.e.prefill(ids, from, cancel, marks);
     }
 
-    /// One chunk of a prompt pass on every rank (Engine.prefillChunk); the first token once the chunk ends the prompt.
-    pub fn chunk(x: Leader, ids: []const u32, start: usize, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 {
+    /// One prompt step on every rank (Engine.prefillChunk): a chunk, or a solo pass if `solo` and prompt_solo.at allow.
+    pub fn chunk(x: Leader, ids: []const u32, start: usize, solo: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk {
         if (start >= ids.len) return error.BadChunk;
         try x.flush();
-        const part = ids[start..@min(ids.len, start + engine.prompt_rows + 1)];
+        const n: usize = if (solo and prompt_solo.at(x.e, start, ids.len, from, marks)) engine.solo_rows else engine.prompt_rows;
+        const part = ids[start..@min(ids.len, start + n + 1)];
         const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world);
         try x.send(.{ .kind = @backingInt(Kind.chunk), .flags = x.rows() | @as(u32, @intFromBool(split)) << 1, .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0, .depth = @intCast(start), .top_k = @intCast(ids.len) }, part);
         x.e.index_split = split;
-        return x.e.prefillChunk(part, start, ids.len, from, cancel, marks);
+        const first = try x.e.prefillChunk(part, start, ids.len, from, cancel, marks);
+        return .{ .first = first, .rows = @min(ids.len - start, n) };
     }
 
     /// `bits`: the rows' grammar words (structured output, with Engine.structured), else 0.

zig/src/families/glm/cuda_forward.zig

@@ -27,6 +27,7 @@ const shared = @import("cuda_shared.zig");
 const index_mod = @import("cuda_index.zig");
 
 pub const prefill_rows = state.prefill_rows;
+pub const solo_rows = state.solo_rows;
 pub const window_rows = state.window_rows;
 pub const State = state.State;
 pub const Buffers = state.Buffers;

zig/src/families/glm/cuda_lanes.zig

@@ -248,13 +248,13 @@ pub const Lanes = struct {
         self.x.e.prompt_logprobs = sink.want();
         defer self.x.e.prompt_logprobs = null;
         var keeps: prefix_mod.Keeps = undefined;
-        if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, self.marksOf(s, &keeps))) |first| {
+        const step = try self.x.chunk(s.prompt(), at, !@import("cuda_prompt_solo.zig").othersDecode(&self.lanes, s), from, .{ .ptr = s, .check = cancelled }, self.marksOf(s, &keeps));
+        if (step.first) |first| {
             l.first = self.take(first);
             l.fill = null;
             return true;
         }
-        // its chunks sit on the 2048 grid (cuda_engine.zig): l.fill stays uncapped; the share still works by time.
-        l.fill = at + engine.prompt_rows;
+        l.fill = at + step.rows; // on the 2048 grid (cuda_engine.zig), so uncapped; the share still works by time
         if (s.settle_step) try self.x.e.stream.synchronize(); // its tail would land in the next decode round's time
         return false;
     }
@@ -266,7 +266,7 @@ pub const Lanes = struct {
         if (at <= s.cached or s.reuse.hook == null) return;
         try self.x.bind(l.slot);
         var keeps: prefix_mod.Keeps = undefined;
-        _ = self.x.chunk(s.prompt(), at, null, .{ .ptr = null, .check = cancelNow }, self.marksOf(s, &keeps)) catch |err| {
+        _ = self.x.chunk(s.prompt(), at, false, null, .{ .ptr = null, .check = cancelNow }, self.marksOf(s, &keeps)) catch |err| {
             if (err == error.Cancelled) return;
             return err;
         };

zig/src/families/glm/cuda_lanes_test.zig

@@ -3,6 +3,8 @@ const std = @import("std");
 const lanes = @import("lanes");
 const engine = @import("cuda_engine.zig");
 const Lanes = @import("cuda_lanes.zig").Lanes;
+const fw = @import("cuda_forward.zig");
+const solo = @import("cuda_prompt_solo.zig");
 
 test "a lane's slot is chosen before its map entry exists" {
     var storage: [8192]u8 = @splat(0);
@@ -40,3 +42,21 @@ test "a failed lane insertion leaves no entry, and a retry takes slot zero" {
     fail.fail_index = std.math.maxInt(usize);
     try std.testing.expectEqual(@as(u32, 0), (try self.claim(&stream)).slot);
 }
+
+test "a prompt pass goes solo only while no other stream is past its prompt pass" {
+    var e: engine.Engine = undefined;
+    e.gpa = std.testing.allocator;
+    var self: Lanes = .{ .x = .{ .e = &e, .link = undefined } };
+    defer self.lanes.deinit(e.gpa);
+    var streams: [3]lanes.Stream = undefined;
+    const first = try self.claim(&streams[0]); // a stream's own lane, fresh as its first chunk starts
+    try std.testing.expect(!solo.othersDecode(&self.lanes, &streams[0]));
+    first.fill = fw.prefill_rows;
+    (try self.claim(&streams[1])).fill = 2 * fw.prefill_rows; // another prompt pass, partway
+    try std.testing.expect(!solo.othersDecode(&self.lanes, &streams[0]));
+    _ = try self.claim(&streams[2]); // past its prompt pass: decoding
+    try std.testing.expect(solo.othersDecode(&self.lanes, &streams[0]) and solo.othersDecode(&self.lanes, &streams[1]));
+    try std.testing.expect(!solo.othersDecode(&self.lanes, &streams[2])); // the other two still fill
+    _ = self.lanes.remove(&streams[2]); // released
+    try std.testing.expect(!solo.othersDecode(&self.lanes, &streams[0]));
+}

zig/src/families/glm/cuda_prompt_solo.zig

@@ -0,0 +1,89 @@
+//! A solo prompt pass (TF_GLM_CHUNK_SOLO): two grid chunks in one pass, while no other stream decodes.
+const std = @import("std");
+const fw = @import("cuda_forward.zig");
+const engine = @import("cuda_engine.zig");
+const Engine = engine.Engine;
+
+/// Whether a stream besides `s` is past its prompt pass (a lane with no fill): a solo pass would hold up its rounds.
+pub fn othersDecode(lanes: anytype, s: anytype) bool {
+    var it = lanes.iterator();
+    while (it.next()) |kv| if (kv.key_ptr.* != s and kv.value_ptr.fill == null) return true;
+    return false;
+}
+
+/// The rows an engine sizes its prompt buffers for: at start, as they hold memory all the time.
+pub fn bufferRows(on: bool) usize {
+    return if (on) fw.solo_rows else fw.prefill_rows;
+}
+
+/// Whether a pass at `start` of a `len`-id prompt can reach past the next grid point: two rows, and no mark at it.
+pub fn takesTwo(start: usize, len: usize, marks: []const u32) bool {
+    const next = start + fw.prefill_rows;
+    return next + 1 < len and std.mem.indexOfScalar(u32, marks, @intCast(next)) == null; // one row: a chunk's part as sent
+}
+
+/// Whether a pass at `start` to `rows` rebases the index rings where chunks one at a time do: at its start, or not.
+pub fn ringsAgree(base: usize, ring: usize, start: usize, rows: usize) bool {
+    return start + fw.prefill_rows > base + ring or start + rows <= base + ring;
+}
+
+/// The rows a solo pass takes at `start` of a `len`-id prompt.
+fn rowsAt(start: usize, len: usize) usize {
+    return @min(len - start, fw.solo_rows);
+}
+
+/// Rank 0's choice for the step at `start` (the others take the ids it sends); not beside what reads a chunk's rows.
+pub fn at(e: *const Engine, start: usize, len: usize, from: ?engine.Resume, marks: ?engine.Marks) bool {
+    if (e.bp.rows < fw.solo_rows or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false;
+    if (e.images) |v| if (v.rowCount() > 0) return false;
+    if (!takesTwo(start, len, if (marks) |m| m.at else &.{})) return false;
+    // a resumed step restores its kept state, and the rings' base, after this choice
+    const base = if (from) |r| (e.snaps.get(r.id) orelse return false).bases[0] else e.cur.st.base;
+    return ringsAgree(base, e.cur.st.ring_rows, start, rowsAt(start, len));
+}
+
+/// Whether a pass of `rows` past a chunk may run: `at`'s rule on every rank, its kept state restored.
+pub fn allowed(e: *const Engine, start: usize, rows: usize) bool {
+    if (e.bp.rows < rows or e.drafter != null or e.prompt_logprobs != null) return false;
+    return ringsAgree(e.cur.st.base, e.cur.st.ring_rows, start, rows);
+}
+
+/// decode._absorb_rows of a step, a chunk's rows at a time: the head's ring holds a window more than a pass.
+pub fn mtpRows(e: *Engine, f: *fw.Forward, ids: []const i32) !void {
+    const mp = f.prof;
+    defer f.prof = mp;
+    f.prof = null; // the head's pass as one section
+    var done: usize = 0;
+    while (done + 1 < ids.len) {
+        const nxt = ids[done + 1 .. @min(ids.len, done + 1 + fw.prefill_rows)];
+        try f.mtp(nxt, &.{e.bp.fnormed + done * e.c.hidden * 2});
+        try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len);
+        done += nxt.len;
+    }
+}
+
+test "a solo pass: rows past the next grid point, and no mark at it" {
+    const C = fw.prefill_rows;
+    try std.testing.expect(takesTwo(0, C + 2, &.{}));
+    try std.testing.expect(!takesTwo(0, C + 1, &.{})); // its part would read as a chunk and the next id
+    try std.testing.expect(!takesTwo(0, C, &.{})); // one chunk to the end
+    try std.testing.expect(!takesTwo(C, 2 * C, &.{}));
+    try std.testing.expect(takesTwo(2 * C, 5 * C, &.{}));
+    try std.testing.expect(!takesTwo(0, 3 * C, &.{C})); // a state kept at the next chunk: that chunk starts a step
+    try std.testing.expect(takesTwo(0, 3 * C, &.{ 2 * C, 3 * C })); // marks past the pass
+}
+
+test "a solo pass rebases the rings where chunks one at a time do" {
+    const C = fw.prefill_rows;
+    try std.testing.expect(ringsAgree(0, 2 * C, 0, 2 * C)); // neither rebases
+    try std.testing.expect(!ringsAgree(0, 2 * C, C, 2 * C)); // chunks rebase at the second, a pass at its first
+    try std.testing.expect(!ringsAgree(0, 2 * C, C, C + 1));
+    try std.testing.expect(ringsAgree(0, 2 * C, 2 * C, 2 * C)); // both at the pass's start
+    try std.testing.expect(ringsAgree(C, 2 * C, C, 2 * C));
+}
+
+test "a solo pass fits a whole ring once it rebases, and its buffers only when asked" {
+    try std.testing.expectEqual(@as(usize, fw.solo_rows), @import("cuda_state.zig").index_ring);
+    try std.testing.expectEqual(@as(usize, fw.prefill_rows), bufferRows(false));
+    try std.testing.expectEqual(@as(usize, fw.solo_rows), bufferRows(true));
+}