glm cuda: a prompt step takes two chunks while no other stream decodes, each one's exchanges under the other's sites

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:

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;

zig/src/families/glm/cuda_engine.zig

@@ -18,11 +18,9 @@ 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.
-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.
@@ -46,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 {
@@ -71,6 +70,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 +141,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);
+        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;
         const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers);
@@ -180,6 +181,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 (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 +230,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 +250,10 @@ 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 };
+        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.
@@ -342,6 +346,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, 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;

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

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 };
@@ -223,13 +224,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 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();
-        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 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);
+        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);

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

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