From 28f3e930f76bf99997d238fee863e96e38d76e50 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:04:35 +0900 Subject: [PATCH 1/2] GLM-5.3-Flash on CUDA: prompt_logprobs on /v1/completions The lane backend fills Stream.prompt_rows, so the server stops refusing prompt_logprobs for GLM. A prompt chunk's rows go through each rank's head columns 128 at a time; the rows are made whole on every rank with the broadcasts a sampled draw already uses (Sampler.gather), and rank 0 runs sample.cu's tf_logprob_rows on them. Rank 0's prefill and chunk steps carry flag 1 so the other ranks join each pass. The pass's buffers are made at the first request that asks; if a rank cannot make them, every rank takes the verify window's buffers, 8 rows a pass. A pass with prompt rows starts from position 0, as Nemotron's does. --- zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_engine.zig | 16 +++- zig/src/families/glm/cuda_follow.zig | 26 ++++- zig/src/families/glm/cuda_lanes.zig | 28 +++++- zig/src/families/glm/cuda_prompt_logprobs.zig | 94 +++++++++++++++++++ zig/src/families/glm/cuda_sampler.zig | 52 ++++++---- 6 files changed, 195 insertions(+), 23 deletions(-) create mode 100644 zig/src/families/glm/cuda_prompt_logprobs.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index af9254de61..93f5932280 100644 --- a/zig/src/families/glm/cuda.zig +++ b/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; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index ee38e00b3e..90310f2ebd 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/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"); diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 8072b7cfe1..9852d824b1 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/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 })); +} diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 0106c383b0..d88bef5d0b 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/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 }; } diff --git a/zig/src/families/glm/cuda_prompt_logprobs.zig b/zig/src/families/glm/cuda_prompt_logprobs.zig new file mode 100644 index 0000000000..41d483a431 --- /dev/null +++ b/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)); + }; +} diff --git a/zig/src/families/glm/cuda_sampler.zig b/zig/src/families/glm/cuda_sampler.zig index 5eea093bb3..9f76b3eb30 100644 --- a/zig/src/families/glm/cuda_sampler.zig +++ b/zig/src/families/glm/cuda_sampler.zig @@ -29,6 +29,7 @@ pub const Sampler = struct { d: *const cuda.Driver, module: cuda.Module, tf_draw: cuda.Function, + tf_logprob_rows: cuda.Function, vocab: usize, rows: usize, full: cuda.DeviceBuffer, // bf16 [rows][vocab]: every rank's columns in place @@ -41,6 +42,7 @@ pub const Sampler = struct { var m = try cuda.Module.load(d, cuda.kernels.sample); errdefer m.unload(); const f = try m.function("tf_draw"); + const lp = try m.function("tf_logprob_rows"); var full = try cuda.DeviceBuffer.alloc(d, rows * vocab * 2); errdefer full.free(); var r = try cuda.DeviceBuffer.alloc(d, @sizeOf(Rule)); @@ -50,7 +52,7 @@ pub const Sampler = struct { var tok = try cuda.DeviceBuffer.alloc(d, rows * 4); errdefer tok.free(); const prob = try cuda.DeviceBuffer.alloc(d, rows * 4); - return .{ .d = d, .module = m, .tf_draw = f, .vocab = vocab, .rows = rows, .full = full, .rule = r, .meta = meta, .tok = tok, .prob = prob }; + return .{ .d = d, .module = m, .tf_draw = f, .tf_logprob_rows = lp, .vocab = vocab, .rows = rows, .full = full, .rule = r, .meta = meta, .tok = tok, .prob = prob }; } pub fn deinit(s: *Sampler) void { @@ -66,21 +68,7 @@ pub const Sampler = struct { pub fn draw(s: *Sampler, stream: cuda.Stream, exchange: Exchange, spans: []const Span, mine: []const Local, rows: usize, first: u64, r: Rule, out: []u32) !void { std.debug.assert(rows <= s.rows and out.len >= rows and first >= 1); const V = s.vocab; - for (mine) |l| { - const span = spans[l.rank]; - for (0..rows) |row| try s.copy(stream, s.full.ptr + (row * V + span.lo) * 2, l.logits + row * span.len() * 2, span.len() * 2); - } - switch (exchange) { - .local => {}, - .nccl => |x| { // each rank's columns of each row from that rank (spans may differ in width: TP=3) - try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart"); - for (0..rows) |row| for (spans, 0..) |span, q| { - const at = s.full.ptr + (row * V + span.lo) * 2; - try x.lib.check(x.lib.api.ncclBroadcast(at, at, span.len(), .bf16, @intCast(q), x.comm, stream.handle), "ncclBroadcast"); - }; - try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd"); - }, - } + try s.gather(stream, exchange, spans, mine, rows, s.full.ptr); try s.rule.uploadAsync(0, std.mem.asBytes(&r), stream.handle); const meta = [4]i32{ @intCast(first - 1), 0, 0, 0 }; // row r keys meta[0] + r + 1 try s.meta.uploadAsync(0, std.mem.asBytes(&meta), stream.handle); @@ -97,6 +85,38 @@ pub const Sampler = struct { try s.tok.download(0, std.mem.sliceAsBytes(out[0..rows])); } + /// `rows` rows of `mine` made whole in `full` (bf16 [rows][vocab]) on every rank, each column from its rank. + pub fn gather(s: *const Sampler, stream: cuda.Stream, exchange: Exchange, spans: []const Span, mine: []const Local, rows: usize, full: u64) !void { + const V = s.vocab; + for (mine) |l| { + const span = spans[l.rank]; + for (0..rows) |row| try s.copy(stream, full + (row * V + span.lo) * 2, l.logits + row * span.len() * 2, span.len() * 2); + } + switch (exchange) { + .local => {}, + .nccl => |x| { // each rank's columns of each row from that rank (spans may differ in width: TP=3) + try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..rows) |row| for (spans, 0..) |span, q| { + const at = full + (row * V + span.lo) * 2; + try x.lib.check(x.lib.api.ncclBroadcast(at, at, span.len(), .bf16, @intCast(q), x.comm, stream.handle), "ncclBroadcast"); + }; + try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + }, + } + } + + /// sample.cu's logprob words of `rows` whole rows at `full`: each row's pick in `picks`, its `k` best, its rank. + pub fn logprobRows(s: *const Sampler, stream: cuda.Stream, full: u64, rows: usize, picks: u64, k: u8, out: u64) !void { + var args: cuda.Args = .{}; + args.add(full); + args.add(@as(u32, @intCast(s.vocab))); + args.add(picks); + args.add(@as(u32, k)); + args.add(@as(u32, 1)); // the rank too: vLLM's prompt_logprobs give it + args.add(out); + try cuda.launch.launch(s.tf_logprob_rows, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, stream, &args); + } + /// The whole rows a draw read (bf16 [rows][vocab]), for checks against the host rule. pub fn rowsRead(s: *const Sampler, rows: usize, host: []u16) !void { try s.full.download(0, std.mem.sliceAsBytes(host[0 .. rows * s.vocab])); From 5a6ad9edf8bf17e4a7dacc75749d4742d2924115 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:52:51 +0900 Subject: [PATCH 2/2] docs(api): prompt_logprobs on GLM-5.3-Flash's CUDA engine --- docs/api.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/api.md b/docs/api.md index 85194d2150..2d29143be0 100644 --- a/docs/api.md +++ b/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,