From acfc648ed8baa9329a31138303c872745261e09f Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 11:46:49 +0900 Subject: [PATCH 1/3] glm cuda: TF_GLM_CHUNK_SOLO takes two grid chunks in one pass while no other stream decodes Prompt chunks of 4,096 rows read the weights cheaper, but a fixed 4,096 slows warm turns and the streams decoding beside a long prompt (#1576). With the setting on, a pass widens only when nothing else decodes; the grid, the rings and prompt_rows stay at 2,048, so resumes land where main's do, and decoding streams see main's passes. Unset, the prompt buffers stay at 2,048 rows and every step is main's. --- zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_engine.zig | 20 ++-- zig/src/families/glm/cuda_follow.zig | 14 ++- zig/src/families/glm/cuda_forward.zig | 1 + zig/src/families/glm/cuda_lanes.zig | 8 +- zig/src/families/glm/cuda_lanes_test.zig | 20 ++++ zig/src/families/glm/cuda_prompt_solo.zig | 89 +++++++++++++++ zig/src/families/glm/cuda_ranks.zig | 14 ++- zig/src/families/glm/cuda_ranks_test.zig | 24 +++- zig/src/families/glm/cuda_state.zig | 2 + zig/tests/cuda/glm_solo_tests.zig | 130 ++++++++++++++++++++++ zig/tests/cuda/glm_split_tests.zig | 4 +- zig/tests/cuda/main.zig | 3 + 13 files changed, 306 insertions(+), 25 deletions(-) create mode 100644 zig/src/families/glm/cuda_prompt_solo.zig create mode 100644 zig/tests/cuda/glm_solo_tests.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 5ae6f66455..ae3b2d7b9d 100644 --- a/zig/src/families/glm/cuda.zig +++ b/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; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 2e1a6eb16e..87a4b07a06 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/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); diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 0e5d9a356d..cb07b3daa9 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/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. diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index f55f6311bd..1ffca0e815 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/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; diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index c199e56646..6bcd436065 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/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; }; diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig index 14e3d15a69..af503aff49 100644 --- a/zig/src/families/glm/cuda_lanes_test.zig +++ b/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])); +} diff --git a/zig/src/families/glm/cuda_prompt_solo.zig b/zig/src/families/glm/cuda_prompt_solo.zig new file mode 100644 index 0000000000..a246432bf8 --- /dev/null +++ b/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)); +} diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index 689254c843..5b6bd58c61 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -15,6 +15,7 @@ pub const Settings = struct { ring: bool = false, // TF_GLM_PREFILL_REDUCE=ring: split, its exchanges passed along the ring (cuda_relay) kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE: 1 (decode windows on KDA's three-kernel chain) or 0 pieces: u32 = 3, // TF_GLM_OVERLAP_PIECES: a prompt chunk's exchanges in this many row pieces (1: whole). + solo: bool = false, // TF_GLM_CHUNK_SOLO: 4096 (prompt buffers for two chunks a pass, while no stream decodes) or 2048 index_split: u32 = 2048, // TF_GLM_INDEX_SPLIT: from this position on, rank 0 splits a prompt chunk's indexer rows (0: never) heat: ?heat.Bands = null, // TF_HEAT_HIGH and TF_HEAT_LOW (upstream's names): a prompt chunk waits on every rank's hottest zone layers: u32 = 0, // TF_GLM_LAYERS: the first this many layers only, a thin model for one-GPU checks (0: all) @@ -58,6 +59,12 @@ pub const Settings = struct { return error.BadSetting; } } + if (get("TF_GLM_CHUNK_SOLO")) |v| { + if (std.mem.eql(u8, v, "4096")) s.solo = true else if (std.mem.eql(u8, v, "2048")) s.solo = false else { + problem.* = "TF_GLM_CHUNK_SOLO: 2048 or 4096"; + return error.BadSetting; + } + } if (get("TF_GLM_OVERLAP_PIECES")) |v| { s.pieces = std.fmt.parseInt(u32, v, 10) catch 0; if (s.pieces < 1 or s.pieces > 8) { @@ -119,7 +126,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 = 13, // this struct's version + layout: u32 = 14, // this struct's version world: u32, rank: u32, kv8: u32 = 0, @@ -137,6 +144,7 @@ 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) + 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 @@ -164,6 +172,7 @@ pub const Fingerprint = extern struct { max_videos, max_frames, roce_multi, + solo, bits, capacity, heat_high, @@ -196,6 +205,7 @@ pub const Fingerprint = extern struct { .layers = s.layers, .mtp = @intFromBool(mtp), .roce_multi = @intFromBool(s.roce_multi), + .solo = @intFromBool(s.solo), .bits = @as(u32, s.dense_bits) << 8 | s.head_bits, .capacity = capacity, .heat_high = if (s.heat) |b| b.high_c else 0, @@ -232,6 +242,7 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "max_videos", "max_frames", "roce_multi", + "solo", "bits", "capacity", "heat_high", @@ -341,6 +352,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .max_videos => "the video count limit", .max_frames => "the frame limit", .roce_multi => "TF_GLM_ROCE_MULTI", + .solo => "TF_GLM_CHUNK_SOLO", .bits => "TF_GLM_DENSE_BITS or TF_GLM_HEAD_BITS", .capacity => "the sequence capacity (the context)", .heat_high => "TF_HEAT_HIGH", diff --git a/zig/src/families/glm/cuda_ranks_test.zig b/zig/src/families/glm/cuda_ranks_test.zig index 1c55285883..09ecd38fd4 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", "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", "capacity", "heat_high", "heat_low" }) |name| { var all = [_]Fingerprint{ zero, one }; @field(all[1], name) += 1; try std.testing.expectEqual(@field(Fingerprint.Field, name), mismatch(&all).?.field); @@ -238,6 +238,28 @@ test "ranks disagree on tensor bytes under identical headers, and agree on a cop try std.testing.expect(mismatch(&all) == null); } +test "TF_GLM_CHUNK_SOLO: off by default, 2048 or 4096, and every rank must agree (the others take rank 0's solo passes)" { + const Reader = struct { + fn absent(_: [:0]const u8) ?[]const u8 { + return null; + } + fn on(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_CHUNK_SOLO")) "4096" else null; + } + fn bad(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_CHUNK_SOLO")) "1" else null; + } + }; + var problem: []const u8 = ""; + try std.testing.expect(!(try Settings.read(Reader.absent, &problem)).solo); + const on = try Settings.read(Reader.on, &problem); + try std.testing.expect(on.solo); + try std.testing.expectError(error.BadSetting, Settings.read(Reader.bad, &problem)); + try std.testing.expectEqualStrings("TF_GLM_CHUNK_SOLO: 2048 or 4096", problem); + const all = [_]Fingerprint{ Fingerprint.of(2, 0, .{}, 8192, true, @splat(0), @splat(0), @splat(0)), Fingerprint.of(2, 1, on, 8192, true, @splat(0), @splat(0), @splat(0)) }; + try std.testing.expectEqual(Mismatch{ .field = .solo, .rank = 1 }, mismatch(&all).?); +} + test "prompt K experiment is explicit and all ranks must agree" { const Reader = struct { fn absent(_: [:0]const u8) ?[]const u8 { diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 68cdd0c29b..72196feeed 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -25,6 +25,8 @@ pub const prompt_heads = 16; pub const dense_reach = 2560; /// Index-key and gate rows a DSA cache holds (a ring): pool_keys reads only from the current pool's first row on. pub const index_ring = 2 * prefill_rows; +/// A solo prompt pass's rows (TF_GLM_CHUNK_SOLO): a whole ring, which a rebase at a grid point leaves it. +pub const solo_rows = index_ring; /// Sub-allocations of one device allocation, 256-byte aligned like the driver's own. const Arena = struct { diff --git a/zig/tests/cuda/glm_solo_tests.zig b/zig/tests/cuda/glm_solo_tests.zig new file mode 100644 index 0000000000..f34b4296d5 --- /dev/null +++ b/zig/tests/cuda/glm_solo_tests.zig @@ -0,0 +1,130 @@ +//! GLM-5.3-Flash's solo prompt passes (TF_GLM_CHUNK_SOLO) against twins a chunk a step, on one Engine and its ranks. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// The KDA records' copy a state reads next, both cleared in `bytes`: a pass flips them, so a solo pass's differ. +fn takeRecords(a: std.mem.Allocator, st: anytype, bytes: []u8) ![]const u8 { + const at = st.rec - st.arena.buf.ptr; + const one = st.kda_layers * st.rec_layer; + const cur: usize = if (st.kda_layers > 0) st.cur[0] else 0; + const out = try a.dupe(u8, bytes[at + cur * one ..][0..one]); + @memset(bytes[at..][0 .. 2 * one], 0); + return out; +} + +/// `glm-solo-check`: prompts a step at a time, solo and not; first token, next window and every state byte equal. +pub fn soloCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { + const glm = @import("glm_cuda"); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dc = try std.json.parseFromSliceLeaky(struct { capacity: usize = 8192 }, a, try std.Io.Dir.cwd().readFileAlloc(gpu.io, dc_path, a, .limited(1 << 26)), .{ .ignore_unknown_fields = true }); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try @import("glm_lanes_tests.zig").openComm(gpu, &lib, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(comm); + var problem: []const u8 = ""; + const set = glm.ranks.Settings.read(&glm.ranks.environment, &problem) catch |err| { + std.debug.print("REFUSE {s}\n", .{problem}); + return err; + }; + try check.expect(set.solo, "glm-solo-check needs TF_GLM_CHUNK_SOLO=4096", .{}); + const C = glm.engine.prompt_rows; + const capacity = @max(dc.capacity, 6 * C); + const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = capacity, .mtp = true, .settings = set }); + defer e.deinit(); + var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, capacity); + defer link.deinit(); + if (rank != 0) { + try glm.follow.follow(e, &link); + check.pass("rank {d} followed rank 0's solo and single prompt steps to the end", .{rank}); + return; + } + const x: glm.follow.Leader = .{ .e = e, .link = &link }; + defer x.end() catch {}; + const V = e.share.vocab.len(); + const Case = struct { len: usize, marks: []const u32 = &.{} }; + const cases = [_]Case{ + .{ .len = 2 * C + 300 }, // a solo pass and a tail + .{ .len = C + 2 }, // the shortest solo pass + .{ .len = C + 1 }, // a chunk and one row, a step each + .{ .len = 3 * C }, // a solo pass and a chunk + .{ .len = 2 * C }, // a solo pass to the end + .{ .len = C }, // one chunk + .{ .len = 5 * C + 7 }, // solo passes past a rebase + .{ .len = 4 * C, .marks = &.{C} }, // a mark at chunk 1: 0 and 1 alone, then a solo pass + }; + const Kept = struct { + fn call(_: *anyopaque, _: u32) void {} // the check keeps no state: its marks only shape the steps + }; + var no_ctx: u8 = 0; + var failed = false; + for (cases, 0..) |case, i| { + const len = case.len; + const marks: ?glm.engine.Marks = if (case.marks.len == 0) null else .{ .at = case.marks, .ctx = &no_ctx, .call = Kept.call }; + const ids = try a.alloc(u32, len); + for (ids, 0..) |*u, j| u.* = @intCast(100 + (j * 7919 + i * 31) % 5000); + var firsts: [2]u32 = undefined; + var outs: [2][4]u32 = undefined; + var logits: [2][]u16 = undefined; + var solo: [2]u64 = @splat(0); + for (0..2) |twin| { // twin 0 a chunk a step, twin 1 solo where it can, both with the same marks + try x.bind(@intCast(2 * i + twin)); + try x.begin(null); + var at: usize = 0; + const t = while (true) { + const r = x.chunk(ids, at, twin == 1, null, null, marks) catch |err| { + std.debug.print("FAIL solo: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); + return err; + }; + solo[twin] += @intFromBool(r.rows > C); + if (r.first) |f| break f; + at += r.rows; + }; + firsts[twin] = t; + const win = [_]u32{ t, 101, 202, 303 }; + try x.verify(&win, &outs[twin], 0); + logits[twin] = try a.alloc(u16, win.len * V); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(logits[twin].ptr, e.bd.logits, win.len * V * 2), "cuMemcpyDtoH"); + } + const st = &e.seqs.items[2 * i].?.st; + const tw = &e.seqs.items[2 * i + 1].?.st; + var want: u64 = 0; // the solo passes prompt_solo.at allows over this prompt, the rings' base followed here + var base: usize = 0; + var at: usize = 0; + while (at < len) { + const reach = @min(len - at, 2 * C); + const two = glm.prompt_solo.takesTwo(at, len, case.marks) and glm.prompt_solo.ringsAgree(base, st.ring_rows, at, reach); + const rows = @min(len - at, @as(usize, if (two) 2 * C else C)); + if (at + rows > base + st.ring_rows) base = at; // State.ready rebases at the step's start, a whole pool + want += @intFromBool(two); + at += rows; + } + if (solo[0] != 0 or solo[1] != want) { + std.debug.print("FAIL solo: prompt of {d}: {d} solo passes a chunk a step, {d} solo (want {d})\n", .{ len, solo[0], solo[1], want }); + failed = true; + } + if (firsts[0] != firsts[1] or !std.mem.eql(u32, &outs[0], &outs[1]) or !std.mem.eql(u16, logits[0], logits[1])) { + std.debug.print("FAIL solo: prompt of {d}: first {d}/{d}, the next window's tokens or logits differ\n", .{ len, firsts[0], firsts[1] }); + failed = true; + } + try e.stream.synchronize(); + const n = st.arena.buf.len; + const h0 = try a.alloc(u8, n); + const h1 = try a.alloc(u8, n); + try st.arena.buf.download(0, h0); + try tw.arena.buf.download(0, h1); + const r0 = try takeRecords(a, st, h0); + const r1 = try takeRecords(a, tw, h1); + const diff = std.mem.indexOfDiff(u8, h0, h1); + if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or st.base != tw.base or st.mtp_base != tw.mtp_base or !std.mem.eql(u8, r0, r1) or diff != null) { + std.debug.print("FAIL solo: prompt of {d}: the solo state differs (pos {d}/{d}, rings' base {d}/{d}, MTP base {d}/{d}, records {}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, st.base, tw.base, st.mtp_base, tw.mtp_base, std.mem.eql(u8, r0, r1), diff, n }); + failed = true; + } + } + if (failed) return error.TestFailed; + check.pass("solo prompt passes at TP={d}: {d} prompts, one with a mark at chunk 1, equal their twins a chunk a step in the first token, the next window's logits and tokens, the rings' bases and every state byte", .{ world, cases.len }); +} diff --git a/zig/tests/cuda/glm_split_tests.zig b/zig/tests/cuda/glm_split_tests.zig index 06d9861bb5..6aeb206763 100644 --- a/zig/tests/cuda/glm_split_tests.zig +++ b/zig/tests/cuda/glm_split_tests.zig @@ -54,11 +54,11 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c const before = e.index_splits; var at: usize = 0; const t = while (true) : (at += C) { - const r = x.chunk(ids, at, null, null, null) catch |err| { + const r = x.chunk(ids, at, false, null, null, null) catch |err| { std.debug.print("FAIL split: prompt of {d}, twin {d}: the chunk at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); return err; }; - if (r) |f| break f; + if (r.first) |f| break f; }; splits[twin] = e.index_splits - before; firsts[twin] = t; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index a3ad2d2dad..eff29cfbdf 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -65,6 +65,7 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); 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_synth = @import("glm_synth.zig"); const tree_round_tests = @import("tree_round_tests.zig"); @@ -249,6 +250,7 @@ const usage = \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte \\ glm-split-check prompt chunks with their indexer rows split across ranks against whole twins, byte for byte + \\ glm-solo-check solo prompt passes against twins a chunk a step, logits and states byte for byte \\ glm-paged-check a paged stream (pages out of order, NaN first) against a whole-window twin \\ dory MODEL LOOPS N TEXT Dory (NVFP4): N greedy tokens from TEXT, then verify windows of 2 to 16 rows bit-equal \\ dory-serve MODEL LOOPS TEXT N Dory served: eager == graphs, copy drafts == plain, 2/4/8 streams == solo @@ -478,6 +480,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-paged-check")) return glm_paged_tests.pagedCheck(gpu, try arg(rest, 0), try std.fmt.parseInt(c_int, try arg(rest, 1), 10), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-split-check")) return glm_split_tests.splitCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-solo-check")) return glm_solo_tests.soloCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-synth")) { const kind = if (rest.len > 1) std.meta.stringToEnum(glm_synth.Kind, rest[1]) orelse return error.UnknownCommand else .nvfp4; return glm_synth.run(gpu, try arg(rest, 0), kind); From 00b3cbbd5b8773b79b5b61462531a8fee3e3df83 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:09:59 +0900 Subject: [PATCH 2/3] docs: the GLM CUDA recipe lists TF_GLM_CHUNK_SOLO with its measured trade-off --- docs/recipes/glm-5.3-flash-cuda.md | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 4a613e97cb..c0263d840d 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/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 From 858f2df3815a98cd745c5edba1cc9af9c3a33579 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:47:57 +0000 Subject: [PATCH 3/3] changelog: #1823 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 91135a5eb3..91bcfefaec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash on CUDA: `TF_GLM_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).