From e05b0447ab4a8effcf92c551a99affcfffda1460 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:24:06 +0900 Subject: [PATCH 1/9] glm cuda: a prompt step takes two chunks while no other stream decodes, each one's exchanges and glue under the other's sites --- zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_engine.zig | 14 ++- zig/src/families/glm/cuda_follow.zig | 8 +- zig/src/families/glm/cuda_forward.zig | 27 +++-- zig/src/families/glm/cuda_forward_pair.zig | 85 ++++++++++++++ zig/src/families/glm/cuda_forward_pieces.zig | 105 +++++++++-------- zig/src/families/glm/cuda_lanes_test.zig | 3 +- zig/src/families/glm/cuda_prompt_pair.zig | 114 +++++++++++++++++++ zig/src/families/glm/cuda_ranks.zig | 14 ++- zig/src/families/glm/cuda_ranks_test.zig | 24 +++- zig/src/families/glm/cuda_state.zig | 27 ++++- zig/tests/cuda/glm_pair_tests.zig | 103 +++++++++++++++++ zig/tests/cuda/main.zig | 3 + 13 files changed, 461 insertions(+), 68 deletions(-) create mode 100644 zig/src/families/glm/cuda_forward_pair.zig create mode 100644 zig/src/families/glm/cuda_prompt_pair.zig create mode 100644 zig/tests/cuda/glm_pair_tests.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index ae3b2d7b9d..b1ba5c4d73 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -31,6 +31,7 @@ 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 prompt_pair = @import("cuda_prompt_pair.zig"); pub const pages = @import("cuda_pages.zig"); pub const open_owner = @import("cuda_open_owner.zig"); pub const images = @import("cuda_images.zig"); @@ -51,6 +52,7 @@ test { _ = @import("cuda_roce_multi.zig"); _ = @import("cuda_lanes_test.zig"); _ = prompt_solo; + _ = prompt_pair; _ = experts; _ = exl3; _ = vision; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 87a4b07a06..724a36746d 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -18,6 +18,7 @@ const mm_mod = @import("cuda_mm.zig"); const dflash = @import("dflash.zig"); const prof_mod = @import("cuda_prof.zig"); const pages_mod = @import("cuda_pages.zig"); +const prompt_pair = @import("cuda_prompt_pair.zig"); pub const Options = @import("cuda_engine_types.zig").Options; /// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. @@ -71,6 +72,7 @@ pub const Engine = struct { pieces: fw.Pieces, w: weights.Weights, bp: fw.Buffers, // prompt chunks + pair: ?prompt_pair.Pair = null, // TF_GLM_PROMPT_PAIR: a paired prompt step's second chunk bd: fw.Buffers, // verify windows bm: fw.Buffers, // the MTP head's rows bs: ?fw.Buffers = null, // shared verify windows (with the MTP head, as the lane core shares only then) @@ -141,7 +143,7 @@ pub const Engine = struct { for (spans, 0..) |*sp, q| sp.* = (try shard.Share.of(&c, q, world)).vocab; var m = try fw.Modules.load(d, ctx); errdefer m.deinit(); - var pieces = try fw.Pieces.init(d, o.settings.pieces); + var pieces = try fw.Pieces.init(d, o.settings.pieces, o.settings.pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers); @@ -180,6 +182,8 @@ pub const Engine = struct { 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 }); + if (o.settings.pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); + errdefer if (e.pair) |*p| p.deinit(); 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); @@ -227,6 +231,7 @@ pub const Engine = struct { if (e.bs) |*b| b.deinit(); e.bm.deinit(); e.bd.deinit(); + if (e.pair) |*p| p.deinit(); e.bp.deinit(); if (e.pages) |p| p.deinit(); for (0..e.seqs.items.len) |i| e.free(i); @@ -246,10 +251,11 @@ pub const Engine = struct { return .{ .nccl = .{ .lib = e.lib, .comm = e.comm, .group = e.group } }; } - fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { + pub fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null }; + const twin = if (e.pair) |*p| (if (b == &p.bp) &p.pieces else null) else null; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else twin, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -342,6 +348,7 @@ pub const Engine = struct { 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, @as(usize, if (part.len > fw.prefill_rows + 1) fw.solo_rows else fw.prefill_rows)); + if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks); 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; @@ -513,6 +520,7 @@ pub const Engine = struct { } pub const holdPrompt = @import("cuda_prompt_hold.zig").save; + pub const pairedAt = prompt_pair.pairedAt; pub const prefillMany = @import("cuda_prompts.zig").prefillMany; /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran). diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index cb07b3daa9..f064c469f6 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -223,13 +223,13 @@ pub const Leader = struct { return x.e.prefill(ids, from, cancel, marks); } - /// 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 { + /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and prompt_solo.at or the pair's rule allow. + pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (solo and prompt_solo.at(x.e, start, ids.len, from, marks)) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, 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); + const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // both chunks of a pair share it 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; const first = try x.e.prefillChunk(part, start, ids.len, from, cancel, marks); diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 1ffca0e815..6cd5a9daf4 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -133,10 +133,21 @@ pub const Pieces = struct { gathered: [max_pieces]cuda.Event, glues: [max_pieces]cuda.Event, // split: a piece's own rows glued shared: [max_pieces]cuda.Event, // split: every rank's glued rows of a piece in + owns_side: bool = true, // a paired step's second chunk (initTwin) runs on the first's side stream - pub fn init(d: *const cuda.Driver, count: usize) !Pieces { + /// With `pair` (TF_GLM_PROMPT_PAIR) the side stream runs at the highest priority, ahead of the main stream's grids. + pub fn init(d: *const cuda.Driver, count: usize, pair: bool) !Pieces { if (count == 0 or count > max_pieces) return error.BadPieces; - var p: Pieces = .{ .count = count, .side = try cuda.Stream.init(d, false), .filled = undefined, .gathered = undefined, .glues = undefined, .shared = undefined }; + return make(d, count, if (pair) try cuda.Stream.initPriority(d, .highest) else try cuda.Stream.init(d, false), true); + } + + /// A paired step's second chunk on `of`'s side stream: every rank posts both chunks' exchanges in one order. + pub fn initTwin(d: *const cuda.Driver, of: *const Pieces) !Pieces { + return make(d, of.count, of.side, false); + } + + fn make(d: *const cuda.Driver, count: usize, side: cuda.Stream, owns_side: bool) !Pieces { + var p: Pieces = .{ .count = count, .side = side, .filled = undefined, .gathered = undefined, .glues = undefined, .shared = undefined, .owns_side = owns_side }; for (&p.filled, &p.gathered, &p.glues, &p.shared) |*a, *b, *g, *s| { a.* = try cuda.Event.init(d, false); b.* = try cuda.Event.init(d, false); @@ -153,7 +164,7 @@ pub const Pieces = struct { g.deinit(); s.deinit(); } - p.side.deinit(); + if (p.owns_side) p.side.deinit(); p.* = undefined; } }; @@ -225,7 +236,7 @@ pub const Forward = struct { } /// forward._pre: hc_pre of every row at a site into b.normed, xs, post and comb. - fn pre(f: *const Forward, r: Rank, hc: weights.Hc, norm: glue.Norm, rows: usize) !void { + pub fn pre(f: *const Forward, r: Rank, hc: weights.Hc, norm: glue.Norm, rows: usize) !void { const stamp = prof_mod.span(f.prof, f.stream, .glue); defer stamp.end(); const c = &r.w.config; @@ -241,7 +252,7 @@ pub const Forward = struct { } /// forward._mixer: the layer's KDA or DSA block on b.normed; `li`/`di`: its index among loaded KDA / DSA layers. - fn mixer(f: *const Forward, r: Rank, l: weights.Layer, li: usize, di: usize, rows: usize, nch: usize, out: ?u64) !void { + pub fn mixer(f: *const Forward, r: Rank, l: weights.Layer, li: usize, di: usize, rows: usize, nch: usize, out: ?u64) !void { const c = &r.w.config; const b = r.b; const st = r.st; @@ -280,7 +291,7 @@ pub const Forward = struct { } /// forward._ffn: the dense MLP or MoE block on b.normed into this rank's slot. - fn ffn(f: *const Forward, r: Rank, l: weights.Layer, rows: usize, out: ?u64) !void { + pub fn ffn(f: *const Forward, r: Rank, l: weights.Layer, rows: usize, out: ?u64) !void { const stamp = prof_mod.span(f.prof, f.stream, .ffn); defer stamp.end(); const c = &r.w.config; @@ -402,6 +413,8 @@ pub const Forward = struct { } } + pub const pairBody = @import("cuda_forward_pair.zig").pairBody; + /// forward.forward of `ids`: logits in b.logits; a prompt chunk's KDA layers committed, a window's at `commit`. pub fn run(f: *const Forward, ids: []const i32) !void { const R = ids.len; @@ -464,7 +477,7 @@ pub const Forward = struct { } /// The final norm and this rank's rows of the head (a prompt buffer's last row only). - fn finish(f: *const Forward, R: usize) !void { + pub fn finish(f: *const Forward, R: usize) !void { prof_mod.context(f.prof, null, null, .chunk); const stamp = prof_mod.span(f.prof, f.stream, .head); defer stamp.end(); diff --git a/zig/src/families/glm/cuda_forward_pair.zig b/zig/src/families/glm/cuda_forward_pair.zig new file mode 100644 index 0000000000..eec6992e70 --- /dev/null +++ b/zig/src/families/glm/cuda_forward_pair.zig @@ -0,0 +1,85 @@ +//! A paired prompt step's pass (TF_GLM_PROMPT_PAIR): two chunks site by site, each one's exchanges under the other's. +const glue = @import("cuda_glue.zig"); +const latent = @import("cuda_latent.zig"); +const fw = @import("cuda_forward.zig"); +const pieces = @import("cuda_forward_pieces.zig"); +const Forward = fw.Forward; +const Pieces = fw.Pieces; +const Site = fw.Forward.Site; +const room = fw.Forward.room; +const share = fw.Forward.share; + +/// Two consecutive prompt chunks of one sequence (`b` a view of `a`'s at its end), each on `pieced`'s kernels and bits. +pub fn pairBody(a: *const Forward, b: *const Forward, ra: usize, rb: usize) !void { + const two = [_]*const Forward{ a, b }; + const rows = [_]usize{ ra, rb }; + var cuts: [2][fw.max_pieces][2]usize = undefined; + var n: [2]usize = undefined; + for (two, rows, 0..) |f, R, i| { + if (f.pieces == null or f.prof != null or f.taps != null) return error.BadPair; + const c = &f.ranks[0].w.config; + for (f.ranks) |r| { + if (r.segs.len > 0 or !r.b.prefill or R == 0 or R > r.b.rows) return error.BadPair; + try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); + } + if (f.images) |im| for (f.ranks) |r| try im.v.splice(f.stream, im.slot, r.st.pos, r.st.pos + R, r.b.x, c.hc); + n[i] = fw.ranges(R, f.pieces.?.count, &cuts[i]); + if (n[i] == 1 and !f.splits(R)) return error.BadPair; // run would take this chunk whole + for (f.ranks) |r| try f.pre(r, r.w.layers[0].attn_hc, r.w.layers[0].in_norm, R); + } + const layers = a.ranks[0].w.layers; + var li: usize = 0; + var di: usize = 0; + for (layers, 0..) |l, j| { + // b's mixer follows a's on the main stream, so a's caches, KDA state and conv rows of this layer are in + for (two, rows, 0..) |f, R, i| { + if (j > 0) try awaitShared(f, f.pieces.?, n[i]); + for (f.ranks) |r| try f.mixer(r, r.w.layers[j], li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), null); + try exchangeGlue(f, f.pieces.?, j, .mixer, cuts[i][0..n[i]]); + } + for (two, rows, 0..) |f, R, i| { + try awaitShared(f, f.pieces.?, n[i]); + for (f.ranks) |r| try f.ffn(r, r.w.layers[j], R, null); + try exchangeGlue(f, f.pieces.?, j, .ffn, cuts[i][0..n[i]]); + } + switch (l.mixer) { + .kda => li += 1, + .mla => di += 1, + } + } + for (two, rows, 0..) |f, R, i| { + try awaitShared(f, f.pieces.?, n[i]); + try f.finish(R); + } +} + +/// A paired site: each piece filled on the main stream, then on the side stream its exchange, glue and (split) sharing. +fn exchangeGlue(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { + const D = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + const last = site == .ffn and j + 1 == f.ranks[0].w.layers.len; + var g = f.*; // the glue's launches go on the side stream + g.stream = p.side; + for (cut, 0..) |piece, q| { + const lo = piece[0]; + const n = piece[1] - lo; + const split = f.splits(n); + for (f.ranks) |r| try f.fillRows(r, r.w.layers[j], site, lo, piece[1], if (split) r.b.part + lo * D * 4 else room(r, lo) + r.w.share.rank * n * D * 4); + try p.filled[q].record(f.stream); + try p.side.wait(p.filled[q]); + try pieces.partials(f, p, lo, n, split); + if (split) { + for (f.ranks) |r| { + const mine = share(n, world, r.w.share.rank); + try g.glueRows(r, j, site, lo + mine[0], lo + mine[1], room(r, lo)); + } + try pieces.shareGlued(f, p, lo, n, last); + } else for (f.ranks) |r| try g.glueRows(r, j, site, lo, piece[1], room(r, lo)); + try p.shared[q].record(p.side); + } +} + +/// Before a paired chunk's next site: its last site's pieces glued and shared on the side stream. +fn awaitShared(f: *const Forward, p: *const Pieces, count: usize) !void { + for (p.shared[0..count]) |e| try f.stream.wait(e); +} diff --git a/zig/src/families/glm/cuda_forward_pieces.zig b/zig/src/families/glm/cuda_forward_pieces.zig index 2198269087..47a68e7558 100644 --- a/zig/src/families/glm/cuda_forward_pieces.zig +++ b/zig/src/families/glm/cuda_forward_pieces.zig @@ -12,7 +12,6 @@ const share = fw.Forward.share; pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { defer prof.context(f.prof, j, null, if (site == .mixer) .mixer else .ffn); const D = f.ranks[0].w.config.hidden; - const world = f.ranks[0].w.share.world; for (cut, 0..) |piece, q| { prof.context(f.prof, j, q, if (site == .mixer) .mixer else .ffn); const lo = piece[0]; @@ -24,33 +23,7 @@ pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, try p.side.wait(p.filled[q]); stamp.end(); stamp = prof.span(f.prof, p.side, .partials); - const bytes = n * D * 4; - if (!split) switch (f.exchange) { - .local => for (f.ranks) |to| for (f.ranks) |from| { - if (to.w.share.rank == from.w.share.rank) continue; - const off = from.w.share.rank * bytes; - try f.copy(p.side, room(to, lo) + off, room(from, lo) + off, bytes); - }, - .nccl => |x| { - const r = f.ranks[0]; - try x.lib.check(x.lib.api.ncclAllGather(room(r, lo) + r.w.share.rank * bytes, room(r, lo), n * D, .f32, x.comm, p.side.handle), "ncclAllGather"); - }, - } else switch (f.exchange) { - .local => for (f.ranks) |to| { - const mine = share(n, world, to.w.share.rank); - const own = mine[1] - mine[0]; - for (f.ranks) |from| try f.copy(p.side, room(to, lo) + from.w.share.rank * own * D * 4, from.b.part + (lo + mine[0]) * D * 4, own * D * 4); - }, - .nccl => |x| { - const r = f.ranks[0]; - const me = r.w.share.rank; - const mine = share(n, world, me); - const own = mine[1] - mine[0]; - const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; - if (f.ring) try relay.partials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D) else try relay.splitPartials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D); - try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4); - }, - } + try partials(f, p, lo, n, split); stamp.end(); try p.gathered[q].record(p.side); } @@ -59,8 +32,6 @@ pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, /// overlap.Overlap.glue: each piece's glue once its exchange is in, its rows then shared on the second stream. pub fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { defer prof.context(f.prof, j, null, if (site == .mixer) .mixer else .ffn); - const c = &f.ranks[0].w.config; - const D: usize = c.hidden; const world = f.ranks[0].w.share.world; const last = site == .ffn and j + 1 == f.ranks[0].w.layers.len; var sent: usize = 0; @@ -84,23 +55,7 @@ pub fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut try p.side.wait(p.glues[q]); stamp.end(); stamp = prof.span(f.prof, p.side, .glued); - switch (f.exchange) { - .local => for (f.ranks) |from| { - const mine = share(n, world, from.w.share.rank); - const bytes = (mine[1] - mine[0]) * D * 2; - for (f.ranks) |to| if (to.w.share.rank != from.w.share.rank) { - const at = (lo + mine[0]) * D * 2; - try f.copy(p.side, (if (last) to.b.hidden else to.b.normed) + at, (if (last) from.b.hidden else from.b.normed) + at, bytes); - }; - }, - .nccl => |x| { - const r = f.ranks[0]; - const me = r.w.share.rank; - const rows_of = if (last) r.b.hidden else r.b.normed; - const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; - if (f.ring) try relay.glued(x.lib, x.comm, p.side, plan, rows_of, lo, D) else try relay.splitGlued(x.lib, x.comm, p.side, plan, rows_of, lo, D); - }, - } + try shareGlued(f, p, lo, n, last); stamp.end(); try p.shared[q].record(p.side); sent |= @as(usize, 1) << @intCast(q); @@ -112,3 +67,59 @@ pub fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut stamp.end(); }; } + +/// A piece's partials on the side stream: gathered whole, or split (each rank's rows of every partial to their owner). +pub fn partials(f: *const Forward, p: *const Pieces, lo: usize, n: usize, split: bool) !void { + const D = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + const bytes = n * D * 4; + if (!split) switch (f.exchange) { + .local => for (f.ranks) |to| for (f.ranks) |from| { + if (to.w.share.rank == from.w.share.rank) continue; + const off = from.w.share.rank * bytes; + try f.copy(p.side, room(to, lo) + off, room(from, lo) + off, bytes); + }, + .nccl => |x| { + const r = f.ranks[0]; + try x.lib.check(x.lib.api.ncclAllGather(room(r, lo) + r.w.share.rank * bytes, room(r, lo), n * D, .f32, x.comm, p.side.handle), "ncclAllGather"); + }, + } else switch (f.exchange) { + .local => for (f.ranks) |to| { + const mine = share(n, world, to.w.share.rank); + const own = mine[1] - mine[0]; + for (f.ranks) |from| try f.copy(p.side, room(to, lo) + from.w.share.rank * own * D * 4, from.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const mine = share(n, world, me); + const own = mine[1] - mine[0]; + const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; + if (f.ring) try relay.partials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D) else try relay.splitPartials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D); + try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + } +} + +/// A split piece's glued rows (normed, or hidden after the last layer) to every rank on the side stream. +pub fn shareGlued(f: *const Forward, p: *const Pieces, lo: usize, n: usize, last: bool) !void { + const D: usize = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + switch (f.exchange) { + .local => for (f.ranks) |from| { + const mine = share(n, world, from.w.share.rank); + const bytes = (mine[1] - mine[0]) * D * 2; + for (f.ranks) |to| if (to.w.share.rank != from.w.share.rank) { + const at = (lo + mine[0]) * D * 2; + try f.copy(p.side, (if (last) to.b.hidden else to.b.normed) + at, (if (last) from.b.hidden else from.b.normed) + at, bytes); + }; + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const rows_of = if (last) r.b.hidden else r.b.normed; + const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; + if (f.ring) try relay.glued(x.lib, x.comm, p.side, plan, rows_of, lo, D) else try relay.splitGlued(x.lib, x.comm, p.side, plan, rows_of, lo, D); + }, + } +} diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig index af503aff49..1eca112ac8 100644 --- a/zig/src/families/glm/cuda_lanes_test.zig +++ b/zig/src/families/glm/cuda_lanes_test.zig @@ -1,7 +1,8 @@ -//! The lane backend's slots: chosen before a stream's map entry exists, and none left behind by a failed insertion. +//! The lane backend's slots, chosen and given back, and when a prompt step may pair. const std = @import("std"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); +const fw = @import("cuda_forward.zig"); const Lanes = @import("cuda_lanes.zig").Lanes; const fw = @import("cuda_forward.zig"); const solo = @import("cuda_prompt_solo.zig"); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig new file mode 100644 index 0000000000..0a1c7f9ef8 --- /dev/null +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -0,0 +1,114 @@ +//! A prompt step of two chunks in one pass (TF_GLM_PROMPT_PAIR): each chunk's exchanges under the other's sites. +const std = @import("std"); +const cuda = @import("cuda"); +const fw = @import("cuda_forward.zig"); +const weights = @import("cuda_weights.zig"); +const snapshots = @import("cuda_snapshots.zig"); +const engine = @import("cuda_engine.zig"); +const Engine = engine.Engine; + +/// The fewest rows a paired step's second chunk takes: a shorter one's sites end before the first's exchanges do. +pub const min_rows = fw.prefill_rows / 4; + +/// A paired step's second chunk: twins of the prompt buffers, pieces on their side stream, and its first position. +pub const Pair = struct { + bp: fw.Buffers, + pieces: fw.Pieces, + pos: cuda.DeviceBuffer, // an i32: where the second chunk's view of the sequence starts + steps: u64 = 0, // paired steps so far (glm-pair-check counts them) + + pub fn init(d: *const cuda.Driver, w: *const weights.Weights, bp: *const fw.Buffers, pieces: *const fw.Pieces, capacity: usize) !Pair { + var b = try fw.Buffers.initTwin(d, w, bp, capacity); + errdefer b.deinit(); + var p = try fw.Pieces.initTwin(d, pieces); + errdefer p.deinit(); + return .{ .bp = b, .pieces = p, .pos = try cuda.DeviceBuffer.alloc(d, 4) }; + } + + pub fn deinit(p: *Pair) void { + p.pos.free(); + p.pieces.deinit(); + p.bp.deinit(); + p.* = undefined; + } +}; + +/// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. +pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool { + const next = start + fw.prefill_rows; + if (next >= len or len - next < @max(min_rows, world)) return false; + return std.mem.indexOfScalar(u32, marks, @intCast(next)) == null; +} + +/// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. +pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks) bool { + if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; + if (e.pieces.count == 1 and !e.settings.split) return false; // run takes such a chunk whole + return takesTwo(start, len, e.share.world, if (marks) |m| m.at else &.{}); +} + +/// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. +pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 { + const ra: usize = fw.prefill_rows; + if (len <= start + ra) return error.BadChunk; + const rb: usize = @min(len - start - ra, fw.prefill_rows); + const pair = if (e.pair) |*p| p else return error.BadChunk; + if (rb < e.share.world or part.len != @min(len - start, ra + rb + 1) or e.drafter != null or e.prompt_logprobs != null) return error.BadChunk; + if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); + e.gate(cancel) catch |err| { + if (err == error.Cancelled) try snapshots.interrupted(e, marks, @intCast(start)); + return err; + }; + var one: [1]fw.Rank = undefined; + var two: [1]fw.Rank = undefined; + var fa = e.view(&e.bp, &one); + var fb = e.view(&pair.bp, &two); + fa.prof = null; + fb.prof = null; + const st = &e.cur.st; + try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base + var sb = st.*; + sb.pos_dev = pair.pos.ptr; + try sb.setPos(e.stream, st.pos + ra); + two[0].st = &sb; + st.pending = ra; + st.prompt = true; + const ids: []const i32 = @ptrCast(part); + try cuda.copy_source.enqueue(e.d, e.bp.ids, std.mem.sliceAsBytes(ids[0..ra]), e.stream); + try cuda.copy_source.enqueue(e.d, pair.bp.ids, std.mem.sliceAsBytes(ids[ra .. ra + rb]), e.stream); + try fw.Forward.pairBody(&fa, &fb, ra, rb); + pair.steps += 1; + try mtpRows(e, &fa, ids[1 .. ra + 1], e.bp.fnormed); + try fa.commit(ra); + st.pending = rb; + st.prompt = true; + var first: ?u32 = null; + if (start + ra + rb == len) { + var t: [1]u32 = undefined; + try e.draw(pair.bp.logits, 1, len, 0, &t); + first = t[0]; + try e.holdPrompt(pair.bp.fnormed + (rb - 1) * e.c.hidden * 2); + } + try mtpRows(e, &fa, ids[ra + 1 ..], pair.bp.fnormed); // the MTP head on a's buffers, over b's rows + try fa.commit(rb); + return first; +} + +/// decode._absorb_rows of a chunk: the MTP head over its next ids `nxt` on its final-norm rows `hidden`. +fn mtpRows(e: *Engine, f: *fw.Forward, nxt: []const i32, hidden: u64) !void { + if (!e.mtp or nxt.len == 0) return; + try f.mtp(nxt, &.{hidden}); + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len); +} + +test "a prompt step takes the next chunk when it holds min_rows or more and no kept state starts there" { + const C = fw.prefill_rows; + try std.testing.expect(takesTwo(0, 2 * C, 3, &.{})); + try std.testing.expect(takesTwo(0, C + min_rows, 3, &.{})); + try std.testing.expect(!takesTwo(0, C + min_rows - 1, 3, &.{})); // a shorter tail goes alone + try std.testing.expect(!takesTwo(0, C, 3, &.{})); + try std.testing.expect(!takesTwo(C, 2 * C, 3, &.{})); // the prompt's last chunk + try std.testing.expect(takesTwo(2 * C, 5 * C, 3, &.{})); + try std.testing.expect(!takesTwo(0, 3 * C, 3, &.{C})); // a state kept at the next chunk's start: that chunk starts a step + try std.testing.expect(takesTwo(0, 3 * C, 3, &.{ 2 * C, 3 * C })); // marks past the pair +} diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index 5b6bd58c61..ead31e774d 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -16,6 +16,7 @@ pub const Settings = struct { 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 + pair: bool = true, // TF_GLM_PROMPT_PAIR: 1 (while no stream decodes, a prompt step takes two chunks, each one's exchanges under the other's sites) or 0 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) @@ -65,6 +66,12 @@ pub const Settings = struct { return error.BadSetting; } } + if (get("TF_GLM_PROMPT_PAIR")) |v| { + if (std.mem.eql(u8, v, "1")) s.pair = true else if (std.mem.eql(u8, v, "0")) s.pair = false else { + problem.* = "TF_GLM_PROMPT_PAIR: 0 or 1"; + 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) { @@ -126,7 +133,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, @@ -136,6 +143,7 @@ pub const Fingerprint = extern struct { split: u32, // TF_GLM_PREFILL_REDUCE: 0 gather, 1 split, 2 ring kda_wide: u32, pieces: u32, + pair: u32, // TF_GLM_PROMPT_PAIR: rank 0's paired steps need the second chunk's buffers on every rank layers: u32, // TF_GLM_LAYERS (0: all) mtp: u32, // 1: the MTP head loads and runs (drafts on); 2: and DFlash2 (--drafter), whose rounds gather vision: u32 = 0, // --vision's image tokens (0: off): rank 0 sends image rows every rank must hold @@ -164,6 +172,7 @@ pub const Fingerprint = extern struct { split, kda_wide, pieces, + pair, layers, mtp, vision, @@ -202,6 +211,7 @@ pub const Fingerprint = extern struct { .split = if (!s.split) 0 else if (s.ring) 2 else 1, .kda_wide = @intFromBool(s.kda_wide), .pieces = s.pieces, + .pair = @intFromBool(s.pair), .layers = s.layers, .mtp = @intFromBool(mtp), .roce_multi = @intFromBool(s.roce_multi), @@ -234,6 +244,7 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "split", "kda_wide", "pieces", + "pair", "layers", "mtp", "vision", @@ -344,6 +355,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .split => "TF_GLM_PREFILL_REDUCE", .kda_wide => "TF_GLM_KDA_DECODE_WIDE", .pieces => "TF_GLM_OVERLAP_PIECES", + .pair => "TF_GLM_PROMPT_PAIR", .layers => "TF_GLM_LAYERS", .mtp => "--no-drafts given to some ranks and not others", .vision => "--vision (or its image tokens) given to some ranks and not others", diff --git a/zig/src/families/glm/cuda_ranks_test.zig b/zig/src/families/glm/cuda_ranks_test.zig index 09ecd38fd4..07b47b963f 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", "pair", "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); @@ -292,6 +292,28 @@ test "prompt K experiment is explicit and all ranks must agree" { try std.testing.expectEqual(@as(?Mismatch, null), mismatch(&all)); } +test "TF_GLM_PROMPT_PAIR: on by default, 0 or 1, and every rank must agree (the others take rank 0's paired steps)" { + const Reader = struct { + fn absent(_: [:0]const u8) ?[]const u8 { + return null; + } + fn off(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_PROMPT_PAIR")) "0" else null; + } + fn bad(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_PROMPT_PAIR")) "2" else null; + } + }; + var problem: []const u8 = ""; + try std.testing.expect((try Settings.read(Reader.absent, &problem)).pair); + const off = try Settings.read(Reader.off, &problem); + try std.testing.expect(!off.pair); + try std.testing.expectError(error.BadSetting, Settings.read(Reader.bad, &problem)); + try std.testing.expectEqualStrings("TF_GLM_PROMPT_PAIR: 0 or 1", problem); + const all = [_]Fingerprint{ Fingerprint.of(2, 0, .{}, 8192, true, @splat(0), @splat(0), @splat(0)), Fingerprint.of(2, 1, off, 8192, true, @splat(0), @splat(0), @splat(0)) }; + try std.testing.expectEqual(Mismatch{ .field = .pair, .rank = 1 }, mismatch(&all).?); +} + test "TF_GLM_INDEX_SPLIT: 2048 by default, 0 for never, a position, and a refusal naming a value that is not one" { 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 72196feeed..2a0214987d 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -167,16 +167,25 @@ pub const Buffers = struct { part: u64, // fp32 [rows][dims]: this rank's partial rows a split exchange sends from pub fn init(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize) !Buffers { - return make(d, w, rows, prefill, capacity, rows); + return make(d, w, rows, prefill, capacity, rows, null); } /// A shared window's scratch: `rows` rows of several streams, each one's attention and KDA chain on `window_rows`. pub fn initShared(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, capacity: usize) !Buffers { - return make(d, w, rows, false, capacity, window_rows); + return make(d, w, rows, false, capacity, window_rows, null); } + /// A paired step's second chunk: its own rows across sites and streams, and `of`'s scratch for one site's work. + pub fn initTwin(d: *const cuda.Driver, w: *const weights.Weights, of: *const Buffers, capacity: usize) !Buffers { + return make(d, w, of.rows, of.prefill, capacity, of.rows, of); + } + + /// The fields a twin holds apart: make's sizes by index, then their names. + const twin_sizes = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7, 62, 63, 64, 65, 70 }; + const twin_fields = .{ "ids", "x", "normed", "xs", "post", "comb", "hcpart", "gathered", "hidden", "fnormed", "fxs", "logits", "part" }; + /// `seg`: the most rows one stream's attention takes (a decode window's rows, or a shared window's segment). - fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize) !Buffers { + fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize, twin: ?*const Buffers) !Buffers { const c = &w.config; const s = &w.share; const D: usize = c.hidden; @@ -242,8 +251,12 @@ pub const Buffers = struct { x3[0], x3[1], x3[2], x3[3], // EXL3 experts: 8-slot picks, rotated inputs, a decode window's partials if (split) rows * @max(3 * n, kw - 3 * n) * 2 else 4, // a split KDA stack's part rows }; + var held = sizes; + if (twin != null) for (&held, 0..) |*bytes, i| { // a twin allocates its own rows only + if (std.mem.indexOfScalar(usize, &twin_sizes, i) == null) bytes.* = 0; + }; var p: [sizes.len]u64 = undefined; - const a = try arena(d, &sizes, &p); + const a = try arena(d, &held, &p); var out: Buffers = .{ .arena = a, .rows = rows, @@ -296,6 +309,12 @@ pub const Buffers = struct { }; out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; out.kda.tmp = p[75]; + if (twin) |of| { + var t = of.*; + t.arena = out.arena; + inline for (twin_fields) |f| @field(t, f) = @field(out, f); + return t; + } return out; } diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig new file mode 100644 index 0000000000..ff03559ed8 --- /dev/null +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -0,0 +1,103 @@ +//! GLM-5.3-Flash's paired prompt steps (TF_GLM_PROMPT_PAIR) 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; + +/// `glm-pair-check`: prompts a step at a time, paired and not; first token, next window and every state byte equal. +pub fn pairCheck(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.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); + const C = glm.engine.prompt_rows; + const capacity = @max(dc.capacity, 4 * 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 paired 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 min = glm.prompt_pair.min_rows; + // a pair and a tail, the shortest second chunk, a pair and a chunk, a pair to the end, a tail too short to pair + const lens = [_]usize{ 2 * C + 300, C + min, 3 * C, 2 * C, C + min - 1 }; + var failed = false; + for (lens, 0..) |len, i| { + 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 paired: [2]u64 = undefined; + for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can + e.settings.pair = twin == 1; + try x.bind(@intCast(2 * i + twin)); + try x.begin(null); + const before = e.pair.?.steps; + var at: usize = 0; + const t = while (true) { + const r = x.chunk(ids, at, true, null, null, null) catch |err| { + std.debug.print("FAIL pair: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); + return err; + }; + if (r.first) |f| break f; + at += r.rows; + }; + paired[twin] = e.pair.?.steps - before; + 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"); + } + e.settings.pair = true; + var want: u64 = 0; // the steps glm.prompt_pair.takesTwo pairs over this prompt + var at: usize = 0; + while (at < len) { + const two = glm.prompt_pair.takesTwo(at, len, e.share.world, &.{}); + want += @intFromBool(two); + at += if (two) 2 * C else C; + } + if (paired[0] != 0 or paired[1] != want) { + std.debug.print("FAIL pair: prompt of {d}: {d} paired steps a chunk a step, {d} paired (want {d})\n", .{ len, paired[0], paired[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 pair: 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 st = &e.seqs.items[2 * i].?.st; + const tw = &e.seqs.items[2 * i + 1].?.st; + 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 diff = std.mem.indexOfDiff(u8, h0, h1); + if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or !std.mem.eql(u1, st.cur, tw.cur) or diff != null) { + std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, diff, n }); + failed = true; + } + } + if (failed) return error.TestFailed; + check.pass("paired prompt steps at TP={d}: prompts of {any} rows equal their twins a chunk a step in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 62a5f82b1a..a27fe6a8de 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -64,6 +64,7 @@ const lfm2_kernel_tests = @import("lfm2_kernel_tests.zig"); 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_pair_tests = @import("glm_pair_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"); @@ -250,6 +251,7 @@ const usage = \\ 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 \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ glm-pair-check paired prompt steps against twins a chunk a step, 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 @@ -482,6 +484,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(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-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-pair-check")) return glm_pair_tests.pairCheck(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")) { From 656f3635b975258c6c449a29cb9f8d05e99d7d4d Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:57:14 +0900 Subject: [PATCH 2/9] glm cuda: a paired step takes rank 0's indexer split for both chunks and counts it once (#1404) Engine.prefillChunk reads and clears the split flag before it hands a paired step to prompt_pair.prefill, which never set it on its two forwards, so with #1404 a paired step ran the indexer whole on every rank. The step's flag now reaches both chunks: the header carries the one choice rank 0 makes at the step's start, and index_splits counts the step once. The twin buffers already share the first chunk's select scratch, whose pools hold #1404's world rows more. --- zig/src/families/glm/cuda_engine.zig | 2 +- zig/src/families/glm/cuda_prompt_pair.zig | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 724a36746d..6bb4ad5fbf 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -348,7 +348,7 @@ pub const Engine = struct { 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, @as(usize, if (part.len > fw.prefill_rows + 1) fw.solo_rows else fw.prefill_rows)); - if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks); + if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks, split); 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; diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index 0a1c7f9ef8..763bf5bc10 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -48,7 +48,7 @@ pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks } /// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. -pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 { +pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks, split: bool) !?u32 { const ra: usize = fw.prefill_rows; if (len <= start + ra) return error.BadChunk; const rb: usize = @min(len - start - ra, fw.prefill_rows); @@ -65,6 +65,8 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: var fb = e.view(&pair.bp, &two); fa.prof = null; fb.prof = null; + fa.index_split = split; // the step's choice (TF_GLM_INDEX_SPLIT) for both chunks + fb.index_split = split; const st = &e.cur.st; try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base var sb = st.*; @@ -78,6 +80,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: try cuda.copy_source.enqueue(e.d, pair.bp.ids, std.mem.sliceAsBytes(ids[ra .. ra + rb]), e.stream); try fw.Forward.pairBody(&fa, &fb, ra, rb); pair.steps += 1; + e.index_splits += @intFromBool(split); try mtpRows(e, &fa, ids[1 .. ra + 1], e.bp.fnormed); try fa.commit(ra); st.pending = rb; From 5f617b95632ea18716bced72a8597fdf368e5617 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:57:15 +0900 Subject: [PATCH 3/9] glm cuda: glm-split-check takes the pair's Leader.chunk, and a fifth prompt pairs a step past the split's position The check called Leader.chunk without the pair argument, so compile-cuda-tests failed on the stacked tree. It now steps by Chunk.rows, pairs where it can and counts the splits a step at a time; pairing leaves the first four prompts fewer split steps, so a fifth prompt of three chunks and min_rows pairs its second step past TF_GLM_INDEX_SPLIT. --- zig/tests/cuda/glm_split_tests.zig | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/zig/tests/cuda/glm_split_tests.zig b/zig/tests/cuda/glm_split_tests.zig index 6aeb206763..46c490c881 100644 --- a/zig/tests/cuda/glm_split_tests.zig +++ b/zig/tests/cuda/glm_split_tests.zig @@ -5,7 +5,7 @@ const cuda = @import("cuda"); const check = @import("check.zig"); const Gpu = check.Gpu; -/// `glm-split-check`: prompts a chunk at a time, split and not; first token, next window and every state byte equal. +/// `glm-split-check`: prompts a step at a time, split and not; first token, next window and every state byte equal. pub fn splitCheck(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); @@ -36,10 +36,12 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c const x: glm.follow.Leader = .{ .e = e, .link = &link }; defer x.end() catch {}; const V = e.share.vocab.len(); - const from = C; // every chunk from the second on: each holds rows past the dense limit - // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks - const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C }; + const from = C; // every step from the second chunk on: each holds rows past the dense limit + // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks, and a second step + // that pairs past `from` (Leader.chunk pairs where it can; a paired step's two chunks take the split of its start) + const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C, 3 * C + glm.prompt_pair.min_rows }; var failed = false; + var paired: u64 = 0; // the split twins' paired steps for (lens, 0..) |len, i| { const ids = try a.alloc(u32, len); for (ids, 0..) |*u, j| u.* = @intCast(100 + (j * 7919 + i * 31) % 5000); @@ -47,20 +49,25 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c var outs: [2][4]u32 = undefined; var logits: [2][]u16 = undefined; var splits: [2]u64 = undefined; + var want: u64 = 0; // the steps glm.sparse.splitAt splits over this prompt for (0..2) |twin| { // twin 0 whole, twin 1 split from `from` e.settings.index_split = if (twin == 1) from else 0; try x.bind(@intCast(2 * i + twin)); try x.begin(null); const before = e.index_splits; + const steps = if (e.pair) |p| p.steps else 0; var at: usize = 0; - const t = while (true) : (at += C) { - 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) }); + const t = while (true) { + if (twin == 1) want += @intFromBool(glm.sparse.splitAt(at, from, e.share.world)); + const r = x.chunk(ids, at, true, null, null, null) catch |err| { + std.debug.print("FAIL split: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); return err; }; if (r.first) |f| break f; + at += r.rows; }; splits[twin] = e.index_splits - before; + if (twin == 1) paired += (if (e.pair) |p| p.steps else 0) - steps; firsts[twin] = t; const win = [_]u32{ t, 101, 202, 303 }; try x.verify(&win, &outs[twin], 0); @@ -68,11 +75,8 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(logits[twin].ptr, e.bd.logits, win.len * V * 2), "cuMemcpyDtoH"); } e.settings.index_split = 0; - var want: u64 = 0; // the chunks glm.sparse.splitAt splits over this prompt - var start: usize = 0; - while (start < len) : (start += C) want += @intFromBool(glm.sparse.splitAt(start, from, e.share.world)); if (splits[0] != 0 or splits[1] != want) { - std.debug.print("FAIL split: prompt of {d}: {d} split chunks whole, {d} split (want {d})\n", .{ len, splits[0], splits[1], want }); + std.debug.print("FAIL split: prompt of {d}: {d} split steps whole, {d} split (want {d})\n", .{ len, splits[0], splits[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])) { @@ -94,5 +98,5 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c } } if (failed) return error.TestFailed; - check.pass("split indexer rows at TP={d}: prompts of {any} rows equal their whole twins in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); + check.pass("split indexer rows at TP={d}: prompts of {any} rows ({d} paired steps) equal their whole twins in the first token, the next window's logits and tokens, and every state byte", .{ world, lens, paired }); } From a18d3d705ea8fc7859c350896b8d3a2b1828f201 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:50:24 +0900 Subject: [PATCH 4/9] glm cuda: a step pairs only where its index rings rebase as a chunk at a time's do prompt_pair.prefill readies the index rings once, at the pair's start. Where a chunk at a time rebases them only at the second chunk (from chunk 1 of rings based at 0: after a resume from a state kept at an odd chunk, or once decoders finish mid-prompt), the pair rebased at the first: the tokens agree, but the rings' bytes and base do not. pairedAt now pairs only where both rebase at the first chunk or neither rebases, judging a resumed step by its kept state's base (restored after the choice), so one unpaired chunk realigns a prompt; prefill refuses any other pair on every rank. glm-pair-check adds a prompt of four chunks with a mark at chunk 1, counts the paired steps with the rings' base followed on the host, and compares both twins' bases. --- zig/src/families/glm/cuda_follow.zig | 4 +-- zig/src/families/glm/cuda_prompt_pair.zig | 25 +++++++++++-- zig/tests/cuda/glm_pair_tests.zig | 44 +++++++++++++++-------- 3 files changed, 54 insertions(+), 19 deletions(-) diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index f064c469f6..833ea2c7dd 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -227,9 +227,9 @@ pub const Leader = struct { pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, marks))) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(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); // both chunks of a pair share it + const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // a paired step's chunks share it 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; const first = try x.e.prefillChunk(part, start, ids.len, from, cancel, marks); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index 763bf5bc10..e68367c73e 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -40,11 +40,20 @@ pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool return std.mem.indexOfScalar(u32, marks, @intCast(next)) == null; } +/// Whether a pair at `start` (b of `rb` rows) rebases the index rings where chunks one at a time do: at a, or not. +pub fn ringsAgree(base: usize, ring: usize, start: usize, rb: usize) bool { + const ra = fw.prefill_rows; + return start + ra > base + ring or start + ra + rb <= base + ring; +} + /// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. -pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks) bool { +pub fn pairedAt(e: *const Engine, start: usize, len: usize, from: ?engine.Resume, marks: ?engine.Marks) bool { if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; if (e.pieces.count == 1 and !e.settings.split) return false; // run takes such a chunk whole - return takesTwo(start, len, e.share.world, if (marks) |m| m.at else &.{}); + if (!takesTwo(start, len, e.share.world, 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, @min(len - start - fw.prefill_rows, fw.prefill_rows)); } /// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. @@ -54,6 +63,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: const rb: usize = @min(len - start - ra, fw.prefill_rows); const pair = if (e.pair) |*p| p else return error.BadChunk; if (rb < e.share.world or part.len != @min(len - start, ra + rb + 1) or e.drafter != null or e.prompt_logprobs != null) return error.BadChunk; + if (!ringsAgree(e.cur.st.base, e.cur.st.ring_rows, start, rb)) return error.BadChunk; // pairedAt's rule, every rank if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); e.gate(cancel) catch |err| { if (err == error.Cancelled) try snapshots.interrupted(e, marks, @intCast(start)); @@ -68,7 +78,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: fa.index_split = split; // the step's choice (TF_GLM_INDEX_SPLIT) for both chunks fb.index_split = split; const st = &e.cur.st; - try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base + try st.ready(e.stream, ra + rb); // one rebase for both chunks, at a as a chunk at a time, before b's view copies it var sb = st.*; sb.pos_dev = pair.pos.ptr; try sb.setPos(e.stream, st.pos + ra); @@ -115,3 +125,12 @@ test "a prompt step takes the next chunk when it holds min_rows or more and no k try std.testing.expect(!takesTwo(0, 3 * C, 3, &.{C})); // a state kept at the next chunk's start: that chunk starts a step try std.testing.expect(takesTwo(0, 3 * C, 3, &.{ 2 * C, 3 * C })); // marks past the pair } + +test "a pair rebases the index rings where a chunk at a time would, so a ring based at chunk 0 pairs from chunk 2" { + const C = fw.prefill_rows; + try std.testing.expect(ringsAgree(0, 2 * C, 0, C)); // neither rebases + try std.testing.expect(!ringsAgree(0, 2 * C, C, C)); // a chunk at a time rebases at b's start, a pair at a's + try std.testing.expect(!ringsAgree(0, 2 * C, C, 1)); + try std.testing.expect(ringsAgree(0, 2 * C, 2 * C, C)); // both at a's start + try std.testing.expect(ringsAgree(C, 2 * C, C, C)); +} diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig index ff03559ed8..e339df1626 100644 --- a/zig/tests/cuda/glm_pair_tests.zig +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -23,7 +23,7 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ }; try check.expect(set.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); const C = glm.engine.prompt_rows; - const capacity = @max(dc.capacity, 4 * C); + const capacity = @max(dc.capacity, 5 * 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); @@ -37,24 +37,37 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ defer x.end() catch {}; const V = e.share.vocab.len(); const min = glm.prompt_pair.min_rows; - // a pair and a tail, the shortest second chunk, a pair and a chunk, a pair to the end, a tail too short to pair - const lens = [_]usize{ 2 * C + 300, C + min, 3 * C, 2 * C, C + min - 1 }; + const Case = struct { len: usize, marks: []const u32 = &.{} }; + const cases = [_]Case{ + .{ .len = 2 * C + 300 }, // a pair and a tail + .{ .len = C + min }, // the shortest second chunk + .{ .len = 3 * C }, // a pair and a chunk + .{ .len = 2 * C }, // a pair to the end + .{ .len = C + min - 1 }, // a tail too short to pair + .{ .len = 4 * C, .marks = &.{C} }, // a mark at chunk 1: 0 and 1 alone (a pair at 1 rebases apart), then a pair + }; + 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 (lens, 0..) |len, i| { + 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 paired: [2]u64 = undefined; - for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can + for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can, both with the same marks e.settings.pair = twin == 1; try x.bind(@intCast(2 * i + twin)); try x.begin(null); const before = e.pair.?.steps; var at: usize = 0; const t = while (true) { - const r = x.chunk(ids, at, true, null, null, null) catch |err| { + const r = x.chunk(ids, at, true, null, null, marks) catch |err| { std.debug.print("FAIL pair: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); return err; }; @@ -69,12 +82,17 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(logits[twin].ptr, e.bd.logits, win.len * V * 2), "cuMemcpyDtoH"); } e.settings.pair = true; - var want: u64 = 0; // the steps glm.prompt_pair.takesTwo pairs over this prompt + const st = &e.seqs.items[2 * i].?.st; + const tw = &e.seqs.items[2 * i + 1].?.st; + var want: u64 = 0; // the steps glm.prompt_pair.pairedAt pairs over this prompt, the rings' base followed here + var base: usize = 0; var at: usize = 0; while (at < len) { - const two = glm.prompt_pair.takesTwo(at, len, e.share.world, &.{}); + const two = glm.prompt_pair.takesTwo(at, len, e.share.world, case.marks) and glm.prompt_pair.ringsAgree(base, st.ring_rows, at, @min(len - at - C, C)); + 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 += if (two) 2 * C else C; + at += rows; } if (paired[0] != 0 or paired[1] != want) { std.debug.print("FAIL pair: prompt of {d}: {d} paired steps a chunk a step, {d} paired (want {d})\n", .{ len, paired[0], paired[1], want }); @@ -85,19 +103,17 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ failed = true; } try e.stream.synchronize(); - const st = &e.seqs.items[2 * i].?.st; - const tw = &e.seqs.items[2 * i + 1].?.st; 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 diff = std.mem.indexOfDiff(u8, h0, h1); - if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or !std.mem.eql(u1, st.cur, tw.cur) or diff != null) { - std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, diff, n }); + 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(u1, st.cur, tw.cur) or diff != null) { + std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, rings' base {d}/{d}, MTP base {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, st.base, tw.base, st.mtp_base, tw.mtp_base, diff, n }); failed = true; } } if (failed) return error.TestFailed; - check.pass("paired prompt steps at TP={d}: prompts of {any} rows equal their twins a chunk a step in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); + check.pass("paired prompt steps 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 }); } From 0a3849f5c4088439d4937362fc61127e5346da1c Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:39:21 +0900 Subject: [PATCH 5/9] glm cuda: only an engine with other ranks makes the pair's twin At TP=1 a pair has no exchanges to hide and no receipt, yet the twin was allocated and the side stream ran at top priority. prompt_pair.wanted(on, world) is now on and world > 1, for both; pairedAt already pairs only with a twin. glm-pair-check needs other ranks. --- zig/src/families/glm/cuda_engine.zig | 5 +++-- zig/src/families/glm/cuda_prompt_pair.zig | 11 +++++++++++ zig/tests/cuda/glm_pair_tests.zig | 1 + 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 6bb4ad5fbf..70adda4b60 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -143,7 +143,8 @@ pub const Engine = struct { for (spans, 0..) |*sp, q| sp.* = (try shard.Share.of(&c, q, world)).vocab; var m = try fw.Modules.load(d, ctx); errdefer m.deinit(); - var pieces = try fw.Pieces.init(d, o.settings.pieces, o.settings.pair); + const pair = prompt_pair.wanted(o.settings.pair, world); + var pieces = try fw.Pieces.init(d, o.settings.pieces, pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers); @@ -182,7 +183,7 @@ pub const Engine = struct { 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 }); - if (o.settings.pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); + if (pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); errdefer if (e.pair) |*p| p.deinit(); e.bd = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); errdefer e.bd.deinit(); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index e68367c73e..f6ce94ba55 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -33,6 +33,11 @@ pub const Pair = struct { } }; +/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, and other ranks whose exchanges a pair hides. +pub fn wanted(on: bool, world: usize) bool { + return on and world > 1; +} + /// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool { const next = start + fw.prefill_rows; @@ -134,3 +139,9 @@ test "a pair rebases the index rings where a chunk at a time would, so a ring ba try std.testing.expect(ringsAgree(0, 2 * C, 2 * C, C)); // both at a's start try std.testing.expect(ringsAgree(C, 2 * C, C, C)); } + +test "only ranks that exchange make the twin: TP=1 has no exchanges to hide" { + try std.testing.expect(!wanted(true, 1)); + try std.testing.expect(wanted(true, 2)); + try std.testing.expect(!wanted(false, 3)); +} diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig index e339df1626..633e1d9c6a 100644 --- a/zig/tests/cuda/glm_pair_tests.zig +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -22,6 +22,7 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ return err; }; try check.expect(set.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); + try check.expect(world > 1, "glm-pair-check needs other ranks: at TP=1 an engine makes no twin", .{}); const C = glm.engine.prompt_rows; const capacity = @max(dc.capacity, 5 * 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 }); From 83e43e8298ccfaebeebdec1b135764ae2ab99eda Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:39:21 +0900 Subject: [PATCH 6/9] glm cuda: one table gives the twin's fields their size indexes, and a size added or taken away stops the build twin_sizes tied the fields a twin holds apart to make's size indexes by hand, beside a list of their names. twin_at now maps each field to its index: make reads those fields from it, the twin's allocation and copy take it too, and make's sizes are an array of twin_checked entries, so a future size entry fails the build until the indexes are checked again. --- zig/src/families/glm/cuda_state.zig | 45 +++++++++++++++++------------ 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 2a0214987d..8d69f28102 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -180,9 +180,16 @@ pub const Buffers = struct { return make(d, w, of.rows, of.prefill, capacity, of.rows, of); } - /// The fields a twin holds apart: make's sizes by index, then their names. - const twin_sizes = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7, 62, 63, 64, 65, 70 }; - const twin_fields = .{ "ids", "x", "normed", "xs", "post", "comb", "hcpart", "gathered", "hidden", "fnormed", "fxs", "logits", "part" }; + /// The fields a twin holds apart, by their index in make's sizes (make reads the fields from these indexes too). + const twin_at = .{ .ids = 0, .x = 1, .normed = 2, .xs = 3, .post = 4, .comb = 5, .hcpart = 6, .gathered = 7, .hidden = 62, .fnormed = 63, .fxs = 64, .logits = 65, .part = 70 }; + /// make's size count when twin_at's indexes were last checked: adding or removing a size fails the build. + const twin_checked = 76; + + /// Whether make's size `i` is one a twin holds apart. + fn twinHolds(i: usize) bool { + inline for (@typeInfo(@TypeOf(twin_at)).@"struct".field_names) |f| if (@field(twin_at, f) == i) return true; + return false; + } /// `seg`: the most rows one stream's attention takes (a decode window's rows, or a shared window's segment). fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize, twin: ?*const Buffers) !Buffers { @@ -212,7 +219,7 @@ pub const Buffers = struct { var split = false; // a KDA stack in two parts moves each part's rows through one more buffer for (w.layers) |l| split = split or (l.mixer == .kda and l.mixer.kda.proj == .split); const hist_len = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (count + 1) else 1; - const sizes = [_]usize{ + const sizes: [twin_checked]usize = .{ rows * 4, rows * S * D * 2, rows * D * 2, rows * (D / 64) * 4, rows * S * 4, rows * S * S * 4, glue.hcPartBytes(rows), s.world * rows * D * 4, if (prefill or seg < rows) rows * kw * 2 else 4, rows * n * 2, rows * n * 2, // kproj, ka, kg @@ -253,7 +260,7 @@ pub const Buffers = struct { }; var held = sizes; if (twin != null) for (&held, 0..) |*bytes, i| { // a twin allocates its own rows only - if (std.mem.indexOfScalar(usize, &twin_sizes, i) == null) bytes.* = 0; + if (!twinHolds(i)) bytes.* = 0; }; var p: [sizes.len]u64 = undefined; const a = try arena(d, &held, &p); @@ -263,14 +270,14 @@ pub const Buffers = struct { .prefill = prefill, .part_rows = part_rows, .nch = nch, - .ids = p[0], - .x = p[1], - .normed = p[2], - .xs = p[3], - .post = p[4], - .comb = p[5], - .hcpart = p[6], - .gathered = p[7], + .ids = p[twin_at.ids], + .x = p[twin_at.x], + .normed = p[twin_at.normed], + .xs = p[twin_at.xs], + .post = p[twin_at.post], + .comb = p[twin_at.comb], + .hcpart = p[twin_at.hcpart], + .gathered = p[twin_at.gathered], .kproj = p[8], .kda = .{ .p = 0, .ka = p[9], .kg = p[10], .saves = .{ .out = p[11], .state_out = 0, .k = p[12], .v = p[13], .g = p[14], .b = p[15] }, .wide = .{ .q = p[16], .y = p[17] }, .mm_part = p[42] }, .dsa = .{ @@ -297,22 +304,22 @@ pub const Buffers = struct { }, .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47], .mm_part = p[42] }, .moe = .{ .mlog = p[48], .router_part = p[26], .pick = p[49], .wts = p[50], .plan = .{ .members = p[51], .items = p[52], .counts = p[53], .rank = p[54], .hist = p[55] }, .eact = p[56], .ey = p[57], .sgu = p[58], .sact = p[59], .sxs = p[60], .sy = p[61], .mm_part = p[42] }, - .hidden = p[62], - .fnormed = p[63], - .fxs = p[64], - .logits = p[65], + .hidden = p[twin_at.hidden], + .fnormed = p[twin_at.fnormed], + .fxs = p[twin_at.fxs], + .logits = p[twin_at.logits], .hin = p[66], .me = p[67], .mcat = p[68], .mx = p[69], - .part = p[70], + .part = p[twin_at.part], }; out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; out.kda.tmp = p[75]; if (twin) |of| { var t = of.*; t.arena = out.arena; - inline for (twin_fields) |f| @field(t, f) = @field(out, f); + inline for (@typeInfo(@TypeOf(twin_at)).@"struct".field_names) |f| @field(t, f) = @field(out, f); return t; } return out; From a3d09d0eedb969f667cf0f6f931ddd08d7cf9f35 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:47:40 +0900 Subject: [PATCH 7/9] glm cuda: glm-split-check's list of prompts is described on one line --- zig/tests/cuda/glm_split_tests.zig | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zig/tests/cuda/glm_split_tests.zig b/zig/tests/cuda/glm_split_tests.zig index 46c490c881..cb1233200f 100644 --- a/zig/tests/cuda/glm_split_tests.zig +++ b/zig/tests/cuda/glm_split_tests.zig @@ -37,8 +37,7 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c defer x.end() catch {}; const V = e.share.vocab.len(); const from = C; // every step from the second chunk on: each holds rows past the dense limit - // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks, and a second step - // that pairs past `from` (Leader.chunk pairs where it can; a paired step's two chunks take the split of its start) + // tails of 300, 2 (one rank without a row) and 5, two whole chunks, and a second step paired past `from` const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C, 3 * C + glm.prompt_pair.min_rows }; var failed = false; var paired: u64 = 0; // the split twins' paired steps From a4c867e15e088d6e5a83e2b9b49dc283561b00e8 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:37:03 +0900 Subject: [PATCH 8/9] glm cuda: the pair's choice and its twin's pieces sit with the pair and Cancel with the engine's types, so the engine and the lanes keep room under the line limit --- zig/src/families/glm/cuda_engine.zig | 8 ++------ zig/src/families/glm/cuda_engine_types.zig | 3 +++ zig/src/families/glm/cuda_follow.zig | 3 ++- zig/src/families/glm/cuda_prompt_pair.zig | 7 +++++++ 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 70adda4b60..041342b976 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -21,9 +21,6 @@ const pages_mod = @import("cuda_pages.zig"); const prompt_pair = @import("cuda_prompt_pair.zig"); pub const Options = @import("cuda_engine_types.zig").Options; -/// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. -pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; - /// The most drafts a chain holds (decode.DepthPolicy's top depth). pub const max_drafts = 3; /// A greedy pick as it crosses ranks: (value, id) and padding, so the RoCE gather (16-byte packs) can carry it. @@ -47,6 +44,7 @@ pub const Kept = types.Kept; const Seq = types.Seq; const Snap = types.Snap; pub const Resume = types.Resume; +pub const Cancel = types.Cancel; pub const Marks = types.Marks; pub const Engine = struct { @@ -255,8 +253,7 @@ pub const Engine = struct { pub fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; - const twin = if (e.pair) |*p| (if (b == &p.bp) &p.pieces else null) else null; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else twin, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null }; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = prompt_pair.piecesOf(e, b), .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -521,7 +518,6 @@ pub const Engine = struct { } pub const holdPrompt = @import("cuda_prompt_hold.zig").save; - pub const pairedAt = prompt_pair.pairedAt; pub const prefillMany = @import("cuda_prompts.zig").prefillMany; /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran). diff --git a/zig/src/families/glm/cuda_engine_types.zig b/zig/src/families/glm/cuda_engine_types.zig index dded605953..125a6ee1ed 100644 --- a/zig/src/families/glm/cuda_engine_types.zig +++ b/zig/src/families/glm/cuda_engine_types.zig @@ -3,6 +3,9 @@ const cuda = @import("cuda"); const fw = @import("cuda_forward.zig"); const sampler = @import("cuda_sampler.zig"); +/// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. +pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; + /// One stream's window of a shared verify: its slot, its ids, and its drawn tokens. pub const Part = struct { slot: u32, ids: []const u32, out: []u32, bits: u64 = 0 }; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 833ea2c7dd..b6b2827472 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -11,6 +11,7 @@ const prompt_wire = @import("cuda_prompt_wire.zig"); const dflash = @import("dflash.zig"); const pages_mod = @import("cuda_pages.zig"); const sparse = @import("cuda_sparse.zig"); +const prompt_pair = @import("cuda_prompt_pair.zig"); const Engine = engine.Engine; pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw, idle, snapshots, pool }; @@ -227,7 +228,7 @@ pub const Leader = struct { pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, from, marks))) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or prompt_pair.pairedAt(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); // a paired step's chunks share it 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); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index f6ce94ba55..d35f77b5e8 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -51,6 +51,13 @@ pub fn ringsAgree(base: usize, ring: usize, start: usize, rb: usize) bool { return start + ra > base + ring or start + ra + rb <= base + ring; } +/// The pieces a forward over `b` exchanges in: the engine's, or the pair's for its twin; none for other buffers. +pub fn piecesOf(e: *const Engine, b: *const fw.Buffers) ?*const fw.Pieces { + if (b == &e.bp) return &e.pieces; + const p = if (e.pair) |*p| p else return null; + return if (b == &p.bp) &p.pieces else null; +} + /// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. pub fn pairedAt(e: *const Engine, start: usize, len: usize, from: ?engine.Resume, marks: ?engine.Marks) bool { if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; From 61d2b808b68cf654840b01752d9fa1ff8b1ceeec Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 20:19:00 +0900 Subject: [PATCH 9/9] glm cuda: TF_GLM_CHUNK_SOLO=4096 takes the pair's place, and its server holds no twin A solo pass and a paired step both reach the other ranks as a part of two chunks, so one server cannot mix them. With the wider buffers set, the pair stands aside and makes no twin. --- docs/recipes/glm-5.3-flash-cuda.md | 3 ++- zig/src/families/glm/cuda_engine.zig | 2 +- zig/src/families/glm/cuda_follow.zig | 2 +- zig/src/families/glm/cuda_lanes_test.zig | 3 +-- zig/src/families/glm/cuda_prompt_pair.zig | 17 +++++++++++------ 5 files changed, 16 insertions(+), 11 deletions(-) diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 6fb85e59bf..c9368351ec 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -69,7 +69,8 @@ A prompt runs in passes of 2,048 rows. Kept prompt states sit on that grid, and 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 `prompt_logprobs` or image rows, the pass takes 2,048. A drafter's steps take one -chunk, so `--drafter` beside it is refused at load. +chunk, so `--drafter` beside it is refused at load. It takes the place of `TF_GLM_PROMPT_PAIR`'s paired steps, whose +second set of buffers it does not hold. 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: diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 041342b976..4c8c705f94 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -141,7 +141,7 @@ pub const Engine = struct { for (spans, 0..) |*sp, q| sp.* = (try shard.Share.of(&c, q, world)).vocab; var m = try fw.Modules.load(d, ctx); errdefer m.deinit(); - const pair = prompt_pair.wanted(o.settings.pair, world); + const pair = prompt_pair.wanted(o.settings.pair, o.settings.solo, world); var pieces = try fw.Pieces.init(d, o.settings.pieces, pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index b6b2827472..862e3394f8 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -224,7 +224,7 @@ pub const Leader = struct { return x.e.prefill(ids, from, cancel, marks); } - /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and prompt_solo.at or the pair's rule allow. + /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and solo or pair rules allow. pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig index 1eca112ac8..af503aff49 100644 --- a/zig/src/families/glm/cuda_lanes_test.zig +++ b/zig/src/families/glm/cuda_lanes_test.zig @@ -1,8 +1,7 @@ -//! The lane backend's slots, chosen and given back, and when a prompt step may pair. +//! The lane backend's slots: chosen before a stream's map entry exists, and none left behind by a failed insertion. const std = @import("std"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); -const fw = @import("cuda_forward.zig"); const Lanes = @import("cuda_lanes.zig").Lanes; const fw = @import("cuda_forward.zig"); const solo = @import("cuda_prompt_solo.zig"); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index d35f77b5e8..c864e93338 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -33,9 +33,9 @@ pub const Pair = struct { } }; -/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, and other ranks whose exchanges a pair hides. -pub fn wanted(on: bool, world: usize) bool { - return on and world > 1; +/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, no solo passes, and ranks to hide exchanges from. +pub fn wanted(on: bool, solo: bool, world: usize) bool { + return on and !solo and world > 1; } /// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. @@ -148,7 +148,12 @@ test "a pair rebases the index rings where a chunk at a time would, so a ring ba } test "only ranks that exchange make the twin: TP=1 has no exchanges to hide" { - try std.testing.expect(!wanted(true, 1)); - try std.testing.expect(wanted(true, 2)); - try std.testing.expect(!wanted(false, 3)); + try std.testing.expect(!wanted(true, false, 1)); + try std.testing.expect(wanted(true, false, 2)); + try std.testing.expect(!wanted(false, false, 3)); +} + +test "a server with solo passes holds no twin: they take the pair's steps" { + try std.testing.expect(!wanted(true, true, 2)); + try std.testing.expect(!wanted(true, true, 3)); }