GLM-5.3-Flash on two CUDA ranks: prompts admitted together (prefill_many)

zig/src/core/lane_host_admit.zig

@@ -23,9 +23,10 @@ fn fits(a: api.Admit, job: *const Job) bool {
     return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens);
 }
 
-/// Whether to admit several at once: prompts batch, none fill by chunks, and two or more wait with two lanes free.
+/// Whether to admit several at once: prompts batch (beside chunked fills up to many_rows), two wait, two lanes free.
 pub fn batches(h: *LaneHost) bool {
-    if (!h.core.batchesPrompts() or h.core.fills()) return false;
+    if (!h.core.batchesPrompts()) return false;
+    if (h.core.fills() and h.core.backend.vtable.many_rows == 0) return false;
     h.lock();
     defer h.unlock();
     const lanes_now = @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes);
@@ -37,6 +38,7 @@ pub fn many(h: *LaneHost) bool {
     var jobs: [lanes.Engine.max_together]*Job = undefined;
     var began: [lanes.Engine.max_together]i96 = undefined;
     var n: usize = 0;
+    var rows: usize = 0; // with chunked fills a batch is one pass of many_rows: decode rounds run between passes
     var took = false;
     while (n < jobs.len) {
         const job = switch (h.takeJob()) {
@@ -49,10 +51,13 @@ pub fn many(h: *LaneHost) bool {
         };
         took = true;
         const at = h.setup(job) orelse continue;
-        if (h.startsAlone(job)) {
+        const fills = h.core.fills();
+        const over = fills and rows + job.stream.prompt_len > h.core.backend.vtable.many_rows;
+        if (h.startsAlone(job) or (fills and (!h.core.batchable(&job.stream) or over))) { // long, image, resumed, overflow: filled
             _ = h.launch(job, at);
             continue;
         }
+        rows += job.stream.prompt_len;
         jobs[n] = job;
         began[n] = at;
         n += 1;

zig/src/core/lane_host_many_test.zig

@@ -49,13 +49,26 @@ fn expectOwn(prompt: []const u32, tokens: []const u32) !void {
 }
 
 /// `n` requests queued before the host starts (`sampled`: one the fake refuses), replies checked, batched passes.
-fn serve(n: usize, memory_lanes: ?usize, sampled: ?usize) !struct { passes: usize, streams: usize } {
+const Passes = struct { passes: usize, streams: usize };
+
+fn serve(n: usize, memory_lanes: ?usize, sampled: ?usize) !Passes {
+    return serveOn(n, memory_lanes, sampled, .whole);
+}
+
+const Mode = enum { whole, steps, tight };
+
+fn serveOn(n: usize, memory_lanes: ?usize, sampled: ?usize, mode: Mode) !Passes {
     var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7);
     defer cfg.deinit(gpa);
     var target: lanes.fake.Fake = .{ .gpa = gpa, .memory_lanes = memory_lanes, .refuse_sampled = sampled != null };
     defer target.deinit();
     var clock: lanes.fake.FixedClock = .{};
-    var core = lanes.Engine.init(gpa, &cfg, target.batched(), clock.clock());
+    const backend = switch (mode) {
+        .whole => target.batched(),
+        .steps => target.batchedSteps(),
+        .tight => target.batchedStepsTight(),
+    };
+    var core = lanes.Engine.init(gpa, &cfg, backend, clock.clock());
     defer core.deinit();
     var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 4 });
     var prompts: [6][5]u32 = undefined;
@@ -89,6 +102,18 @@ test "queued requests share one prompt pass, each reply the target's own continu
     try std.testing.expectEqual(@as(usize, 4), r.streams);
 }
 
+test "a backend that also fills by chunks: short queued prompts still share one pass" {
+    const r = try serveOn(4, null, null, .steps);
+    try std.testing.expectEqual(@as(usize, 1), r.passes);
+    try std.testing.expectEqual(@as(usize, 4), r.streams);
+}
+
+test "beside chunked fills a batch is one pass: prompts past its rows fill by chunks, every reply its own" {
+    const r = try serveOn(4, null, null, .tight); // 5-token prompts, a 12-token pass: two batch, two fill
+    try std.testing.expectEqual(@as(usize, 1), r.passes);
+    try std.testing.expectEqual(@as(usize, 2), r.streams);
+}
+
 test "more requests than lanes: the free lanes' prompts in one pass, the rest after" {
     const r = try serve(6, null, null);
     try std.testing.expect(r.passes >= 1);

zig/src/core/lanes/backend.zig

@@ -79,6 +79,8 @@ pub const Backend = struct {
         prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null,
         /// Several prompts in one pass, each as after `prefill`; false or an error leaves none prefilled.
         prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null,
+        /// Prompts of at most this many tokens prefill_many takes while prompts otherwise fill by chunks (0: none).
+        many_rows: u32 = 0,
         /// Draw the first token from the prompt's last row at `position`; a handle to it.
         first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64,
         /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it.

zig/src/core/lanes/engine.zig

@@ -94,6 +94,7 @@ pub const Engine = struct {
     pub const addStream = open.add;
     pub const addStreamCommitted = open.addCommitted;
     pub const addStreams = open.addStreams;
+    pub const batchable = open.batchable;
     pub const max_together = open.max_together;
 
     /// Whether the backend prefills several prompts in one pass (addStreams batches them).

zig/src/core/lanes/fake.zig

@@ -123,6 +123,16 @@ pub const Fake = struct {
         return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } };
     }
 
+    /// Both: prompts a chunk a call, and short ones several in one pass of up to 32 tokens (`many_rows`).
+    pub fn batchedSteps(x: *Fake) be.Backend {
+        return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .prefill_many = prefillMany, .many_rows = 32, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } };
+    }
+
+    /// `batchedSteps` with a pass of 12 tokens: two 5-token prompts fill it, a third fills by chunks.
+    pub fn batchedStepsTight(x: *Fake) be.Backend {
+        return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .prefill_many = prefillMany, .many_rows = 12, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } };
+    }
+
     /// The backend with prompts in pieces (begin / piece / finish) beside whole prompts.
     pub fn pieceBackend(x: *Fake) be.Backend {
         return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish } };

