diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index fde33f7010..4e3834c262 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/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; diff --git a/zig/src/core/lane_host_many_test.zig b/zig/src/core/lane_host_many_test.zig index e58f760353..f2f165fb3a 100644 --- a/zig/src/core/lane_host_many_test.zig +++ b/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); diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 330109ee73..03c5e46d12 100644 --- a/zig/src/core/lanes/backend.zig +++ b/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. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 539d938e1f..583c14a0ad 100644 --- a/zig/src/core/lanes/engine.zig +++ b/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). diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index e6f92c5d17..bca2468b87 100644 --- a/zig/src/core/lanes/fake.zig +++ b/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 } }; diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 1af320dcc5..5346275b79 100644 --- a/zig/src/core/lanes/open.zig +++ b/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; diff --git a/zig/src/families/glm/cuda_b16.zig b/zig/src/families/glm/cuda_b16.zig index 31dcb73127..77e6fb6171 100644 --- a/zig/src/families/glm/cuda_b16.zig +++ b/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 = .{}; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 4971aca5bd..4a25cefca9 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/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"); } }; diff --git a/zig/src/families/glm/cuda_engine_types.zig b/zig/src/families/glm/cuda_engine_types.zig index 77ed0e1b11..c7716f0d18 100644 --- a/zig/src/families/glm/cuda_engine_types.zig +++ b/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 }; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index fb7caa1e3d..e3695b812a 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/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); diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index a646c70899..86a438635f 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/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); diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index a4a2467e03..6398464656 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -4,8 +4,13 @@ const std = @import("std"); const cuda = @import("cuda"); const mm_mod = @import("cuda_mm.zig"); -/// What the forward's kernels launch from: our glue's functions. -pub const Set = struct { glue: *const Kernels, mm: ?*const mm_mod.Kernels = null, packs: ?*const mm_mod.Packed = null }; // mm: linears tiled at load +/// What the forward's kernels launch from: our glue's functions, glm_mm's (tiled weights), a prompt pass's flag. +pub const Set = struct { + glue: *const Kernels, + mm: ?*const mm_mod.Kernels = null, // linears tiled at load + packs: ?*const mm_mod.Packed = null, + prompt: bool = false, // glm_mm's one-slice blocks at any rows: a prompt's sums alike alone or batched +}; /// Our glm_glue.cu's functions. pub const Kernels = struct { diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 6e75236690..1dd8a511cf 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -3,6 +3,7 @@ const std = @import("std"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); +const types = @import("cuda_engine_types.zig"); const follow = @import("cuda_follow.zig"); const fw = @import("cuda_forward.zig"); const dflash = @import("dflash.zig"); @@ -35,6 +36,7 @@ const Lane = struct { chances: [max_held]f64 = undefined, // DFlash2's: each held draft's chance from the selector held_n: usize = 0, fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + first: ?u64 = null, // its first token's handle once its prompt pass ends }; pub const Lanes = struct { @@ -56,6 +58,7 @@ pub const Lanes = struct { const whole_vt: be.Backend.VTable = .{ .prefill = prefillFn, + .prefill_many = prefillManyFn, .first = firstFn, .queue = queueFn, .read = readFn, @@ -68,6 +71,7 @@ pub const Lanes = struct { const steps_vt: be.Backend.VTable = vt: { var v = whole_vt; v.prefill_step = prefillStepFn; + v.many_rows = engine.prompt_rows; // short fresh prompts go together; longer ones fill by chunks break :vt v; }; @@ -156,7 +160,34 @@ pub const Lanes = struct { const self = of(ptr); const from = try self.opening(s); try self.images(s, from); - _ = self.take(try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s))); + const first = try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s)); + self.lanes.getPtr(s).?.first = self.take(first); + } + + /// Fresh text prompts together, in passes of at most prompt_rows rows; false: any image or kept state (each alone). + fn prefillManyFn(ptr: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const self = of(ptr); + if (ss.len > types.max_together) return false; + for (ss) |s| if (s.images.len > 0 or s.reuse.saved != null or s.prompt_len > engine.prompt_rows or + s.prompt_len + s.max_new + max_rows >= self.x.e.capacity) return false; // opening's limit, before any pass + var opened: usize = 0; + errdefer for (ss[0..opened]) |s| releaseFn(ptr, s); + var parts: [types.max_together]types.Prompt = undefined; + var firsts: [types.max_together]u32 = undefined; + var lens: [types.max_together]usize = undefined; + for (ss, lens[0..ss.len]) |s, *n| n.* = s.prompt_len; + var cut: [types.max_together][2]usize = undefined; + for (cut[0..passes(lens[0..ss.len], engine.prompt_rows, &cut)]) |p| { + for (ss[p[0]..p[1]], p[0]..) |s, i| { + if (try self.opening(s) != null) return error.BadPrompts; // no kept state (checked above) + opened += 1; + parts[i] = .{ .slot = self.lanes.getPtr(s).?.slot, .ids = s.prompt() }; + } + // stream i's hold row is i: every pass's rows stay until the core asks each stream's first drafts + try self.x.prefillMany(parts[p[0]..p[1]], firsts[p[0]..p[1]], p[0]); + for (ss[p[0]..p[1]], firsts[p[0]..p[1]]) |s, t| self.lanes.getPtr(s).?.first = self.take(t); + } + return true; } /// One chunk of the stream's prompt pass (fillOne runs other rounds between); true once the first token is drawn. @@ -170,7 +201,7 @@ pub const Lanes = struct { const at = going orelse if (from) |r| r.at else 0; errdefer l.fill = null; if (try self.x.chunk(s.prompt(), at, from, .{ .ptr = s, .check = cancelled }, marksOf(s))) |first| { - _ = self.take(first); + l.first = self.take(first); l.fill = null; return true; } @@ -186,6 +217,7 @@ pub const Lanes = struct { l.pending_rows = null; l.held_n = 0; l.fill = null; + l.first = null; try self.x.bind(l.slot); if (self.x.e.images) |v| v.clear(); // no earlier request's rows splice into this one try self.x.begin(s.sampling); @@ -227,7 +259,8 @@ pub const Lanes = struct { fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const self = of(ptr); if (position != s.prompt_len) return error.PositionMismatch; - return self.next - 1; + const l = self.lanes.getPtr(s) orelse return error.UnknownStream; + return l.first orelse error.NoFirstToken; } fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { @@ -389,6 +422,47 @@ pub const Lanes = struct { } }; +/// prefill_many's passes over prompts of `lens` tokens, in order, at most `rows` tokens a pass: [lo, hi) of each. +fn passes(lens: []const usize, rows: usize, out: *[types.max_together][2]usize) usize { + var n: usize = 0; + var lo: usize = 0; + var used: usize = 0; + for (lens, 0..) |len, i| { + if (i > lo and used + len > rows) { // this pass is full: the next starts here + out[n] = .{ lo, i }; + n += 1; + lo = i; + used = 0; + } + used += len; + } + if (lens.len > lo) { + out[n] = .{ lo, lens.len }; + n += 1; + } + return n; +} + +test "prefill_many's passes: in order, each within the rows, every prompt once, hold rows distinct across passes" { + var cut: [types.max_together][2]usize = undefined; + const lens = [_]usize{ 1500, 800, 300, 2048, 10, 10 }; + const n = passes(&lens, 2048, &cut); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 3 }, .{ 3, 4 }, .{ 4, 6 } }, cut[0..n]); + var seen: [types.max_together]bool = @splat(false); + for (cut[0..n]) |p| { + var used: usize = 0; + for (p[0]..p[1]) |i| { + used += lens[i]; + try std.testing.expect(!seen[i]); // a stream's hold row is its index in the batch: no two share one + seen[i] = true; + } + try std.testing.expect(used <= 2048); + } + for (seen[0..lens.len]) |s| try std.testing.expect(s); + try std.testing.expectEqual(@as(usize, 1), passes(&[_]usize{ 28, 28, 28, 28 }, 2048, &cut)); // the ruler's case + try std.testing.expectEqual(@as(usize, 0), passes(&[_]usize{}, 2048, &cut)); +} + comptime { std.debug.assert(max_rows <= fw.window_rows); } diff --git a/zig/src/families/glm/cuda_mm.zig b/zig/src/families/glm/cuda_mm.zig index d3edd83550..06cad43f82 100644 --- a/zig/src/families/glm/cuda_mm.zig +++ b/zig/src/families/glm/cuda_mm.zig @@ -96,13 +96,13 @@ pub fn partBytes(m: usize, w: Tiled) usize { return if (sk > 1) sk * @min(m, 64) * w.n * 4 else 0; } -/// x [m][k] (rows `ldx` apart) by `w` -> out [m][n], bf16 or fp32 (`wide`); `part`: partBytes(m, w). -pub fn matmul(k: *const Kernels, s: cuda.Stream, x: u64, ldx: usize, w: Tiled, out: u64, part: u64, m: usize, wide: bool) !void { +/// x [m][k] (`ldx` apart) by `w` -> out [m][n] (`wide`: fp32); `one`: prompt blocks at any rows, sums per row alike. +pub fn matmul(k: *const Kernels, s: cuda.Stream, x: u64, ldx: usize, w: Tiled, out: u64, part: u64, m: usize, wide: bool, one: bool) !void { if (w.k % 64 != 0 or w.n % 8 != 0 or x % 16 != 0 or ldx % 8 != 0) return error.BadShape; var done: usize = 0; const wi = try widthIndex(w.bits); while (done < m) { // windows to 64 rows a launch with K slices; prompts in one launch of 128-row blocks, one slice - const big = m - done >= 128; + const big = one or m - done >= 128; const rows = if (big) m - done else @min(m - done, 64); const i: usize = if (big) 3 else if (rows <= 16) 0 else if (rows <= 32) 1 else 2; const bm = bucket(i); diff --git a/zig/src/families/glm/cuda_prompts.zig b/zig/src/families/glm/cuda_prompts.zig new file mode 100644 index 0000000000..b8b57d2a80 --- /dev/null +++ b/zig/src/families/glm/cuda_prompts.zig @@ -0,0 +1,72 @@ +//! GLM-5.3-Flash's prefill_many: fresh streams' whole prompts in one prompt pass, each left as its own pass leaves it. + +const std = @import("std"); +const fw = @import("cuda_forward.zig"); +const engine = @import("cuda_engine_types.zig"); +const Engine = @import("cuda_engine.zig").Engine; +const dflash = @import("dflash.zig"); + +comptime { // a pass's rows go to DFlash2 in one ingest, every prompt within its window + std.debug.assert(fw.prefill_rows <= dflash.MAX_CTX and fw.prefill_rows <= dflash.WIN); +} + +/// Whether `parts` go in one pass: fresh prompts, prompt_rows in all, under the capacity, hold rows from hold0 fit. +pub fn fits(parts: []const engine.Prompt, hold0: usize, capacity: usize) bool { + if (parts.len == 0 or hold0 + parts.len > engine.max_together) return false; + var rows: usize = 0; + for (parts, 0..) |p, i| { + if (p.ids.len == 0 or p.ids.len + fw.window_rows >= capacity) return false; + for (parts[0..i]) |q| if (q.slot == p.slot) return false; + rows += p.ids.len; + } + return rows <= fw.prefill_rows; +} + +/// The prompts' pass (each begun at 0): prompt k's first token in firsts[k], draft source in hold row hold0 + k. +pub fn prefillMany(e: *Engine, parts: []const engine.Prompt, firsts: []u32, hold0: usize) !void { + if (!fits(parts, hold0, e.capacity) or firsts.len < parts.len) return error.BadPrompts; + var segs: [engine.max_together]fw.Seg = undefined; + var ids: [fw.prefill_rows]u32 = undefined; + var rows: usize = 0; + for (parts, 0..) |p, k| { + try e.bind(p.slot); + if (e.cur.st.pos != 0 or e.cur.st.mtp_len != 0) return error.BadPrompts; + segs[k] = .{ .st = &e.cur.st, .row0 = rows, .rows = p.ids.len }; + @memcpy(ids[rows..][0..p.ids.len], p.ids); + rows += p.ids.len; + } + try e.gate(null); + var one: [1]fw.Rank = .{.{ .w = &e.w, .st = segs[0].st, .b = &e.bp, .segs = segs[0..parts.len] }}; + var f: fw.Forward = .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = &one, .exchange = e.exchange(), .pieces = &e.pieces, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .taps = if (e.drafter) |dr| dr.taps(true) else null }; + try f.run(@ptrCast(ids[0..rows])); + if (e.drafter) |dr| { // every prompt is within DFlash2's window (prompt_rows): all its rows into its slot's ring + var at: [engine.max_together][3]usize = undefined; + for (parts, segs[0..parts.len], at[0..parts.len]) |p, g, *a| a.* = .{ p.slot, g.row0, g.rows }; + try dr.prompts(e, at[0..parts.len]); + } + const D = e.c.hidden; + for (parts, segs[0..parts.len], 0..) |p, g, k| { + try e.bind(p.slot); + try e.draw(e.bp.logits + k * e.share.vocab.len() * 2, 1, p.ids.len, firsts[k..][0..1]); + if (!e.mtp) continue; + try hold(e, hold0 + k, e.bp.fnormed + (g.row0 + g.rows - 1) * D * 2); + } + f.pieces = null; + f.taps = null; + for (parts, segs[0..parts.len]) |p, g| { + one[0] = .{ .w = &e.w, .st = g.st, .b = &e.bp }; + // decode._absorb_rows; the head's pass writes final-norm row 0, which the hold rows above already copied + if (e.mtp and g.rows > 1) { + try f.mtp(@ptrCast(p.ids[1..]), &.{e.bp.fnormed + g.row0 * D * 2}); + try g.st.setMtpLen(e.stream, g.st.mtp_len + g.rows - 1); + } + try f.commit(g.rows); + } +} + +/// The bound stream's first-draft source: its prompt's last hidden row (bf16 at `src`) into hold row `row`. +pub fn hold(e: *Engine, row: usize, src: u64) !void { + const at = e.hold.ptr + row * e.c.hidden * 2; + try e.copy(at, src, e.c.hidden * 2); + e.cur.held_at = at; +} diff --git a/zig/src/families/glm/cuda_shared.zig b/zig/src/families/glm/cuda_shared.zig index 1ac7cea033..1899ed31b2 100644 --- a/zig/src/families/glm/cuda_shared.zig +++ b/zig/src/families/glm/cuda_shared.zig @@ -1,4 +1,4 @@ -//! GLM-5.3-Flash's shared verify windows: several streams' rows in one forward, chains per segment. +//! GLM-5.3-Flash's shared windows: several streams' verify windows or whole prompts in one forward, chains per segment. const std = @import("std"); const cuda = @import("cuda"); @@ -8,25 +8,36 @@ const kda = @import("cuda_kda.zig"); const latent = @import("cuda_latent.zig"); const kda_block = @import("cuda_kda_block.zig"); const dsa_block = @import("cuda_dsa_block.zig"); +const glue = @import("cuda_glue.zig"); +const b16 = @import("cuda_b16.zig"); -/// Checks a shared window's segments against its `rows` rows, then each sequence holds its rows until its commit. +/// Checks a shared window's segments (a prompt buffer's: whole prompts) against its `rows` rows; each holds its own. pub fn begin(r: fw.Rank, rows: usize, stream: cuda.Stream) !void { - if (r.b.prefill) return error.BadWindow; + const most: usize = if (r.b.prefill) fw.prefill_rows else fw.window_rows; + if (r.b.prefill and r.segs.len > @import("cuda_state.zig").prompt_heads) return error.BadWindow; var at: usize = 0; for (r.segs) |g| { - if (g.row0 != at or g.rows == 0 or g.rows > fw.window_rows) return error.BadWindow; + if (g.row0 != at or g.rows == 0 or g.rows > most) return error.BadWindow; try g.st.ready(stream, g.rows); at += g.rows; } if (at != rows) return error.BadWindow; for (r.segs) |g| { g.st.pending = g.rows; - g.st.prompt = false; + g.st.prompt = r.b.prefill; } } -/// Forward.mixer on a shared window: all rows' projections, each segment's chain or attention alone, all o_proj. -pub fn mixer(f: *const fw.Forward, r: fw.Rank, l: weights.Layer, li: usize, di: usize, rows: usize, out: u64) !void { +/// Forward.finish's head on a prompt buffer: each segment's last final-norm row, logits row k for segment k. +pub fn heads(f: *const fw.Forward, r: fw.Rank) !void { + const D = r.w.config.hidden; + const b = r.b; + for (r.segs, 0..) |g, k| try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(b.normed + k * D * 2, b.fnormed + (g.row0 + g.rows - 1) * D * 2, D * 2, f.stream.handle), "cuMemcpyDtoDAsync"); + try b16.matmul(f.set(), f.stream, b.normed, D, r.w.head, b.logits, b.kda.mm_part, r.segs.len, r.w.share.vocab.len(), D, false); +} + +/// Forward.mixer on a shared window: all rows' projections, each segment's chain or attention, o_proj (null: pieced). +pub fn mixer(f: *const fw.Forward, r: fw.Rank, l: weights.Layer, li: usize, di: usize, rows: usize, out: ?u64) !void { const c = &r.w.config; const b = r.b; switch (l.mixer) { @@ -40,20 +51,28 @@ pub fn mixer(f: *const fw.Forward, r: fw.Rank, l: weights.Layer, li: usize, di: const st = g.st; const cur = st.cur[li]; var v = s; - v.p = st.proj + li * st.proj_layer * 2; // the window's rows stay with the layer until its commit - try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(v.p, b.kproj + g.row0 * w * 2, g.rows * w * 2, f.stream.handle), "cuMemcpyDtoDAsync"); v.ka += g.row0 * n * 2; v.kg += g.row0 * n * 2; v.saves.out += g.row0 * n * 2; v.saves.state_out = st.recAt(1 - cur, li); + const conv = st.conv + li * 3 * st.conv_c * 2; + if (b.prefill) { // a whole prompt: its rows in place, the layer committed now (as Forward.mixer) + v.p = b.kproj + g.row0 * w * 2; + try kda_block.chainRows(&f.m.kda, f.stream, k, v, g.rows, conv, st.recAt(cur, li), false); + try glue.convShift(f.set(), f.stream, conv, v.p, g.rows, g.rows * w, w, 1, 3 * n); + st.cur[li] = 1 - cur; + continue; + } + v.p = st.proj + li * st.proj_layer * 2; // the window's rows stay with the layer until its commit + try f.d.check(f.d.api.cuMemcpyDtoDAsync_v2(v.p, b.kproj + g.row0 * w * 2, g.rows * w * 2, f.stream.handle), "cuMemcpyDtoDAsync"); const kv = li * fw.window_rows * st.heads * kda.dk; v.saves.k = st.save_k + kv * 4; v.saves.v = st.save_v + kv * 2; v.saves.g = st.save_g + kv * 4; v.saves.b = st.save_b + li * fw.window_rows * st.heads * 4; - try kda_block.chainRows(&f.m.kda, f.stream, k, v, g.rows, st.conv + li * 3 * st.conv_c * 2, st.recAt(cur, li), f.kda_wide); + try kda_block.chainRows(&f.m.kda, f.stream, k, v, g.rows, conv, st.recAt(cur, li), f.kda_wide); } - try kda_block.fill(f.set(), f.stream, k, s, 0, rows, out); + if (out) |o| try kda_block.fill(f.set(), f.stream, k, s, 0, rows, o); }, .mla => |a| { var sparse_any = false; @@ -64,7 +83,7 @@ pub fn mixer(f: *const fw.Forward, r: fw.Rank, l: weights.Layer, li: usize, di: try dsa_block.attend(.{ .set = f.set(), .ops = f.m.ops() }, f.stream, a, g.st.caches[di], b.dsa, g.row0, g.rows, g.st.pos_dev, g.st.pos, nch, b.part_rows); } try latent.expandV(f.set(), f.stream, b.dsa.ol, a.wv, b.dsa.vn, rows, a.heads, a.dv, a.lw); - try dsa_block.fill(f.set(), f.stream, a, b.dsa, 0, rows, out); + if (out) |o| try dsa_block.fill(f.set(), f.stream, a, b.dsa, 0, rows, o); }, } } diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 54bc08c856..e51b4a22ad 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -19,6 +19,8 @@ const kda_block = @import("cuda_kda_block.zig"); pub const prefill_rows = 2048; /// A decode window's rows (Engine's max_rows). pub const window_rows = 8; +/// Prompts one prompt pass takes together (cuda_prompts): a prompt buffer's logits rows, one a prompt. +pub const prompt_heads = 16; /// The context the buffers' dense attention partials reach (Buffers' min(capacity, 2560)). pub const dense_reach = 2560; /// Index-key and gate rows a DSA cache holds (a ring): pool_keys reads only from the current pool's first row on. @@ -163,7 +165,7 @@ pub const Buffers = struct { rows * sh * 2, rows * (sh / 64) * 4, rows * D * 4, // sxs, sy rows * D * 2, rows * D * 2, - rows * (D / 64) * 4, (if (prefill) 1 else rows) * s.vocab.len() * 2, + rows * (D / 64) * 4, (if (prefill) prompt_heads else rows) * s.vocab.len() * 2, rows * D * 2, rows * D * 2, rows * 2 * D * 2, rows * D * 2, // hin, me, mcat, mx rows * D * 4, // part: this rank's fp32 partial rows a split exchange sends from x3[0], x3[1], x3[2], x3[3], // EXL3 experts: 8-slot picks, rotated inputs, a decode window's partials diff --git a/zig/src/families/glm/dflash.zig b/zig/src/families/glm/dflash.zig index a91799fa26..d488d8a0cc 100644 --- a/zig/src/families/glm/dflash.zig +++ b/zig/src/families/glm/dflash.zig @@ -215,6 +215,17 @@ pub const Drafter = struct { try dr.ingest(e, &.{run}); } + /// A prompt pass's whole prompts, each (slot, first row, rows): every row into its slot's ring, one ingest. + pub fn prompts(dr: *Drafter, e: anytype, parts: []const [3]usize) !void { + var runs: [16]Run = undefined; + if (parts.len > runs.len) return error.DrafterContextTooWide; + for (parts, runs[0..parts.len]) |p, *r| { + const s = try dr.slot(@intCast(p[0])); + r.* = .{ .ring = s.ring.?.ptr, .taps = dr.taps_prefill + p[1] * TAP_ROW, .n = p[2], .pos0 = 0 }; + } + try dr.ingest(e, runs[0..parts.len]); + } + /// The runs' rows through fc, the hidden norm and each layer's k, v into their rings. pub fn ingest(dr: *Drafter, e: anytype, runs: []const Run) !void { var N: usize = 0; diff --git a/zig/src/families/glm/dflash_block.zig b/zig/src/families/glm/dflash_block.zig index 90ab7fbfb4..c9ceb9442d 100644 --- a/zig/src/families/glm/dflash_block.zig +++ b/zig/src/families/glm/dflash_block.zig @@ -20,7 +20,7 @@ fn set(e: anytype) glue.Set { /// out[m, n] = x[m, k] . w^T, bf16 out (`wide`: fp32, glm_mm only): 4 bits 16 rows a launch, or glm_mm. fn lin(dr: *const Drafter, e: anytype, x: u64, w: df.Lin, out: u64, m: usize, wide: bool) !void { switch (w) { - .bf16 => |t| try mm.matmul(&e.m.mm, e.stream, x, t.k, t, out, dr.part, m, wide), + .bf16 => |t| try mm.matmul(&e.m.mm, e.stream, x, t.k, t, out, dr.part, m, wide, false), .q4 => |q| { if (wide) return error.WideFromFourBits; var a: cuda.Args = .{}; diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index c825d65f4b..8992a81b82 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -24,6 +24,7 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); } }.f else null, + .many_rows = v.many_rows, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); @@ -134,12 +135,12 @@ test "every entry the family's backend gives is forwarded, and none it leaves ou }; var fake: lanes.fake.Fake = .{ .gpa = std.testing.allocator, .prefill_chunks = 3 }; defer fake.deinit(); - for ([_]be.Backend{ fake.backend(), fake.stepped() }) |b| { + for ([_]be.Backend{ fake.backend(), fake.stepped(), fake.batchedSteps() }) |b| { const w = vtable(H, b.vtable.*); const info = @typeInfo(be.Backend.VTable).@"struct"; inline for (info.field_names, info.field_types) |name, T| { if (@typeInfo(T) == .optional) try std.testing.expectEqual(@field(b.vtable, name) != null, @field(w, name) != null); - if (T == bool) try std.testing.expectEqual(@field(b.vtable, name), @field(w, name)); + if (T == bool or @typeInfo(T) == .int) try std.testing.expectEqual(@field(b.vtable, name), @field(w, name)); // e.g. many_rows } } try std.testing.expect(vtable(H, fake.stepped().vtable.*).prefill_step != null);