GLM-5.3-Flash on CUDA: prompt_logprobs on /v1/completions

docs/api.md

@@ -148,7 +148,8 @@ sequences. `/v1/completions` takes the legacy `logprobs`, an integer from 0 thro
 the legacy form, streamed or not: `tokens`, `token_logprobs`, `top_logprobs` (each alternative's text to its logprob)
 and `text_offset` (each token's first character in the completion, counting each token's own text). Other engines,
 the Responses route and the request modes above return HTTP 400 when probabilities are requested.
-Nonstreamed `/v1/completions` also takes vLLM's `prompt_logprobs`, an integer from 0 through 20, on the same engine.
+Nonstreamed `/v1/completions` also takes vLLM's `prompt_logprobs`, an integer from 0 through 20, on the same engine
+and on GLM-5.3-Flash's CUDA engine.
 `choices[0].prompt_logprobs` has one entry per prompt token: `null` for the first, then an object keyed by token id
 (a string) whose first key is the prompt token itself, followed by the requested best tokens; each value has
 `logprob`, `rank` and `decoded_token`. The prompt token's `rank` counts every vocabulary token at or above it,

zig/src/families/glm/cuda.zig

@@ -18,6 +18,7 @@ pub const quant = @import("cuda_quant.zig");
 pub const q4 = @import("cuda_q4.zig");
 pub const ranks = @import("cuda_ranks.zig");
 pub const sampler = @import("cuda_sampler.zig");
+pub const prompt_logprobs = @import("cuda_prompt_logprobs.zig");
 pub const engine = @import("cuda_engine.zig");
 pub const follow = @import("cuda_follow.zig");
 pub const lanes = @import("cuda_lanes.zig");
@@ -41,6 +42,7 @@ test {
     _ = q4;
     _ = ranks;
     _ = sampler;
+    _ = prompt_logprobs;
     _ = engine;
     _ = follow;
     _ = lanes;

zig/src/families/glm/cuda_engine.zig

@@ -9,6 +9,7 @@ const weights = @import("cuda_weights.zig");
 const fw = @import("cuda_forward.zig");
 const ranks = @import("cuda_ranks.zig");
 const sampler = @import("cuda_sampler.zig");
+const prompt_lp = @import("cuda_prompt_logprobs.zig");
 const heat = @import("heat");
 const modelopt = @import("core").modelopt;
 
@@ -86,6 +87,10 @@ pub const Engine = struct {
     next_snap: u64 = 1, // rank 0: the next id
     announce: u64 = 0, // rank 0: a state kept at this chunk's mark, for the other ranks to copy at its gate
     dropped: std.ArrayList(u32) = .empty, // rank 0: ids dropped since the other ranks last heard (cuda_follow)
+    prompt_logprobs: ?prompt_lp.Want = null, // a prompt pass's rows asked for (every rank runs the head on them)
+    prompt_bufs: ?prompt_lp.Buffers = null, // their head passes' buffers, made at the first pass that asks
+    prompt_tried: bool = false, // that pass came: without buffers after it, the passes take the window's
+    prompt_out: cuda.HostBuffer, // a head pass's logprob words, then its picks, host-mapped (rank 0 reads them)
 
     /// Opens rank `rank` of `world` over `comm`; the ranks agree on settings (cuda_ranks) before the weights load.
     pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, ctx: *const cuda.Context, checkpoint: []const u8, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, rank: usize, world: usize, o: Options) !*Engine {
@@ -129,7 +134,7 @@ pub const Engine = struct {
         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);
         for (layers, 0..) |*l, i| l.* = @intCast(i);
-        e.* = .{ .gpa = gpa, .io = io, .heat_gate = .{ .bands = o.settings.heat, .root = heatRoot() }, .shared_buf = undefined, .d = d, .c = c, .share = share, .spans = spans, .lib = lib, .comm = comm, .settings = o.settings, .m = m, .stream = stream, .pieces = pieces, .w = undefined, .capacity = o.capacity, .bp = undefined, .bd = undefined, .bm = undefined, .hold = undefined, .pairs = undefined, .smp = undefined, .mtp = o.mtp };
+        e.* = .{ .gpa = gpa, .io = io, .heat_gate = .{ .bands = o.settings.heat, .root = heatRoot() }, .shared_buf = undefined, .d = d, .c = c, .share = share, .spans = spans, .lib = lib, .comm = comm, .settings = o.settings, .m = m, .stream = stream, .pieces = pieces, .w = undefined, .capacity = o.capacity, .bp = undefined, .bd = undefined, .bm = undefined, .hold = undefined, .pairs = undefined, .smp = undefined, .prompt_out = undefined, .mtp = o.mtp };
         e.w = try weights.load(gpa, io, d, stream, e.m.tools(), checkpoint, c, share, layers, o.mtp);
         errdefer e.w.deinit();
         try e.bind(0);
@@ -151,6 +156,8 @@ pub const Engine = struct {
         errdefer e.pairs.free();
         e.smp = try sampler.Sampler.init(d, c.vocab, fw.window_rows);
         errdefer e.smp.deinit();
+        e.prompt_out = try cuda.HostBuffer.allocMapped(d, prompt_lp.out_bytes);
+        errdefer e.prompt_out.free();
         e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8);
         try e.stream.synchronize();
         return e;
@@ -163,6 +170,8 @@ pub const Engine = struct {
         e.snaps.deinit(e.gpa);
         e.dropped.deinit(e.gpa);
         e.shared_buf.free();
+        if (e.prompt_bufs) |*b| b.free();
+        e.prompt_out.free();
         e.smp.deinit();
         e.pairs.free();
         e.hold.free();
@@ -181,7 +190,7 @@ pub const Engine = struct {
         gpa.destroy(e);
     }
 
-    fn exchange(e: *const Engine) fw.Exchange {
+    pub fn exchange(e: *const Engine) fw.Exchange {
         return .{ .nccl = .{ .lib = e.lib, .comm = e.comm } };
     }
 
@@ -264,6 +273,7 @@ pub const Engine = struct {
         if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start));
         try e.gate(cancel);
         try f.run(as_i32[0..rows]);
+        if (e.prompt_logprobs) |p| if (part.len > 1) try prompt_lp.run(e, p, part[1..]);
         const last = start + rows == len;
         var first: ?u32 = null;
         if (last) {
@@ -349,7 +359,7 @@ pub const Engine = struct {
     }
 
     /// The ranks' maximum of `v` (an all-reduce on the engine's stream).
-    fn shared(e: *Engine, v: f64) !f64 {
+    pub fn shared(e: *Engine, v: f64) !f64 {
         var x = v;
         try e.shared_buf.uploadAsync(0, std.mem.asBytes(&x), e.stream.handle);
         try e.lib.check(e.lib.api.ncclAllReduce(e.shared_buf.ptr, e.shared_buf.ptr, 1, .f64, .max, e.comm, e.stream.handle), "ncclAllReduce");

zig/src/families/glm/cuda_follow.zig

@@ -13,7 +13,7 @@ pub const Header = extern struct {
     kind: u32,
     n: u32 = 0, // ids that follow
     depth: u32 = 0, // draft: drafts to chain; chunk: where it starts; shared: its parts
-    flags: u32 = 0, // begin: sampled; draft: from the prompt
+    flags: u32 = 0, // begin: sampled; draft: from the prompt; prefill, chunk: prompt_logprobs (the head on each row)
     a: u64 = 0, // keep: rows; prefill, chunk: the kept state it resumes (0: none)
     seed: u64 = 0, // begin: the sampling seed; prefill, chunk: the kept state's position
     temperature: f64 = 0,
@@ -114,6 +114,11 @@ pub const Leader = struct {
         try x.link.send(x.e.m.glue.put_words, hh, ids);
     }
 
+    /// A prompt step's flags: 1 when its rows go through the head (the other ranks join each pass's broadcasts).
+    fn rows(x: Leader) u32 {
+        return @intFromBool(x.e.prompt_logprobs != null);
+    }
+
     /// The stream's sequence for the steps that follow (each step names it, so the other ranks bind it too).
     pub fn bind(x: Leader, slot: u32) !void {
         try x.e.bind(slot);
@@ -156,7 +161,7 @@ pub const Leader = struct {
     /// `cancel` (rank 0's), asked before each chunk, decides for every rank; `from` every rank holds; `marks` rank 0's.
     pub fn prefill(x: Leader, ids: []const u32, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !u32 {
         try x.flush();
-        try x.send(.{ .kind = @backingInt(Kind.prefill), .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0 }, ids);
+        try x.send(.{ .kind = @backingInt(Kind.prefill), .flags = x.rows(), .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0 }, ids);
         return x.e.prefill(ids, from, cancel, marks);
     }
 
@@ -165,7 +170,7 @@ pub const Leader = struct {
         if (start >= ids.len) return error.BadChunk;
         try x.flush();
         const part = ids[start..@min(ids.len, start + engine.prompt_rows + 1)];
-        try x.send(.{ .kind = @backingInt(Kind.chunk), .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);
+        try x.send(.{ .kind = @backingInt(Kind.chunk), .flags = x.rows(), .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);
         return x.e.prefillChunk(part, start, ids.len, from, cancel, marks);
     }
 
@@ -210,6 +215,12 @@ pub const Leader = struct {
     }
 };
 
+/// A prompt step whose rows go through the head (Leader.rows): this rank's columns join each pass's broadcasts.
+fn promptRows(h: Header) bool {
+    const kind: Kind = @fromBackingInt(@intCast(h.kind));
+    return (kind == .prefill or kind == .chunk) and h.flags & 1 != 0;
+}
+
 /// A rank other than 0 runs rank 0's steps until it ends; a cancelled prefill is error.Cancelled and the loop goes on.
 pub fn follow(e: *Engine, link: *Link) !void {
     var drafts: [engine.max_drafts]u32 = undefined;
@@ -219,6 +230,7 @@ pub fn follow(e: *Engine, link: *Link) !void {
         const h = step.h;
         const kind: Kind = @fromBackingInt(@intCast(h.kind));
         if (kind != .end and kind != .drop and kind != .shared) try e.bind(h.slot);
+        e.prompt_logprobs = if (promptRows(h)) .{} else null; // rank 0 alone reads the rows
         switch (kind) {
             .begin => try e.begin(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 => {
@@ -277,3 +289,11 @@ test "a step's header is 64 bytes and keeps its fields" {
     try std.testing.expectEqual(h, back);
     try std.testing.expectEqual(Kind.draft, @as(Kind, @fromBackingInt(@intCast(back.kind))));
 }
+
+test "flag 1 asks for prompt rows on prefill and chunk steps only (begin's is sampling, draft's from the prompt)" {
+    try std.testing.expect(promptRows(.{ .kind = @backingInt(Kind.chunk), .flags = 1, .depth = 2048, .top_k = 4000 }));
+    try std.testing.expect(promptRows(.{ .kind = @backingInt(Kind.prefill), .flags = 1 }));
+    try std.testing.expect(!promptRows(.{ .kind = @backingInt(Kind.chunk), .depth = 2048, .top_k = 4000 }));
+    try std.testing.expect(!promptRows(.{ .kind = @backingInt(Kind.begin), .flags = 1, .seed = 7 }));
+    try std.testing.expect(!promptRows(.{ .kind = @backingInt(Kind.draft), .flags = 1, .depth = 3 }));
+}

zig/src/families/glm/cuda_lanes.zig

@@ -5,6 +5,7 @@ const lanes = @import("lanes");
 const engine = @import("cuda_engine.zig");
 const follow = @import("cuda_follow.zig");
 const fw = @import("cuda_forward.zig");
+const prompt_lp = @import("cuda_prompt_logprobs.zig");
 
 const be = lanes.backend;
 
@@ -49,6 +50,7 @@ pub const Lanes = struct {
         .keep = keepFn,
         .draft = draftFn,
         .release = releaseFn,
+        .prompt_rows = true,
     };
     const steps_vt: be.Backend.VTable = vt: {
         var v = whole_vt;
@@ -121,6 +123,9 @@ pub const Lanes = struct {
     fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void {
         const self = of(ptr);
         const from = try self.opening(s);
+        var sink: Sink = .{ .gpa = self.x.e.gpa, .s = s };
+        self.x.e.prompt_logprobs = sink.want();
+        defer self.x.e.prompt_logprobs = null;
         _ = self.take(try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s)));
     }
 
@@ -133,6 +138,9 @@ pub const Lanes = struct {
         if (going != null) try self.x.bind(l.slot);
         const at = going orelse if (from) |r| r.at else 0;
         errdefer l.fill = null;
+        var sink: Sink = .{ .gpa = self.x.e.gpa, .s = s };
+        self.x.e.prompt_logprobs = sink.want();
+        defer self.x.e.prompt_logprobs = null;
         if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, marksOf(s))) |first| {
             _ = self.take(first);
             l.fill = null;
@@ -154,9 +162,27 @@ pub const Lanes = struct {
         try self.x.bind(gop.value_ptr.slot);
         try self.x.begin(s.sampling);
         s.cached = 0;
-        return self.resumeFrom(s, ids.len);
+        s.prompt_rows.clearRetainingCapacity();
+        return if (s.prompt_logprobs != null) null else self.resumeFrom(s, ids.len); // prompt rows need every row
     }
 
+    /// A stream's prompt rows as rank 0 reads them back, in prompt order.
+    const Sink = struct {
+        gpa: std.mem.Allocator,
+        s: *lanes.Stream,
+
+        fn want(p: *Sink) ?prompt_lp.Want {
+            const k = p.s.prompt_logprobs orelse return null;
+            return .{ .k = k, .reader = .{ .ctx = p, .call = read } };
+        }
+
+        fn read(ctx: *anyopaque, words: []const u32, picks: []const u32) anyerror!void {
+            const p: *Sink = @ptrCast(@alignCast(ctx));
+            const w = prompt_lp.row_words;
+            for (picks, 0..) |pick, i| try p.s.prompt_rows.append(p.gpa, lanes.LogprobRow.fromWords(words[i * w ..][0..w], pick, p.s.prompt_logprobs.?));
+        }
+    };
+
     fn marksOf(s: *lanes.Stream) ?engine.Marks {
         return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt };
     }

zig/src/families/glm/cuda_prompt_logprobs.zig

@@ -0,0 +1,94 @@
+//! prompt_logprobs at TP: a chunk's rows through each rank's head columns, made whole (Sampler.gather) for rank 0.
+const std = @import("std");
+const cuda = @import("cuda");
+const lanes = @import("lanes");
+const b16 = @import("cuda_b16.zig");
+const fw = @import("cuda_forward.zig");
+const config = @import("config.zig");
+const shard = @import("cuda_shard.zig");
+const Engine = @import("cuda_engine.zig").Engine;
+
+/// The words the kernel writes for one row (lanes.logprob's).
+pub const row_words = lanes.logprob.words;
+
+/// Rows a head pass takes with buffers of its own (the head's columns read once for all); else a verify window's.
+pub const batch_rows = 128;
+
+/// A prompt pass with prompt_logprobs: every rank runs the head on its rows; rank 0 has the reader of their words.
+pub const Want = struct {
+    k: u8 = 0,
+    reader: ?Reader = null,
+};
+
+/// Rank 0's: `words` (rows of row_words) for `picks`, the prompt tokens the rows score, in prompt order.
+pub const Reader = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, words: []const u32, picks: []const u32) anyerror!void };
+
+/// This rank's columns of batch_rows rows and the whole rows, made at the first pass that asks.
+pub const Buffers = struct {
+    columns: cuda.DeviceBuffer,
+    full: cuda.DeviceBuffer,
+
+    fn make(e: *const Engine) !Buffers {
+        var columns = try cuda.DeviceBuffer.alloc(e.d, batch_rows * e.share.vocab.len() * 2);
+        errdefer columns.free();
+        return .{ .columns = columns, .full = try cuda.DeviceBuffer.alloc(e.d, batch_rows * e.c.vocab * 2) };
+    }
+
+    pub fn free(b: *Buffers) void {
+        b.full.free();
+        b.columns.free();
+    }
+};
+
+/// The host-mapped words of a pass's rows, then their picks (Engine.prompt_out).
+pub const out_bytes = batch_rows * (row_words + 1) * 4;
+
+/// The chunk's rows in e.bp.fnormed that score `next` (row r scores next[r]), a head pass at a time.
+pub fn run(e: *Engine, want: Want, next: []const u32) !void {
+    const own = try buffers(e);
+    const rows: usize = if (own != null) batch_rows else fw.window_rows;
+    const columns = if (own) |b| b.columns.ptr else e.bd.logits;
+    const full = if (own) |b| b.full.ptr else e.smp.full.ptr;
+    const D: usize = e.c.hidden;
+    const words = e.prompt_out.slice(u32);
+    const picks = words[batch_rows * row_words ..][0..batch_rows];
+    const out = try e.prompt_out.device();
+    var r: usize = 0;
+    while (r < next.len) : (r += rows) {
+        const n = @min(rows, next.len - r);
+        // as the head's last-row pass (Forward.finish); its column count takes no K slices (the test below)
+        try b16.matmul(.{ .glue = &e.m.glue }, e.stream, e.bp.fnormed + r * D * 2, D, e.w.head, columns, e.bp.kda.mm_part, n, e.share.vocab.len(), D, false);
+        try e.smp.gather(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = columns }}, n, full);
+        const reader = want.reader orelse continue;
+        @memcpy(picks[0..n], next[r..][0..n]);
+        try e.smp.logprobRows(e.stream, full, n, out + batch_rows * row_words * 4, want.k, out);
+        try e.stream.synchronize();
+        try reader.call(reader.ctx, words[0 .. n * row_words], picks[0..n]);
+    }
+}
+
+/// The passes' own buffers from the first that asks; if a rank has no room, every rank takes the window's (8 rows).
+fn buffers(e: *Engine) !?*Buffers {
+    if (!e.prompt_tried) {
+        e.prompt_tried = true;
+        e.prompt_bufs = Buffers.make(e) catch null;
+        if (try e.shared(@floatFromInt(@intFromBool(e.prompt_bufs == null))) > 0) {
+            if (e.prompt_bufs) |*b| b.free();
+            e.prompt_bufs = null;
+            std.log.warn("prompt_logprobs: a rank has no room for {d}-row head passes; {d} rows a pass", .{ batch_rows, fw.window_rows });
+        }
+    }
+    return if (e.prompt_bufs) |*b| b else null;
+}
+
+comptime {
+    std.debug.assert(batch_rows >= fw.window_rows); // the window's passes use the same host-mapped words
+}
+
+test "the head's columns at TP=2 and TP=3 take no K slices: a pass of any rows needs no partial sums" {
+    const c: config.Config = .{};
+    for ([_]usize{ 2, 3 }) |world| for (0..world) |rank| {
+        const s = try shard.Share.of(&c, rank, world);
+        try std.testing.expectEqual(@as(usize, 1), b16.splitK(s.vocab.len(), c.hidden));
+    };
+}