zig/src/core/lanes/open.zig

@@ -1,4 +1,5 @@
 //! A stream's way into the rounds: its prompt pass whole or a chunk a round, then its first token and drafts.
+const std = @import("std");
 const Engine = @import("engine.zig").Engine;
 const sm = @import("stream.zig");
 const be = @import("backend.zig");
@@ -40,13 +41,22 @@ pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void {
     if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs);
     _ = e.arena.reset(.retain_capacity);
     for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {};
-    const together = many(e.backend.ptr, ss) catch false; // one stream's refusal is its own: each goes alone
+    const together = many(e.backend.ptr, ss) catch |err| blk: { // one stream's refusal is its own: each goes alone
+        std.log.warn("{d} prompts together: {s}; each goes alone", .{ ss.len, @errorName(err) });
+        break :blk false;
+    };
     if (!together) return addEach(e, ss, errs);
     for (ss, errs) |s, *x| opened(e, s) catch |err| { // a stream cancelled in the pass gets no first token
         x.* = err;
     };
 }
 
+/// Whether `s` may join a batched prompt pass beside chunked fills: short, text only, nothing kept to resume from.
+pub fn batchable(e: *const Engine, s: *const Stream) bool {
+    const n = e.backend.vtable.many_rows;
+    return n > 0 and s.prompt_len <= n and s.images.len == 0 and s.reuse.saved == null;
+}
+
 fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void {
     for (ss, errs) |s, *x| add(e, s) catch |err| {
         x.* = err;

zig/src/families/glm/cuda_b16.zig

@@ -58,7 +58,7 @@ pub fn matmul(set: glue.Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u6
     if (k % bk != 0) return error.KNotMultipleOf64;
     if (set.mm) |mk| { // tiled weights, packed or bf16
         const t: mm.Tiled = if (set.packs) |p| p.tiled(w, n, k) else .{ .w = w, .n = n, .k = k };
-        return mm.matmul(mk, stream, x, x_stride, t, out, part, m, wide);
+        return mm.matmul(mk, stream, x, x_stride, t, out, part, m, wide, set.prompt);
     }
     const sk = splitK(n, k);
     var a: cuda.Args = .{};

zig/src/families/glm/cuda_engine.zig

@@ -68,7 +68,7 @@ pub const Engine = struct {
     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)
-    hold: cuda.DeviceBuffer, // the prompt's last row's hidden (decode's e.last_hidden), the first draft's source
+    hold: cuda.DeviceBuffer, // prompts' last rows' hidden (decode's e.last_hidden), a row a prompt pass's prompt
     pairs: cuda.DeviceBuffer, // every rank's (value, id) of a greedy draw
     smp: sampler.Sampler,
     seqs: std.ArrayList(?*Seq) = .empty, // every stream's sequence by slot (the lane backend's), slot 0 made at init
@@ -168,7 +168,7 @@ pub const Engine = struct {
         errdefer e.bm.deinit();
         if (o.mtp) e.bs = try fw.Buffers.initShared(d, &e.w, shared_rows, o.capacity);
         errdefer if (e.bs) |*b| b.deinit();
-        e.hold = try cuda.DeviceBuffer.alloc(d, c.hidden * 2);
+        e.hold = try cuda.DeviceBuffer.alloc(d, types.max_together * c.hidden * 2);
         errdefer e.hold.free();
         e.pairs = try cuda.DeviceBuffer.alloc(d, pair_bytes * world);
         errdefer e.pairs.free();
@@ -230,7 +230,7 @@ pub const Engine = struct {
         }
     }
 
-    fn exchange(e: *const Engine) fw.Exchange {
+    pub fn exchange(e: *const Engine) fw.Exchange {
         return .{ .nccl = .{ .lib = e.lib, .comm = e.comm, .group = e.group } };
     }
 
@@ -339,7 +339,7 @@ pub const Engine = struct {
         }
         if (f.prof) |p| p.mark(e.stream, .draw);
         if (e.mtp) {
-            if (last) try e.copy(e.hold.ptr, e.bp.fnormed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2);
+            if (last) try @import("cuda_prompts.zig").hold(e, 0, e.bp.fnormed + (rows - 1) * e.c.hidden * 2);
             const nxt = as_i32[1..];
             if (nxt.len > 0) { // decode._absorb_rows
                 const mp = f.prof;
@@ -356,7 +356,7 @@ pub const Engine = struct {
     }
 
     /// The chunk gate on every rank: a kept state at the mark, the heat wait, rank 0's cancel; the id exchanged first.
-    fn gate(e: *Engine, cancel: ?Cancel) !void {
+    pub fn gate(e: *Engine, cancel: ?Cancel) !void {
         const kept = try e.shared(@floatFromInt(e.announce));
         e.announce = 0;
         if (kept > 0 and e.share.rank != 0) try e.keepAs(@intFromFloat(kept));
@@ -515,6 +515,7 @@ pub const Engine = struct {
         try dr.draft(e, reqs);
     }
 
+    pub const prefillMany = @import("cuda_prompts.zig").prefillMany; // whole prompts in one pass
     /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran).
     pub fn keep(e: *Engine, n: usize) !void {
         var one: [1]fw.Rank = undefined;
@@ -529,7 +530,7 @@ pub const Engine = struct {
         var one: [1]fw.Rank = undefined;
         const f = e.view(&e.bm, &one);
         const D = e.c.hidden;
-        const hidden = if (from_prompt) e.hold.ptr else e.cur.rows_at;
+        const hidden = if (from_prompt) e.cur.held_at else e.cur.rows_at;
         try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len - e.cur.in_head);
         e.cur.in_head = 0;
         const next: []const i32 = @ptrCast(follow);
@@ -550,7 +551,7 @@ pub const Engine = struct {
     }
 
     /// `rows` logits rows of this rank's vocabulary -> the tokens at positions first, first + 1, ...
-    fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void {
+    pub fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void {
         if (e.cur.sampling) |s| {
             try e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out);
             return e.checkGathers(); // the draw read rows the forward's gathers made
@@ -593,7 +594,7 @@ pub const Engine = struct {
         return if (s.len == 0) heat.default_root else s;
     }
 
-    fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void {
+    pub fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void {
         try e.d.check(e.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, e.stream.handle), "cuMemcpyDtoDAsync");
     }
 };

zig/src/families/glm/cuda_engine_types.zig

@@ -15,6 +15,7 @@ pub const Seq = struct {
     sampling: ?sampler.Sampling = null, // null: greedy
     in_head: usize = 0, // drafted rows the MTP caches hold past the kept ones (st.mtp_drafted)
     rows_at: u64 = 0, // its last verify's final-norm rows (bf16 [rows][hidden]), which the next draft reads
+    held_at: u64 = 0, // its prompt's last row's hidden (a row of `hold`), which the first draft reads
 };
 
 /// One rank's copy of a kept prompt state: the state arena's bytes and the host fields beside them.
@@ -36,3 +37,8 @@ pub const Options = struct {
     vision: ?images_mod.Limits = null, // --vision's limits (null: off): every rank must hold the rows rank 0 sends
     drafter: bool = false, // DFlash2 loads after rank agreement
 };
+
+/// Whole prompts one prompt pass takes together (cuda_prompts).
+pub const max_together = @import("cuda_state.zig").prompt_heads;
+/// One of them: its stream's slot and its ids.
+pub const Prompt = struct { slot: u32, ids: []const u32 };

zig/src/families/glm/cuda_follow.zig

@@ -3,13 +3,15 @@
 const std = @import("std");
 const cuda = @import("cuda");
 const engine = @import("cuda_engine.zig");
+const types = @import("cuda_engine_types.zig");
 const images_mod = @import("cuda_images.zig");
 const vcfg = @import("vision_config.zig");
 const sampler = @import("cuda_sampler.zig");
 const dflash = @import("dflash.zig");
+const prompts = @import("cuda_prompts.zig");
 const Engine = engine.Engine;
 
-pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block };
+pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many };
 
 /// One step as it crosses: 64 bytes, then `n` u32 ids.
 pub const Header = extern struct {
@@ -214,6 +216,25 @@ pub const Leader = struct {
         try x.e.verifyShared(parts);
     }
 
+    /// Whole prompts in one pass on every rank (Engine.prefillMany): slots, lengths, ids; hold0 (field a).
+    pub fn prefillMany(x: Leader, parts: []const types.Prompt, firsts: []u32, hold0: usize) !void {
+        if (!prompts.fits(parts, hold0, x.e.capacity)) return error.BadPrompts;
+        var ids: usize = 2 * parts.len;
+        for (parts) |p| ids += p.ids.len;
+        if (ids > x.link.ids.len) return error.BadPrompts; // refused before any broadcast: each prompt then goes alone
+        try x.flush();
+        var buf: [2 * types.max_together + engine.prompt_rows]u32 = undefined;
+        var n: usize = 2 * parts.len;
+        for (parts, 0..) |p, i| {
+            buf[2 * i] = p.slot;
+            buf[2 * i + 1] = @intCast(p.ids.len);
+            @memcpy(buf[n..][0..p.ids.len], p.ids);
+            n += p.ids.len;
+        }
+        try x.send(.{ .kind = @backingInt(Kind.many), .depth = @intCast(parts.len), .a = hold0 }, buf[0..n]);
+        try x.e.prefillMany(parts, firsts, hold0);
+    }
+
     pub fn keep(x: Leader, n: usize) !void {
         try x.flush();
         try x.send(.{ .kind = @backingInt(Kind.keep), .a = n }, &.{});
@@ -250,7 +271,7 @@ pub fn follow(e: *Engine, link: *Link) !void {
         const step = try link.receive();
         const h = step.h;
         const kind: Kind = @fromBackingInt(@intCast(h.kind));
-        if (kind != .end and kind != .drop and kind != .shared and kind != .block) try e.bind(h.slot);
+        if (kind != .end and kind != .drop and kind != .shared and kind != .block and kind != .many) try e.bind(h.slot);
         switch (kind) {
             .begin => try e.beginFollowing(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null),
             .prefill => {
@@ -288,6 +309,20 @@ pub fn follow(e: *Engine, link: *Link) !void {
                 }
                 try e.verifyShared(parts[0..k]);
             },
+            .many => {
+                const k = h.depth;
+                if (k == 0 or k > types.max_together or step.ids.len < 2 * k) return error.BadStep;
+                var parts: [types.max_together]types.Prompt = undefined;
+                var at: usize = 2 * k;
+                for (parts[0..k], 0..) |*p, i| {
+                    const n = step.ids[2 * i + 1];
+                    if (at + n > step.ids.len) return error.BadStep;
+                    p.* = .{ .slot = step.ids[2 * i], .ids = step.ids[at..][0..n] };
+                    at += n;
+                }
+                var firsts: [types.max_together]u32 = undefined;
+                try e.prefillMany(parts[0..k], &firsts, @intCast(h.a));
+            },
             .keep => try e.keep(@intCast(h.a)),
             .draft => {
                 const ids = try e.gpa.dupe(u32, step.ids);

zig/src/families/glm/cuda_forward.zig

@@ -184,7 +184,7 @@ pub const Forward = struct {
     taps: ?taps_mod.Taps = null, // DFlash2 (--drafter): the tapped layers' rows, after each such layer
 
     pub fn set(f: *const Forward) glue.Set {
-        return .{ .glue = &f.m.glue, .mm = if (f.m.tiled) &f.m.mm else null, .packs = f.m.packs };
+        return .{ .glue = &f.m.glue, .mm = if (f.m.tiled) &f.m.mm else null, .packs = f.m.packs, .prompt = f.ranks[0].b.prefill };
     }
 
     /// reduce.sums' gather: every rank's partial of `rows` rows into every rank's gathered buffer, rank order.
@@ -435,7 +435,7 @@ pub const Forward = struct {
         var li: usize = 0;
         var di: usize = 0;
         for (layers, 0..) |_, j| {
-            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);
+            for (f.ranks) |r| if (r.segs.len > 0) try shared.mixer(f, r, r.w.layers[j], li, di, R, null) else try f.mixer(r, r.w.layers[j], li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), null);
             try f.exchangePieces(p, j, .mixer, cut);
             try f.gluePieces(p, j, .mixer, cut);
             for (f.ranks) |r| try f.ffn(r, r.w.layers[j], R, null);
@@ -512,6 +512,10 @@ pub const Forward = struct {
         for (f.ranks) |r| {
             const b = r.b;
             try glue.rmsnorm(f.set(), f.stream, b.hidden, c.hidden, r.w.norm, b.fnormed, c.hidden, b.fxs, c.eps, R, c.hidden);
+            if (b.prefill and r.segs.len > 0) {
+                try shared.heads(f, r);
+                continue;
+            }
             const head_rows: usize = if (b.prefill) 1 else R;
             const first = if (b.prefill) R - 1 else 0;
             try b16.matmul(f.set(), f.stream, b.fnormed + first * c.hidden * 2, c.hidden, r.w.head, b.logits, b.kda.mm_part, head_rows, r.w.share.vocab.len(), c.hidden, false);