diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 243b3386a4..8e1cf73b03 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -345,8 +345,8 @@ pub const Decision = struct { arena: ?std.mem.Allocator = null, // the request's, for what the engine hands back }; -/// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +/// An opened engine; ``close`` stops its thread and frees its backend; ``halt`` (null: none) a drain's deadline. +pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque, halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null }; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; @@ -559,8 +559,9 @@ pub const LoneHooks = struct { resumed: bool = false, // the stream comes from the lane core between rounds: no prompt pass, its drafts the core's }; -/// The lane core served to the HTTP threads (lane_host.zig). +/// The lane core served to the HTTP threads (lane_host.zig), and a family's opt-in round planner (lane_rounds.zig). pub const LaneHost = @import("lane_host.zig").LaneHost; +pub const Rounds = @import("lane_host.zig").Rounds; /// Exact prompt reuse between requests, for any family (prompt_cache.zig). pub const prompt_cache = @import("prompt_cache.zig"); @@ -590,8 +591,7 @@ pub const keepalive = @import("keepalive.zig"); pub const originals = @import("originals.zig"); test { - _ = @import("lane_host.zig"); - _ = @import("lane_host_test.zig"); + _ = @import("lane_host.zig"); // its test block imports lane_host_test.zig and the other lane host tests _ = @import("prompt_cache.zig"); _ = @import("prompt_radix.zig"); _ = @import("prompt_imprint.zig"); diff --git a/zig/src/core/lane_halt.zig b/zig/src/core/lane_halt.zig new file mode 100644 index 0000000000..71f875f6c5 --- /dev/null +++ b/zig/src/core/lane_halt.zig @@ -0,0 +1,33 @@ +//! A stop's drain deadline on the lane host: the rounds end at a boundary and what is left fails with a reason. +const std = @import("std"); +const LaneHost = @import("lane_host.zig").LaneHost; + +/// `LaneHost.halt`: the engine thread stops at the next round boundary; what is left fails with `reason`. +pub fn halt(h: *LaneHost, reason: []const u8) void { + h.mutex.lockUncancelable(h.io); + h.closing = true; + h.halt_reason = reason; + h.wake.broadcast(h.io); + h.mutex.unlock(h.io); + if (h.thread) |t| t.join(); + h.thread = null; +} + +/// `halt`: between rounds, every request queued, filling or running fails with `reason` (its stream discarded). +pub fn haltAll(h: *LaneHost, reason: []const u8) void { + h.lock(); + var queued = h.queued; + h.queued = .empty; + var jobs = h.admitted; + h.admitted = .empty; + h.filling.clearRetainingCapacity(); // a job still filling is admitted: it fails below, its lane freed + h.unlock(); + if (queued.items.len + jobs.items.len > 0) std.log.warn("lane host: stopping at a round boundary: {d} running and {d} queued request(s) end ({s})", .{ jobs.items.len, queued.items.len, reason }); + for (jobs.items) |job| { + if (job.started and !job.stream.finished) h.core.discard(&job.stream); + h.finish(job, .failed, reason); + } + for (queued.items) |job| h.finish(job, .failed, reason); + queued.deinit(h.gpa); + jobs.deinit(h.gpa); +} diff --git a/zig/src/core/lane_halt_test.zig b/zig/src/core/lane_halt_test.zig new file mode 100644 index 0000000000..87ba6710ab --- /dev/null +++ b/zig/src/core/lane_halt_test.zig @@ -0,0 +1,210 @@ +//! The lane host's halt: a stop's drain deadline ends the rounds at a boundary, what is left fails with the reason. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; + +test "halt (a stop's drain deadline): the rounds end at a round boundary, every request queued or running fails with the reason" { + const gpa = std.testing.allocator; + 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 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + rounds: u32 = 0, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) return .wait; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + /// the planner's rounds when the finish arrived (on the engine thread, between rounds) + rounds_at_end: u32 = 0, + plan: *Plan, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + b.rounds_at_end = b.plan.rounds; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + var long_prompt: [600]u32 = undefined; + for (&long_prompt, 0..) |*t, i| t.* = @intCast(1 + i % 13); + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + var boxes: [3]Box = .{ .{ .plan = &pl }, .{ .plan = &pl }, .{ .plan = &pl } }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: prefilling a long prompt in pieces, 3: queued behind 2 + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &long_prompt, .max_tokens = 1 << 30 }; + const r3: Request = .{ .prompt = &p1, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (boxes[0].count() < 20) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + const rounds = pl.rounds; + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + // no round was planned after the finishes: the loop stopped at the boundary they were sent at + try std.testing.expectEqual(rounds, b.rounds_at_end); + } + // the decoding stream's tokens are whole rounds' (the fake's own sequence), the prefilling one had none yet + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, &p1); + for (boxes[0].tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); + // nothing runs after it; the stop that follows (a close) finds the thread gone + std.Io.sleep(std.testing.io, .fromMilliseconds(20), .awake) catch {}; + try std.testing.expectEqual(rounds, pl.rounds); + host.stop(); +} + +test "halt fails a request still filling its prompt with the reason and frees its lane; a decoding one fails too" { + const gpa = std.testing.allocator; + 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, .prefill_chunks = 2 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.stepped(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + const Chunks = struct { + n: std.atomic.Value(usize) = .init(0), // chunks the second request's pass ran + fn call(ctx: *anyopaque, _: *lanes.Stream, _: usize) void { + const c: *@This() = @ptrCast(@alignCast(ctx)); + _ = c.n.fetchAdd(1, .monotonic); + } + }; + var chunks: Chunks = .{}; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var boxes: [3]Box = .{ .{}, .{}, .{} }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: filling a prompt that never ends, 3: queued behind both lanes + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const r3: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + // nothing fills now: the next pass takes chunks until the halt + target.prefill_chunks = std.math.maxInt(usize); + target.prefill_hook = Chunks.call; + target.prefill_hook_ctx = &chunks; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (chunks.n.load(.monotonic) < 3) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + } + try std.testing.expect(boxes[0].tokens.items.len >= 4); + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + // every lane freed: the decoding stream's and the filling one's + try std.testing.expectEqual(@as(usize, 0), target.lanes.count()); + try std.testing.expectEqual(@as(usize, 0), host.filling.items.len); + host.stop(); +} diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4a8f7d3493..d51d995313 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -21,12 +21,20 @@ const Info = api.Info; const Reason = api.Reason; const Stats = api.Stats; +pub const Rounds = rounds.Rounds; +pub const Proposers = rounds.Proposers; +const rounds = @import("lane_rounds.zig"); +const lane_stream = @import("lane_stream.zig"); +const lane_halt = @import("lane_halt.zig"); + pub const LaneHost = struct { gpa: Allocator, io: std.Io, core: *lanes.Engine, info_: Info, min_match: i64 = 4, + /// A family's own copy proposer for each request (null: lanes.SuffixLookup at `min_match`) + proposers: ?Proposers = null, mutex: std.Io.Mutex = .init, wake: std.Io.Condition = .init, queued: std.ArrayList(*Job) = .empty, @@ -35,6 +43,8 @@ pub const LaneHost = struct { cancels: std.ArrayList(Id) = .empty, closing: bool = false, stale: bool = false, // the learner's last step: a put-back has not reached the device, so replies are refused + /// `halt`'s reason: what is left finishes failed with it (null: a stop cancels what is left) + halt_reason: ?[]const u8 = null, thread: ?std.Thread = null, decoded: std.ArrayList(status.Mark) = .empty, // tokens a round landed, for the 2 s decode rate prefill_rate: f64 = 0, @@ -57,6 +67,14 @@ pub const LaneHost = struct { memory_lanes: ?usize = null, // lanes the backend had memory for at its last LaneMemory, until a job ends warm_quiet: u64 = 0, // ns a warm pass waits for other work before its first chunk (0: none), set by its family disk: rows_host.Disk = .{}, // the Mac app's questions about kept states, and the disk writer's quiet + /// Runs first on the engine thread: a backend whose API state is per thread (a CUDA context) binds it there. + on_thread: ?ThreadInit = null, + /// opt-in: a family round planner (prompt pieces, decode rows in one pass); null: whole or `prefill_step` prompts + rounds: ?Rounds = null, + pieces: std.ArrayList(Rounds.Piece) = .empty, + finals: std.ArrayList(usize) = .empty, + + pub const ThreadInit = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void }; const Lesson = struct { request: *const api.LearnRequest, sink: api.LearnSink }; @@ -67,11 +85,14 @@ pub const LaneHost = struct { sink: Sink, stream: lanes.Stream = undefined, proposer: lanes.SuffixLookup = undefined, + /// the family's proposer (LaneHost.proposers), in place of `proposer` + custom: ?lanes.proposer.Proposer = null, recall: recalls.Slot = .{}, // its reads of Recall's record (lane_host_recall.zig) delivered: usize = 0, started: bool = false, prefill_sent: bool = false, // a lone driver's prefilled event went out fill_began: bool = false, // its first prompt chunk ran + submitted: i96 = 0, began: i96 = 0, /// The engine thread says the prompt pass is going: status readers may count this job's fill atomics. fill_on: std.atomic.Value(bool) = .init(false), @@ -83,7 +104,7 @@ pub const LaneHost = struct { waits_since: ?u64 = null, // the admissions count when it first waited at the head for pool room /// The backend's prompt pass stands at a mark: the cache keeps the stream's state there. - fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { + pub fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { const job: *Job = @ptrCast(@alignCast(ptr)); job.reported(); // before a keep can evict the entry the pass restored if (job.host.cache) |store| _ = store.keepBorn(job.request.cacheKey(), job.request.bornLen(), at, s, job.request.chunks, job.host.spans(job.request)); @@ -108,10 +129,11 @@ pub const LaneHost = struct { } pub fn start(h: *LaneHost) !void { + if (h.rounds != null and !h.core.backend.pieces()) return error.PiecesUnsupported; // a drafter wrapper runs whole prompts h.thread = try std.Thread.spawn(.{ .stack_size = 16 << 20 }, run, .{h}); } - /// Stops admitting, cancels what is left and joins the engine thread. + /// Stops admitting, cancels what is left and joins the engine thread (after `halt`: the thread is gone already). pub fn stop(h: *LaneHost) void { rows_host.drainThread(h); rows_host.close(h); @@ -124,8 +146,13 @@ pub const LaneHost = struct { h.cancels.deinit(h.gpa); h.decoded.deinit(h.gpa); h.live_tokens.deinit(h.gpa); + h.pieces.deinit(h.gpa); + h.finals.deinit(h.gpa); } + /// A drain's deadline: the round in flight completes on every rank, so the ranks stop at the same boundary. + pub const halt = lane_halt.halt; + pub fn engine(h: *LaneHost) Engine { return .{ .ctx = h, .vtable = &.{ .info = infoFn, @@ -157,7 +184,7 @@ pub const LaneHost = struct { fn submitFn(ctx: *anyopaque, id: Id, request: *const Request, sink: Sink) api.SubmitError!void { const h = self(ctx); const job = h.gpa.create(Job) catch return error.Busy; - job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; + job.* = .{ .host = h, .id = id, .request = request, .sink = sink, .submitted = std.Io.Clock.awake.now(h.io).toNanoseconds() }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); if (h.closing or h.stale) { @@ -192,7 +219,7 @@ pub const LaneHost = struct { } /// A job's cancel hook for its prompt pass: its id is in `cancels`, or the host is closing (read under the lock). - fn cancelled(ctx: *anyopaque) bool { + pub fn cancelled(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -200,7 +227,7 @@ pub const LaneHost = struct { } /// A background pass's yield hook: a foreground request waits or a stream decodes (read under the lock). - fn yieldsTo(ctx: *anyopaque) bool { + pub fn yieldsTo(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -251,6 +278,7 @@ pub const LaneHost = struct { pub fn finish(h: *LaneHost, job: *Job, reason: Reason, message: []const u8) void { h.memory_lanes = null; // a lane may have freed: the next admission asks the backend for memory again + if (h.rounds) |rd| rd.vtable.left(rd.ctx, @intFromPtr(job)); job.reported(); h.gpa.free(job.marks); const s = &job.stream; @@ -260,8 +288,7 @@ pub const LaneHost = struct { if (job.started) { if (s.grammar) |g| g.deinit(h.gpa); s.deinit(h.gpa); - job.proposer.deinit(); - recalls.release(job); + lane_stream.freeProposer(h, job); } h.gpa.destroy(job); } @@ -297,57 +324,8 @@ pub const LaneHost = struct { h.unlock(); const r = job.request; if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran - var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass - if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), r.bornLen(), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; - job.kept0 = store.counts.kept; - job.marks = l.marks; - reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; - } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); - const recall = recalls.reader(h, job); - job.stream = lanes.Stream.init(h.gpa, .{ - .id = "request", - .prompt = r.prompt, - .max_new = r.max_tokens, - .eos = r.eos, - .sampling = r.sampling, - .drafts = r.drafts, - .proposer = job.proposer.proposer(), - .recall = recall, - .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, - .cancel_check = .{ .ptr = job, .check = cancelled }, - .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = yieldsTo } else null, - .think_budget = r.think_budget, - .think_close = r.think_close, - .think_end = if (r.think_end) |t| t else -1, - .logprobs = r.logprobs, - .prompt_logprobs = r.prompt_logprobs, - .loop_guard = r.loop_guard, - .chunks = r.chunks, - .history_len = r.history_len, - .shared_prefixes = r.shared_prefixes, - .reuse = reuse, - .grammar = if (r.structure) |st| h.constraint(st) catch |e| { - job.proposer.deinit(); - recalls.release(job); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); - } else null, - .decode_spans = h.spans(r), - .decide = r.decide, - .untaught = r.untaught, - .images = r.images, - .key = r.cache_key, - }) catch { - job.proposer.deinit(); - recalls.release(job); - return h.drop(job, "out of memory"); - }; - job.started = true; - job.fill_on.store(true, .release); // status readers count the fill from here until the pass reports - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; + if (!lane_stream.open(h, job)) return true; + const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { @@ -389,7 +367,7 @@ pub const LaneHost = struct { } /// A reply's grammar from the server's compiler (Info.structures engines get requests that carry one). - fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { + pub fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { const c = st.compiler orelse return error.NoGrammarCompiler; return c.constraint(h.gpa, h.io, @fromBackingInt(@intCast(@backingInt(st.kind))), st.text, st.after); } @@ -411,7 +389,7 @@ pub const LaneHost = struct { } /// A request's decoded prompt rows, for a backend that keeps their bits (`decoded_rows`); otherwise none. - fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { + pub fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { return if (h.decoded_rows) r.decode_spans else &.{}; } @@ -464,8 +442,7 @@ pub const LaneHost = struct { job.reported(); if (job.stream.grammar) |g| g.deinit(h.gpa); // the next admission compiles its own job.stream.deinit(h.gpa); - job.proposer.deinit(); - recalls.release(job); + lane_stream.freeProposer(h, job); job.started = false; job.fill_began = false; job.entry = null; // looked up again at its next admission @@ -477,10 +454,13 @@ pub const LaneHost = struct { const alone_words = "the engine has no memory for this request's caches with nothing else running: lower --context or --parallel"; - /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + /// A job cancelled in its prompt pass, its lane already released (by a `halt`: failed with its reason). + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); - h.finish(job, .cancelled, ""); + h.lock(); + const why = h.halt_reason; + h.unlock(); + h.finish(job, if (why == null) .cancelled else .failed, why orelse ""); return true; } @@ -495,6 +475,9 @@ pub const LaneHost = struct { } fn run(h: *LaneHost) void { + if (h.on_thread) |t| t.run(t.ctx) catch |e| std.log.err("lane host: the engine thread's init failed ({t})", .{e}); + // a family planner: prompt pieces and decode rows in one pass; else whole or `prefill_step` chunked prompts + if (h.rounds) |rd| return rounds.run(h, rd); while (true) { ends.takeCancels(h); rows_host.answer(h); @@ -548,4 +531,7 @@ pub const LaneHost = struct { test { _ = @import("lane_host_test.zig"); + _ = @import("lane_rounds_test.zig"); + _ = @import("lane_halt_test.zig"); + _ = @import("lane_review_test.zig"); } diff --git a/zig/src/core/lane_host_end.zig b/zig/src/core/lane_host_end.zig index 7163197979..0250579532 100644 --- a/zig/src/core/lane_host_end.zig +++ b/zig/src/core/lane_host_end.zig @@ -2,6 +2,7 @@ const std = @import("std"); const api = @import("engine_api.zig"); const fill = @import("lane_host_fill.zig"); +const lane_halt = @import("lane_halt.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; const Id = api.Id; @@ -40,6 +41,7 @@ pub fn takeCancels(h: *LaneHost) void { /// A failed round ends every stream it held, with the backend's error. pub fn failAll(h: *LaneHost, message: []const u8) void { + std.log.warn("lane host: a round failed ({s}): every admitted stream fails", .{message}); h.lock(); const jobs = h.gpa.dupe(*Job, h.admitted.items) catch &.{}; h.admitted.clearRetainingCapacity(); @@ -53,6 +55,7 @@ pub fn failAll(h: *LaneHost, message: []const u8) void { } pub fn closeAll(h: *LaneHost) void { + if (h.halt_reason) |reason| return lane_halt.haltAll(h, reason); h.lock(); for (h.queued.items) |job| h.cancels.append(h.gpa, job.id) catch {}; for (h.admitted.items) |job| h.cancels.append(h.gpa, job.id) catch {}; diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 85a34c0e15..de718c3c81 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -257,7 +257,7 @@ test "a lane host serves the core's own tokens, in order, and cancels between ro } /// A learner that takes `total` steps over a lesson, then reports it learned in that many. -const FakeLearner = struct { +pub const FakeLearner = struct { sink: ?api.LearnSink = null, steps: u32 = 0, total: u32 = 3, @@ -266,11 +266,11 @@ const FakeLearner = struct { fn emit(l: *FakeLearner, event: api.LearnEvent) void { l.sink.?.event(l.sink.?.ctx, &event); } - fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { + pub fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.* = .{ .sink = sink, .total = l.total, .examples = request.train.len }; } - fn step(ctx: *anyopaque) api.Learner.Step { + pub fn step(ctx: *anyopaque) api.Learner.Step { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.steps += 1; if (l.steps < l.total) return .{ .done = false, .changed = false }; @@ -278,20 +278,20 @@ const FakeLearner = struct { l.emit(.{ .done = .{} }); return .{ .done = true, .changed = true }; } - fn abort(ctx: *anyopaque) void { + pub fn abort(ctx: *anyopaque) void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.emit(.{ .done = .{ .message = "aborted" } }); } }; /// Learn events as tags, read on the test's thread once `done` arrives. -const LearnBox = struct { +pub const LearnBox = struct { mutex: std.Io.Mutex = .init, tags: std.ArrayList(std.meta.Tag(api.LearnEvent)) = .empty, steps: u32 = 0, done: bool = false, - fn sink(b: *LearnBox) api.LearnSink { + pub fn sink(b: *LearnBox) api.LearnSink { return .{ .ctx = b, .event = event }; } fn event(ctx: *anyopaque, e: *const api.LearnEvent) void { @@ -302,7 +302,7 @@ const LearnBox = struct { if (e.* == .learned) b.steps = e.learned.steps; if (e.* == .done) b.done = true; } - fn wait(b: *LearnBox) void { + pub fn wait(b: *LearnBox) void { while (true) { b.mutex.lockUncancelable(std.testing.io); const d = b.done; diff --git a/zig/src/core/lane_review_test.zig b/zig/src/core/lane_review_test.zig new file mode 100644 index 0000000000..04e6302041 --- /dev/null +++ b/zig/src/core/lane_review_test.zig @@ -0,0 +1,239 @@ +//! The lane host under a family round planner: prompt pieces, failed piece runs, learning, refusal. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Memory = api.Memory; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const Engine = api.Engine; +const FakeLearner = @import("lane_host_test.zig").FakeLearner; +const LearnBox = @import("lane_host_test.zig").LearnBox; + +test "more than sixteen planned pieces retain their job until its final" { + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true, .piece_runs = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.runBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + var host = LaneHost.init(std.heap.page_allocator, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + for (0..x.n - 1) |row| try pieces.append(std.heap.page_allocator, .{ .key = x.key, .start = row, .end = row + 1, .save = false }); + try finals.append(std.heap.page_allocator, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1: [18]u32 = @splat(3); + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // The first request alone needs seventeen pieces before its final. + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.length, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(usize, 2), target.joins); + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); + try std.testing.expectEqual(@as(usize, 4), target.runs); +} + +test "halt transfers queued and admitted jobs when every subsequent allocation fails" { + const gpa = std.testing.allocator; + 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 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var failing = std.testing.FailingAllocator.init(gpa, .{}); + var host = LaneHost.init(failing.allocator(), std.testing.io, &core, .{ .lanes = 2 }); + defer host.stop(); + const Box = struct { + finished: usize = 0, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + if (e.* == .finished and e.finished.reason == .failed) b.finished += 1; + } + }; + var box: Box = .{}; + const request: Request = .{ .prompt = &.{ 1, 2, 3 }, .max_tokens = 2 }; + const engine = host.engine(); + try engine.submit(1, &request, .{ .ctx = &box, .event = Box.event }); + try engine.submit(2, &request, .{ .ctx = &box, .event = Box.event }); + try host.admitted.append(host.gpa, host.queued.pop().?); + failing.fail_index = failing.alloc_index; + host.closing = true; + @import("lane_halt.zig").haltAll(&host, "drain deadline"); + try std.testing.expectEqual(@as(usize, 2), box.finished); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len + host.admitted.items.len); +} + +test "failed repeated pieces drop their job once and skip its remaining batches" { + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true, .piece_runs = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_runs = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.runBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + var host = LaneHost.init(std.heap.page_allocator, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + for (0..x.n - 1) |row| try pieces.append(std.heap.page_allocator, .{ .key = x.key, .start = row, .end = row + 1, .save = false }); + try finals.append(std.heap.page_allocator, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1: [18]u32 = @splat(3); + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // The first request alone needs seventeen pieces before its final. + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(usize, 1), target.joins); + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); + try std.testing.expectEqual(@as(usize, 3), target.runs); +} diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig new file mode 100644 index 0000000000..4ffe1ddb25 --- /dev/null +++ b/zig/src/core/lane_rounds.zig @@ -0,0 +1,289 @@ +//! A family's round planner under the lane host: prompts in pieces between the other streams' decode rounds. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const host = @import("lane_host.zig"); +const lane_stream = @import("lane_stream.zig"); +const learn = @import("lane_host_learn.zig"); +const status = @import("lane_host_status.zig"); +const ends = @import("lane_host_end.zig"); +const rows_host = @import("lane_host_rows.zig"); +const LaneHost = host.LaneHost; +const Job = LaneHost.Job; +const Allocator = std.mem.Allocator; + +/// A family's round planner: admissions, then `admitted`, begins, prompt pieces, a decode round, finals, `after`. +pub const Rounds = struct { + ctx: *anyopaque, + vtable: *const VTable, + + pub const Verdict = enum { admit, wait, skip, refuse }; + pub const Piece = struct { key: usize, start: u64, end: u64, save: bool, run: bool = true }; + pub const Ask = struct { key: usize, prompt: []const u32, max_tokens: u32, submitted_s: f64 }; + + pub const VTable = struct { + /// An admission pass starts (a lane is free and requests wait). + begin_admit: *const fn (ctx: *anyopaque) void, + /// The next queued request: start it, wait (it and those behind it), skip past it, or refuse it (`words`). + admit: *const fn (ctx: *anyopaque, ask: Ask, words: *[]const u8) Verdict, + /// The pass is decided: what runs before the admitted streams begin. + admitted: *const fn (ctx: *anyopaque) anyerror!void, + /// Around an admitted request's stream begin (`damaged`: its saved prefix could not be restored). + arm: *const fn (ctx: *anyopaque, key: usize) void, + began: *const fn (ctx: *anyopaque, key: usize, damaged: bool) void, + /// This round's pieces, in order, and the requests whose prompts are in. + plan: *const fn (ctx: *anyopaque, pieces: *std.ArrayList(Piece), finals: *std.ArrayList(usize)) anyerror!void, + /// A request left: finished, cancelled, failed or refused, started or not. + left: *const fn (ctx: *anyopaque, key: usize) void, + /// After the round: the pieces' seconds and the decode round's (finals included). + after: *const fn (ctx: *anyopaque, piece_s: f64, decode_s: f64) void, + }; +}; + +/// Makes and frees a request's copy proposer (`prompt` and `eos` outlive it). +pub const Proposers = struct { + ctx: *anyopaque, + make: *const fn (ctx: *anyopaque, gpa: Allocator, prompt: []const u32, eos: []const u32) anyerror!lanes.proposer.Proposer, + free: *const fn (ctx: *anyopaque, p: lanes.proposer.Proposer) void, +}; + +/// One round under the family's planner; false when it did nothing (nothing to admit, prefill or decode). +fn round(h: *LaneHost, rd: Rounds) bool { + var did = false; + // the admission pass, decided whole before any of it runs + var admits: std.ArrayList(*Job) = .empty; + defer admits.deinit(h.gpa); + var refused: std.ArrayList(struct { job: *Job, why: []const u8 }) = .empty; + defer refused.deinit(h.gpa); + h.lock(); + if (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes) { + rd.vtable.begin_admit(rd.ctx); + var free = h.info_.lanes - h.admitted.items.len; + var i: usize = 0; + while (i < h.queued.items.len and free > 0) { + const job = h.queued.items[i]; + var why: []const u8 = ""; + const v = rd.vtable.admit(rd.ctx, .{ .key = @intFromPtr(job), .prompt = job.request.prompt, .max_tokens = job.request.max_tokens, .submitted_s = seconds(job.submitted) }, &why); + switch (v) { + .admit => { + if (h.admitted.append(h.gpa, job)) |_| {} else |_| break; + if (admits.append(h.gpa, job)) |_| {} else |_| { + _ = h.admitted.pop(); + break; + } + _ = h.queued.orderedRemove(i); + free -= 1; + }, + .wait => break, + .skip => i += 1, + .refuse => { + if (refused.append(h.gpa, .{ .job = job, .why = why })) |_| {} else |_| break; + _ = h.queued.orderedRemove(i); + }, + } + } + } + h.unlock(); + for (refused.items) |x| h.finish(x.job, .failed, x.why); + rd.vtable.admitted(rd.ctx) catch |e| { + ends.failAll(h, h.words(e)); + return true; + }; + for (admits.items) |job| { + did = true; + if (!lane_stream.open(h, job)) continue; + rd.vtable.arm(rd.ctx, @intFromPtr(job)); + const b = h.core.beginStream(&job.stream) catch |e| { + _ = h.drop(job, h.words(e)); + continue; + }; + rd.vtable.began(rd.ctx, @intFromPtr(job), b.damaged); + } + // this round's pieces, a decode round of the streams decoding, then the prompts whose rows are in + rd.vtable.plan(rd.ctx, &h.pieces, &h.finals) catch |e| { + ends.failAll(h, h.words(e)); + return true; + }; + const t0 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + if (h.core.cfg.piece_runs and h.core.backend.vtable.pieces != null) { + // bounded calls for the round's pieces; a failed batch drops each affected job once + var list: [16]lanes.backend.Backend.Piece = undefined; + var jobs: [16]*Job = undefined; + var at: usize = 0; + while (at < h.pieces.items.len) { + var n: usize = 0; + while (at < h.pieces.items.len and n < list.len) : (at += 1) { + const piece = h.pieces.items[at]; + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + list[n] = .{ .stream = &job.stream, .start = piece.start, .end = piece.end, .save = piece.save }; + jobs[n] = job; + n += 1; + } + if (n > 0) h.core.pieces(list[0..n]) catch |e| { + for (jobs[0..n], 0..) |job, i| { + const seen = for (jobs[0..i]) |previous| { + if (previous == job) break true; + } else false; + if (!seen) dropPiece(h, job, h.words(e)); + } + }; + } + } else for (h.pieces.items) |piece| { + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + h.core.piece(&job.stream, piece.start, piece.end, piece.save) catch |e| { + dropPiece(h, job, h.words(e)); + }; + } + const t1 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + // Config.join_tail: a joined prompt's last token decodes in this round's window; else its own runs after + const join = h.core.cfg.join_tail; + if (join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + if (h.core.activeCount() > 0) { + did = true; + h.core.step() catch |e| { + ends.failAll(h, @errorName(e)); + return true; + }; + } + if (!join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + const t2 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + var i: usize = 0; + while (i < h.admitted.items.len) { + const job = h.admitted.items[i]; + if (job.started and h.deliver(job)) continue; // it left the list: the next job slid into slot i + i += 1; + } + rd.vtable.after(rd.ctx, seconds(t1 - t0), seconds(t2 - t1)); + return did; +} + +/// A job whose piece failed: dropped and taken out of this round's finals (`drop` frees the job). +fn dropPiece(h: *LaneHost, job: *Job, message: []const u8) void { + const key = @intFromPtr(job); + for (h.pieces.items) |*piece| if (piece.key == key) { + piece.run = false; + }; + var i: usize = 0; + while (i < h.finals.items.len) { + if (h.finals.items[i] == key) _ = h.finals.orderedRemove(i) else i += 1; + } + _ = h.drop(job, message); +} + +/// The prompts whose rows are in: each stream's last row (Engine.finishStream), then it decodes with the rounds. +fn runFinals(h: *LaneHost) void { + if (h.core.cfg.replay_runs) { + // the finals' shared work first (one forward for all); a failure fails each + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var n: usize = 0; + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished or n == ss.len) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + h.core.prepareFinals(ss[0..n]) catch |e| { + for (jobs[0..n]) |job| _ = h.drop(job, h.words(e)); + return; + }; + } + if (h.core.cfg.join_tail and h.core.cfg.join_drafts) { + // the round's finals together: the joined prompts' first drafts in one pass (Engine.finishStreams) + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var errs: [16]?anyerror = undefined; + var at: usize = 0; + while (at < h.finals.items.len) { + var n: usize = 0; + while (at < h.finals.items.len and n < ss.len) : (at += 1) { + const job: *Job = @ptrFromInt(h.finals.items[at]); + if (job.stream.finished) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + const joined = h.core.finishStreams(ss[0..n], errs[0..n]) catch |e| blk: { + for (errs[0..n]) |*x| x.* = e; + break :blk 0; + }; + // a round's line when its finished prompts' first drafts joined in one pass + if (joined > 0) std.log.info("join drafts: {d} of {d} finished prompts drafted in one pass", .{ joined, n }); + for (jobs[0..n], errs[0..n]) |job, x| { + if (x) |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + } + h.prefilled(job, job.began); + } + } + return; + } + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished) continue; + h.core.finishStream(&job.stream) catch |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + }; + h.prefilled(job, job.began); + } +} + +/// The engine thread under a round planner; an idle round waits for an arrival, a cancel or a short timeout. +pub fn run(h: *LaneHost, rd: Rounds) void { + while (true) { + ends.takeCancels(h); + rows_host.answer(h); + learn.refuse(h); + h.lock(); + if (h.closing) { + const left = h.queued.items.len + h.admitted.items.len; + h.unlock(); + if (h.lesson_open) { + h.learner.?.abort(h.learner.?.ctx); + h.lesson_open = false; + } + if (left == 0) return rows_host.flush(h); + ends.closeAll(h); + continue; + } + const work = h.admitted.items.len > 0 or (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes); + if (!work and learn.ready(h)) { + h.unlock(); + learn.step(h); + continue; + } + if (!work) { + if (rows_host.idle(h)) continue; + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {}; + h.unlock(); + continue; + } + h.unlock(); + const did = round(h, rd); + status.noteLive(h); + if (!did) { + h.lock(); + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(20), .clock = .awake } }) catch {}; + h.unlock(); + } + } +} + +fn seconds(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e9; +} diff --git a/zig/src/core/lane_rounds_test.zig b/zig/src/core/lane_rounds_test.zig new file mode 100644 index 0000000000..4a59e38040 --- /dev/null +++ b/zig/src/core/lane_rounds_test.zig @@ -0,0 +1,392 @@ +//! The lane host under a family round planner: prompt pieces, failed piece runs, learning, refusal. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Memory = api.Memory; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const Engine = api.Engine; +const FakeLearner = @import("lane_host_test.zig").FakeLearner; +const LearnBox = @import("lane_host_test.zig").LearnBox; + +test "a round planner: prompts in pieces between the other streams' decode rounds, queued admissions, the same tokens" { + const gpa = std.testing.allocator; + 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 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round, a final when its rows are in + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + waits: u32 = 0, + rounds: u32 = 0, + mixed: u32 = 0, // rounds that ran a piece while another stream decoded + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) { + p.waits += 1; + return .wait; + }; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + var decoding = false; + for (p.seqs.items) |x| decoding = decoding or x.decoding; + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + if (decoding) p.mixed += 1; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) !Reason { + for (0..10_000) |_| { // 10 s + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8, 1, 8, 2, 8, 4, 5, 9 }; + var b1: Box = .{}; + var b2: Box = .{}; + defer b1.tokens.deinit(gpa); + defer b2.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &p1, .max_tokens = 40 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both queued before the first round, so the second waits whatever the scheduler does + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + { + errdefer host.stop(); // a timed-out wait stops the host thread before this frame goes + try std.testing.expectEqual(Reason.length, try b1.wait()); + try std.testing.expectEqual(Reason.length, try b2.wait()); + } + host.stop(); + // the replies are the whole-prompt path's tokens + for ([_]struct { p: []const u32, b: *Box }{ .{ .p = &p1, .b = &b1 }, .{ .p = &p2, .b = &b2 } }) |c| { + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, c.p); + for (c.b.tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + } + try std.testing.expectEqual(@as(usize, 40), b1.tokens.items.len); + try std.testing.expectEqual(@as(usize, 12), b2.tokens.items.len); + // the second prompt waited queued, then ran its pieces while the first decoded + try std.testing.expect(pl.waits > 0 and pl.mixed >= 3); + try std.testing.expectEqual(@as(usize, 3 + 4), target.pieces_run); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); +} + +test "a failed piece run fails its requests alone: their finals of the same round do not run, the host keeps serving" { + // a run's error drops its jobs, whose finals are listed in the same round; runFinals must not read them + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true, .piece_runs = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_runs = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.runBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + var host = LaneHost.init(std.heap.page_allocator, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(std.heap.page_allocator, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(std.heap.page_allocator, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: one run of their pieces, which fails + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.failed, b2.wait()); + try std.testing.expectEqual(@as(usize, 0), target.joins); // no final of theirs ran + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); + try std.testing.expectEqual(@as(usize, 1), target.runs); +} + +test "under a round planner a learn request still steps while no request waits or runs" { + const gpa = std.testing.allocator; + 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 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + const Idle = struct { + fn none(_: *anyopaque) void {} + fn admit(_: *anyopaque, _: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + return .wait; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn key(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(_: *anyopaque, _: *std.ArrayList(api.Rounds.Piece), _: *std.ArrayList(usize)) anyerror!void {} + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var ctx: u8 = 0; + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + host.rounds = .{ .ctx = &ctx, .vtable = &.{ .begin_admit = Idle.none, .admit = Idle.admit, .admitted = Idle.admitted, .arm = Idle.key, .began = Idle.began, .plan = Idle.plan, .left = Idle.key, .after = Idle.after } }; + var fake: FakeLearner = .{}; + host.learner = .{ .ctx = &fake, .begin = FakeLearner.begin, .step = FakeLearner.step, .abort = FakeLearner.abort }; + try host.start(); + defer host.stop(); + const example: api.Example = .{ .ids = &.{ 1, 2, 3 }, .start = 2, .decide = 0 }; + const learn: api.LearnRequest = .{ .train = &.{example} }; + var box: LearnBox = .{}; + defer box.tags.deinit(gpa); + try host.engine().learn(&learn, box.sink()); + box.wait(); + try std.testing.expectEqualSlices(std.meta.Tag(api.LearnEvent), &.{ .learned, .done }, box.tags.items); + try std.testing.expectEqual(@as(u32, 3), box.steps); +} + +test "a round planner over a backend without prompt pieces is refused at start" { + const gpa = std.testing.allocator; + 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 }; + defer target.deinit(); + var whole = target.backend(); + var vt = whole.vtable.*; + vt.begin = null; + whole.vtable = &vt; + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, whole, clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var ctx: u8 = 0; + host.rounds = .{ .ctx = &ctx, .vtable = undefined }; + try std.testing.expectError(error.PiecesUnsupported, host.start()); +} + +test "a failed piece fails its request alone: its final of the same round does not run, the host keeps serving" { + // a job dropped for a failed piece must leave the round's finals: `drop` frees it, runFinals must not read it + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_pieces = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + const pa = std.heap.page_allocator; + var host = LaneHost.init(pa, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(pa, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(pa, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: the first piece fails, the second's request is served + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(usize, 12), b2.tokens); + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); +} diff --git a/zig/src/core/lane_stream.zig b/zig/src/core/lane_stream.zig new file mode 100644 index 0000000000..8330912ac1 --- /dev/null +++ b/zig/src/core/lane_stream.zig @@ -0,0 +1,69 @@ +//! A request's stream on the lane host: its kept prefix looked up, its drafter made, the stream set up. +const std = @import("std"); +const lanes = @import("lanes"); +const recalls = @import("lane_host_recall.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Job = LaneHost.Job; + +/// A job's drafters freed: the family's (LaneHost.proposers) or the built-in one, and its Recall reader. +pub fn freeProposer(h: *LaneHost, job: *Job) void { + if (job.custom) |c| h.proposers.?.free(h.proposers.?.ctx, c) else job.proposer.deinit(); + job.custom = null; + recalls.release(job); +} + +/// An admitted job's stream (its kept prefix looked up, its drafter): false when it failed (the job is finished). +pub fn open(h: *LaneHost, job: *Job) bool { + const r = job.request; + var reuse: lanes.stream.Reuse = .{}; + // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass + if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), r.bornLen(), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + job.entry = l.entry; + job.kept0 = store.counts.kept; + job.marks = l.marks; + reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; + } else |_| {}; + if (h.proposers) |ps| { + job.custom = ps.make(ps.ctx, h.gpa, r.prompt, r.eos) catch return !h.drop(job, "the drafter could not start"); + } else job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return !h.drop(job, "the drafter could not start"); + const recall = recalls.reader(h, job); + job.stream = lanes.Stream.init(h.gpa, .{ + .id = "request", + .prompt = r.prompt, + .max_new = r.max_tokens, + .eos = r.eos, + .sampling = r.sampling, + .drafts = r.drafts, + .proposer = job.custom orelse job.proposer.proposer(), + .recall = recall, + .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, + .cancel_check = .{ .ptr = job, .check = LaneHost.cancelled }, + .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = LaneHost.yieldsTo } else null, + .think_budget = r.think_budget, + .think_close = r.think_close, + .think_end = if (r.think_end) |t| t else -1, + .logprobs = r.logprobs, + .prompt_logprobs = r.prompt_logprobs, + .loop_guard = r.loop_guard, + .chunks = r.chunks, + .history_len = r.history_len, + .shared_prefixes = r.shared_prefixes, + .reuse = reuse, + .grammar = if (r.structure) |st| h.constraint(st) catch |e| { + freeProposer(h, job); + return !h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + } else null, + .decode_spans = h.spans(r), + .decide = r.decide, + .untaught = r.untaught, + .images = r.images, + .key = r.cache_key, + }) catch { + freeProposer(h, job); + return !h.drop(job, "out of memory"); + }; + job.started = true; + job.fill_on.store(true, .release); // status readers count the fill from here until the pass reports + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 8378c44900..10260248d4 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -104,8 +104,28 @@ pub const Backend = struct { reports_fill: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, + /// A prompt in pieces between rounds (`Rounds`) instead of `prefill`: begin, pieces, finish; null: whole only. + begin: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!Begun = null, + piece: ?*const fn (ptr: *anyopaque, s: *Stream, start: u64, end: u64, save: bool) anyerror!void = null, + finish: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!void = null, + /// `finish` minus the last row (Config.join_tail): true when that token is left for the first round. + join: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// A round's pieces in one call (Config.piece_runs), maybe one forward; a failure fails them all. + pieces: ?*const fn (ptr: *anyopaque, list: []const Piece) anyerror!void = null, + /// Before a round's `finish` / `join` calls (Config.replay_runs): their shared work in one go. + finals: ?*const fn (ptr: *anyopaque, streams: []const *Stream) anyerror!void = null, }; + /// One stream's prompt rows [start, end) of a round (`save`: the prompt's snapshot right after them). + pub const Piece = struct { stream: *Stream, start: u64, end: u64, save: bool }; + + /// `begin`'s outcome: restored prefix tokens, or a damaged saved entry dropped (the prompt runs from 0). + pub const Begun = struct { cached: u64 = 0, damaged: bool = false }; + + pub fn pieces(b: Backend) bool { + return b.vtable.begin != null and b.vtable.piece != null and b.vtable.finish != null; + } + pub fn prefill(b: Backend, s: *Stream) !void { return b.vtable.prefill(b.ptr, s); } diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 85361bd839..c215c87f86 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -32,6 +32,10 @@ pub const Model = struct { draft_probabilities: bool = false, // the head gives each drafted node's chance of landing draft_streams: bool = false, // the head drafts every stream of a shared round in one batch head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them + join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join) + piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them) + join_drafts: bool = false, // with join_tail: a round's joined prompts' first drafts in one call (Engine.finishStreams) + replay_runs: bool = false, // a round's finals go to the backend together first (Backend.finals) lane_costs: []const Cost = &.{}, // a lone stream's window ms by rows, up to its widest tree (timed at load) lone_copies: bool = true, // a lone stream proposes copies (false: only streams in shared rounds do) }; @@ -55,9 +59,18 @@ pub const Config = struct { fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round draft_streams: bool, head_trees: bool, + /// a prompt in pieces leaves its last token as the pending row of its first shared round's window + join_tail: bool = false, + /// a round's pieces of several streams in one backend call (Backend.pieces), which may run them as one forward + piece_runs: bool = false, + /// with join_tail: a round's joined prompts get their first drafts in one backend call + join_drafts: bool = false, + /// a round's finished prompts to the backend together before their finishes (Backend.finals) + replay_runs: bool = false, head_lanes: u32 = 0, // a lone stream's lanes at most, the planner picking each round's (0: the depth rule's chain) head_depth: u32 = 16, // and their depth at most own_prices: bool = true, // the planner prices each pick from the stream's own measured rounds (seeded from the table) + timed: bool = true, // false: rounds priced from the cost tables alone, no wall clock (a request drafts the same whatever the load or the requests before it) depth_prior: []f64, family_costs: Table, shared_costs: Table, @@ -115,6 +128,10 @@ pub const Config = struct { .node_probabilities = m.draft_probabilities, .draft_streams = m.draft_streams, .head_trees = m.head_trees, + .join_tail = m.join_tail, + .piece_runs = m.piece_runs, + .join_drafts = m.join_drafts, + .replay_runs = m.replay_runs, .depth_prior = depth_prior, .family_costs = family_costs, .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows), diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index f11d8dbec5..0e58efcd78 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -13,6 +13,7 @@ const plan_lanes = @import("plan_lanes.zig"); const fill = @import("fill.zig"); const trail = @import("trail.zig"); const open = @import("open.zig"); +const trim = @import("trim.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -26,11 +27,20 @@ const Result = struct { rows: u32, keep: u32, got: []const u32 }; const Unread = struct { event: usize, handle: u64 }; pub const Engine = struct { + const pieces_ = @import("pieces.zig"); + pub const beginStream = pieces_.beginStream; + pub const piece = pieces_.piece; + pub const pieces = pieces_.pieces; + pub const prepareFinals = pieces_.prepareFinals; + pub const finishStream = pieces_.finishStream; + pub const finishStreams = pieces_.finishStreams; + gpa: Allocator, cfg: *const Config, backend: be.Backend, clock: be.Clock, log: ?*ev.Log = null, + trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig) rule: depth.Rule, live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`) drafted: u64 = 0, @@ -209,7 +219,7 @@ pub const Engine = struct { s.priced = null; const a = e.arena.allocator(); var plans = [_]Plan{try win.plan(e, s, copied)}; - if (plans[0].kind == .head and e.cfg.node_probabilities) try win.allocate(e, &plans); + if (plans[0].kind == .head and (e.cfg.node_probabilities or e.trim != null)) try win.allocate(e, &plans); try win.constrain(e, &plans[0]); const plan = plans[0]; const early = e.cfg.family_mtp and s.drafts and plan.kind != .forced and e.cfg.speculate_early and plan.parents == null and s.grammar == null; @@ -231,14 +241,14 @@ pub const Engine = struct { if (early) if (e.backend.vtable.unspeculate) |undo| try undo(e.backend.ptr, s); try e.draftLate(s, plan.position, o.follow, o.path, null); } - if (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard)) { + if (e.cfg.timed and (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard))) { const ms = e.clock.elapsedMs(.cost); const first = s.rounds == 1; // a stream's first round carries the prefill-to-decode switch try trail.event(e, &.{ f("ev", str("cost")), f("stream", str(s.id)), f("drafts", int(rows - 1)), f("init", .{ .bool = first }), f("ms", ev.bits(ms)) }); try e.rule.observeCost(@intCast(rows - 1), ms, first, win.who(s)); } } - if (s.graft_room != null and s.rounds > 1) { + if (e.cfg.timed and s.graft_room != null and s.rounds > 1) { // the stream's round time beyond its window and head levels, for the lane planner if (e.cfg.lane_costs.get(@intCast(rows))) |w| { const over = e.clock.elapsedMs(.cost) - w - e.cfg.mtp_step_ms * @as(f64, @floatFromInt(held_levels)); @@ -248,7 +258,7 @@ pub const Engine = struct { // the pick that shaped this round, priced from its own measured rounds next time if (priced) |p| s.measured[p.choice].add(e.clock.elapsedMs(.cost) - p.ms); } - if (plan.kind == .head) { + if (e.cfg.timed and plan.kind == .head) { const ms = e.clock.elapsedMs(.overhead); try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(1)), f("rows", int(rows)), f("ms", ev.bits(ms)) }); try e.rule.observeOverhead(1, rows, ms, s.sampling != null, 0.0); @@ -306,9 +316,11 @@ pub const Engine = struct { for (heads.items, budgets) |i, d| try e.draftLate(plans[i].stream, plans[i].position, outcomes[i].follow, outcomes[i].path, d); } try e.backend.keep(windows, paths); - const ms = e.clock.elapsedMs(.overhead); - try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); - try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets)); + if (e.cfg.timed) { + const ms = e.clock.elapsedMs(.overhead); + try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); + try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets)); + } } /// Commit the target-sampled path (after the thinking budget's cut), each token with its path row's logprob row. @@ -326,6 +338,7 @@ pub const Engine = struct { path = try accept.acceptPath(a, tokens, parents, sampled); for (path[1..]) |r| try committed.append(a, tokens[r]); try committed.append(a, sampled[path[path.len - 1]]); + if (s.proposer) |pr| pr.round(@intCast(rows), @intCast(path.len - 1)); if (rows > 1) { const kept = path.len - 1; e.drafted += rows - 1; @@ -379,6 +392,10 @@ pub const Engine = struct { } if (s.recall) |r| if (p.kind != .copy) r.noteHead(path.len); // its own rounds: the bar a copy must beat } + if (e.trim) |t| if (p.kind != .forced) { + const chain = parents.len == 0 or accept.isChain(parents); + try t.commit(.{ .stream = s, .drafted = p.kind == .head, .rows = @intCast(rows), .parents = if (chain) null else parents, .path = path, .bonus = sampled[path[path.len - 1]] }); + }; const budget_cut = s.thinkCut(committed.items); const loop_cut = sm.loopCut(s, committed.items); const loop_wins = if (loop_cut) |c| budget_cut == null or c < budget_cut.? else false; @@ -435,7 +452,7 @@ pub const Engine = struct { const s = requests[0].stream; if (s.odds == null) s.odds = shape.Odds.init(shape.start_odds, shape.start_weight); const cap: usize = @intCast(@max(1, @min(@min(shape.max_depth, e.cfg.head_depth), s.draftRoom() - 1))); - const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices) &s.measured else null }; + const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices and e.cfg.timed) &s.measured else null }; // every eighth round one size wider than the best, so the odds of deeper lanes stay measured const p = try plan_lanes.pick(e.gpa, s.odds.?, costs, e.cfg.head_lanes, cap, try e.graftProspect(s), s.rounds % 8 == 4); s.picks[p.choice] += 1; diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig index c17758753b..4e358fd267 100644 --- a/zig/src/core/lanes/engine_test.zig +++ b/zig/src/core/lanes/engine_test.zig @@ -460,3 +460,124 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).tokens.len >= 2); } } + +/// The cases' prompts in pieces, case i opened before round 2 i, others decoding; `join`: Config.join_tail. +fn runPieces(cases: []const Case, join: bool, joins: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = join; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + var opened: usize = 0; + var round: usize = 0; + while (opened < streams.len or engine.activeCount() > 0) : (round += 1) { + if (opened < streams.len and round == 2 * opened) { + const s = &streams[opened]; + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + try engine.finishStream(s); + opened += 1; + } + if (engine.activeCount() > 0) try engine.step(); + } + joins.* = target.joins; + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a prompt's last token joins its first round's shared window: the same tokens, greedy and sampled" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var j0: usize = 0; + var j1: usize = 0; + const off = try runPieces(&cases, false, &j0); + defer free(off); + const on = try runPieces(&cases, true, &j1); + defer free(on); + try std.testing.expectEqual(@as(usize, 0), j0); + try std.testing.expectEqual(@as(usize, 3), j1); // not the undrafted stream nor the budget cutting the first token + for (off, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // and the whole prompts' tokens + const whole = try run(&cases); + defer free(whole); + for (whole, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); +} + +/// Every prompt in one piece, all finished in one round (`together`: Engine.finishStreams), then rounds until done. +fn runFinals(cases: []const Case, together: bool, calls: *usize, most: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = true; + cfg.join_drafts = together; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + const ptrs = try gpa.alloc(*sm.Stream, cases.len); + defer gpa.free(ptrs); + for (streams, ptrs) |*s, *p| { + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + p.* = s; + } + // the finish phase alone: its draft calls and their largest + target.drafts_calls = 0; + target.drafts_most = 0; + if (together) { + const errs = try gpa.alloc(?anyerror, cases.len); + defer gpa.free(errs); + _ = try engine.finishStreams(ptrs, errs); + for (errs) |x| if (x) |e| return e; + } else for (ptrs) |s| try engine.finishStream(s); + calls.* = target.drafts_calls; + most.* = target.drafts_most; + while (engine.activeCount() > 0) try engine.step(); + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a round's finished prompts get their first drafts in one call (Python's _windows), the same tokens" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var c1: usize = 0; + var m1: usize = 0; + var c0: usize = 0; + var m0: usize = 0; + const together = try runFinals(&cases, true, &c1, &m1); + defer free(together); + const apart = try runFinals(&cases, false, &c0, &m0); + defer free(apart); + for (apart, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + const whole = try run(&cases); + defer free(whole); + for (whole, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // three joined prompts in one call, the budget-cut one its own, the undrafted none: 2 calls against 4 + try std.testing.expectEqual(@as(usize, 2), c1); + try std.testing.expectEqual(@as(usize, 3), m1); + try std.testing.expectEqual(@as(usize, 4), c0); + try std.testing.expectEqual(@as(usize, 1), m0); +} diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index 7827f286ac..8c3ed43d01 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -53,6 +53,7 @@ pub const Fake = struct { answer_cycles: bool = false, prefill_chunks: usize = 0, prefill_count: usize = 0, + fail_pieces: usize = 0, // the next this many `piece` calls fail with error.Unbound prefill_hook: ?*const fn (ctx: *anyopaque, s: *Stream, chunk: usize) void = null, prefill_hook_ctx: ?*anyopaque = null, refuse_sampled: bool = false, // prefill refuses a sampled stream with error.SamplingRefused @@ -61,6 +62,12 @@ pub const Fake = struct { fail_round: ?u64 = null, // this round's verify fails with error.GpuFailed, as a round the GPU fails does last_rows: usize = 0, // the last verify's rows, for test clocks last_sampled: bool = false, // and whether one of its streams drew keyed samples + pieces_run: usize = 0, // prompt pieces run (`pieceBackend`) + joins: usize = 0, // prompts whose last token joined a round (`joinBackend`) + drafts_calls: usize = 0, // `draft` calls, and the most requests one took + drafts_most: usize = 0, + fail_runs: usize = 0, // the next this many `pieces` calls fail with error.Unbound (`runBackend`) + runs: usize = 0, // `pieces` calls that ran pub fn deinit(x: *Fake) void { var it = x.lanes.valueIterator(); @@ -84,6 +91,66 @@ 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 } }; } + /// 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 } }; + } + + /// `pieceBackend` whose prompts may leave their last token to the stream's first round (Config.join_tail). + pub fn joinBackend(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, .join = join } }; + } + + /// `joinBackend` that takes a round's pieces in one call (Config.piece_runs): a piece at a time, or a failure. + pub fn runBackend(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, .join = join, .pieces = pieces } }; + } + + fn pieces(ptr: *anyopaque, list: []const be.Backend.Piece) anyerror!void { + const x = self(ptr); + if (x.fail_runs > 0) { + x.fail_runs -= 1; + return error.Unbound; + } + x.runs += 1; + for (list) |p| try piece(ptr, p.stream, p.start, p.end, p.save); + } + + fn join(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + if (x.lane(s).history.items.len + 1 != s.prompt_len) return error.PositionMismatch; + x.joins += 1; + return true; + } + + fn begin(ptr: *anyopaque, s: *Stream) anyerror!be.Backend.Begun { + const x = self(ptr); + const got = try x.lanes.getOrPut(x.gpa, s); + if (got.found_existing) freeLane(x.gpa, got.value_ptr); + got.value_ptr.* = .{}; + return .{}; + } + + fn piece(ptr: *anyopaque, s: *Stream, start: u64, end: u64, snapshot: bool) anyerror!void { + _ = snapshot; + const x = self(ptr); + if (x.fail_pieces > 0) { + x.fail_pieces -= 1; + return error.Unbound; + } + const h = &x.lane(s).history; + if (h.items.len != start or end + 1 > s.prompt_len) return error.PositionMismatch; + try h.appendSlice(x.gpa, s.prompt()[start..end]); + x.pieces_run += 1; + } + + fn finish(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const h = &x.lane(s).history; + if (h.items.len + 1 != s.prompt_len) return error.PositionMismatch; + try h.append(x.gpa, s.prompt()[s.prompt_len - 1]); + } + fn self(ptr: *anyopaque) *Fake { return @ptrCast(@alignCast(ptr)); } @@ -295,6 +362,8 @@ pub const Fake = struct { fn draft(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { const x = self(ptr); + x.drafts_calls += 1; + x.drafts_most = @max(x.drafts_most, requests.len); for (requests) |r| { const l = x.lane(r.stream); // the head reads the last verify's kept rows (a shared round drafts before its rollback) diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index e943bf9d46..6fbd946441 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -16,6 +16,7 @@ pub const cost_rule = @import("cost_rule.zig"); pub const gpu_rule = @import("gpu_rule.zig"); pub const gpu_full = @import("gpu_full.zig"); pub const depth = @import("depth.zig"); +pub const trim = @import("trim.zig"); pub const config = @import("config.zig"); pub const stream = @import("stream.zig"); pub const backend = @import("backend.zig"); @@ -37,6 +38,7 @@ pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; pub const SuffixLookup = proposer.SuffixLookup; pub const LogprobRow = logprob.Row; +pub const Trim = trim.Trim; test { std.testing.refAllDecls(@This()); diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..cd832256ab 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -23,8 +23,16 @@ pub fn add(e: *Engine, s: *Stream) !void { try opened(e, s); } +/// How a stream's first drafts are sized: the depth rule, or the head's whole reach under node probabilities. +pub const Reach = enum { rule, head }; + /// After the prompt pass: the first token drawn and committed, the first drafts asked; rounds from the next. pub fn opened(e: *Engine, s: *Stream) !void { + return openedWith(e, s, .rule); +} + +/// `opened` with its first drafts sized by `reach` (a family planner's prompts in pieces: `.head`). +pub fn openedWith(e: *Engine, s: *Stream, reach: Reach) !void { if (s.isCancelled()) { // cancelled in its last chunk: no first token e.backend.release(s); return error.Cancelled; @@ -40,7 +48,7 @@ pub fn opened(e: *Engine, s: *Stream) !void { var asked: ?u32 = null; if (e.cfg.family_mtp and s.drafts) { // the head reads the prompt's last row and the first token, and drafts the one after it - const d: u32 = @intCast(try e.rule.depth(win.who(s))); + const d = try firstDepth(e, s, reach); asked = d; try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = feed, .rows = null, .start = s.prompt_len, .position = position + 1, .depth = d }}); s.dropHeld(e.gpa); @@ -70,6 +78,13 @@ pub fn opened(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } +/// The first drafts' depth for `reach`. +pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { + var who = win.who(s); + if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e))); + return @intCast(try e.rule.depth(who)); +} + /// Whether the backend can run a prompt pass a chunk at a time between rounds (`fillStream`). pub fn fills(e: *const Engine) bool { return e.backend.vtable.prefill_step != null; diff --git a/zig/src/core/lanes/pieces.zig b/zig/src/core/lanes/pieces.zig new file mode 100644 index 0000000000..cd64bb1eb0 --- /dev/null +++ b/zig/src/core/lanes/pieces.zig @@ -0,0 +1,144 @@ +//! Prompts in pieces on the lane core: a stream begun, its prompt rows run in pieces, then its first token. +const std = @import("std"); +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const ev = @import("events.zig"); +const trail = @import("trail.zig"); +const open = @import("open.zig"); +const Engine = @import("engine.zig").Engine; +const Stream = sm.Stream; +const f = ev.f; +const str = trail.str; +const int = trail.int; + +/// A stream whose prompt runs in pieces between rounds: `piece`s, then `finishStream`; a failed begin holds nothing. +pub fn beginStream(e: *Engine, s: *Stream) !be.Backend.Begun { + try trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }); + const begin = e.backend.vtable.begin orelse return error.NoPieces; + const b = try begin(e.backend.ptr, s); + s.cached = @intCast(b.cached); + return b; +} + +/// Prompt rows [start, end) of a begun stream (`save`: the prompt's snapshot right after them). +pub fn piece(e: *Engine, s: *Stream, start: u64, end: u64, save: bool) !void { + const run = e.backend.vtable.piece orelse return error.NoPieces; + return run(e.backend.ptr, s, start, end, save); +} + +/// A round's pieces of several streams in one backend call (Config.piece_runs, Backend.pieces); else one at a time. +pub fn pieces(e: *Engine, list: []const be.Backend.Piece) !void { + if (e.cfg.piece_runs) if (e.backend.vtable.pieces) |run| return run(e.backend.ptr, list); + for (list) |p| try piece(e, p.stream, p.start, p.end, p.save); +} + +/// A round's finished prompts before their `finishStream` calls (Config.replay_runs, Backend.finals). +pub fn prepareFinals(e: *Engine, streams: []const *Stream) !void { + if (!e.cfg.replay_runs or streams.len < 2) return; + const run = e.backend.vtable.finals orelse return; + return run(e.backend.ptr, streams); +} + +/// A begun stream's prompt is in: first token and drafts; a failure is the host's to discard (Cancelled: released). +pub fn finishStream(e: *Engine, s: *Stream) !void { + _ = e.arena.reset(.retain_capacity); + // Config.join_tail: the last token waits for the next round (a forced or budget-cut first token keeps its row) + if (e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext()) { + if (e.backend.vtable.join) |join| { + if (join(e.backend.ptr, s) catch |x| return released(e, s, x)) return joined(e, s); + return open.openedWith(e, s, .head); + } + } + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| return released(e, s, x); + return open.openedWith(e, s, .head); +} + +/// A round's finished prompts as `finishStream`, joined ones' first drafts in one call; returns how many joined. +pub fn finishStreams(e: *Engine, streams: []const *Stream, errs: []?anyerror) !usize { + _ = e.arena.reset(.retain_capacity); + const a = e.arena.allocator(); + var reqs: std.ArrayList(be.DraftRequest) = .empty; + var joins: std.ArrayList(usize) = .empty; + for (streams, errs, 0..) |s, *err, i| { + err.* = null; + const can = e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext(); + const join = if (can) e.backend.vtable.join else null; + if (join) |j| { + const ok = j(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + if (ok) { + try reqs.append(a, joinRequest(e, s) catch |x| { + err.* = x; + continue; + }); + try joins.append(a, i); + continue; + } + } else { + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + } + open.openedWith(e, s, .head) catch |x| { + err.* = x; + }; + } + if (reqs.items.len == 0) return 0; + e.backend.draft(reqs.items) catch |x| { + for (joins.items) |i| errs[i] = x; + return 0; + }; + for (reqs.items) |r| joinedDrafted(e, r) catch |x| { + for (joins.items) |i| if (streams[i] == r.stream) { + errs[i] = x; + }; + }; + return reqs.items.len; +} + +/// A cancelled stream's lane is released here: the host finishes it without the core. +fn released(e: *Engine, s: *Stream, err: anyerror) anyerror { + if (err == error.Cancelled) e.backend.release(s); + return err; +} + +/// A joined prompt (Backend.join): its last token pending, drafts asked now, verified in its first shared round. +fn joined(e: *Engine, s: *Stream) !void { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const position: u64 = s.prompt_len; + const d = try open.firstDepth(e, s, .head); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = position, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(d)), f("position", int(position)), f("follow", .{ .u32s = &.{last} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// After its first draft: the joined stream holds it and takes part from the next round. +fn joinedDrafted(e: *Engine, r: be.DraftRequest) !void { + const s = r.stream; + s.dropHeld(e.gpa); + s.next = .{ .count = r.depth }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(r.depth)), f("position", int(r.position)), f("follow", .{ .u32s = &.{s.pending.?} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// A joined prompt's state and its first draft request (the last token pending at the prompt's last position). +fn joinRequest(e: *Engine, s: *Stream) !be.DraftRequest { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const d = try open.firstDepth(e, s, .head); + return .{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = s.prompt_len, .depth = d }; +} diff --git a/zig/src/core/lanes/proposer.zig b/zig/src/core/lanes/proposer.zig index f49626e74a..3df9a73393 100644 --- a/zig/src/core/lanes/proposer.zig +++ b/zig/src/core/lanes/proposer.zig @@ -11,6 +11,10 @@ pub const Proposer = struct { propose: *const fn (ptr: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32, last_match: *const fn (ptr: *anyopaque) i64, observe: ?*const fn (ptr: *anyopaque, proposed: i64, accepted: i64) void = null, + /// The proposer prices its own drafts: `propose`'s result is taken as is, no width, length or match gate. + priced: bool = false, + /// Every non-forced round of the stream, whatever drafted it: the window's rows and its drafts kept. + round: ?*const fn (ptr: *anyopaque, rows: u32, kept: u32) void = null, }; /// Tokens valid until the next `propose`. @@ -25,6 +29,10 @@ pub const Proposer = struct { pub fn observe(p: Proposer, proposed: i64, accepted: i64) void { if (p.vtable.observe) |f| f(p.ptr, proposed, accepted); } + + pub fn round(p: Proposer, rows: u32, kept: u32) void { + if (p.vtable.round) |f| f(p.ptr, rows, kept); + } }; pub const max_ngram = 8; diff --git a/zig/src/core/lanes/recall_engine_test.zig b/zig/src/core/lanes/recall_engine_test.zig index 2aab54e6a2..0fb9456102 100644 --- a/zig/src/core/lanes/recall_engine_test.zig +++ b/zig/src/core/lanes/recall_engine_test.zig @@ -6,6 +6,8 @@ const sm = @import("stream.zig"); const fake = @import("fake.zig"); const recall = @import("recall.zig"); const SuffixLookup = @import("proposer.zig").SuffixLookup; +const Proposer = @import("proposer.zig").Proposer; +const win = @import("windows.zig"); const Sampling = @import("sampling.zig").Sampling; const gpa = std.testing.allocator; @@ -222,3 +224,62 @@ test "Recall's continuations graft beside the head's chain and keep the reply's try std.testing.expectEqualSlices(u32, plain.emitted[0], with.emitted[0]); try std.testing.expect(with.recall_kept > before.recall_kept and with.rounds < before.rounds); } + +/// A proposer that gives `out` whatever the context, its match always 0; `priced`: Proposer.VTable.priced. +const Stub = struct { + out: []const u32, + priced: bool, + + fn propose(ptr: *anyopaque, _: []const u32, max_draft: i64) anyerror![]const u32 { + const x: *Stub = @ptrCast(@alignCast(ptr)); + return x.out[0..@min(x.out.len, @as(usize, @intCast(@max(0, max_draft))))]; + } + fn lastMatch(_: *anyopaque) i64 { + return 0; + } + fn proposer(x: *Stub) Proposer { + const plain: Proposer.VTable = .{ .propose = propose, .last_match = lastMatch }; + const priced: Proposer.VTable = .{ .propose = propose, .last_match = lastMatch, .priced = true }; + return .{ .ptr = x, .vtable = if (x.priced) &priced else &plain }; + } +}; + +/// A shared round's copy for a stream of `p` drafting with `stub`, a record continuing `p` (tokens gpa-owned). +fn copyWith(stub: *Stub, shared: *recall.Shared, p: []const u32) !win.Copy { + var cfg = try model(.{}); + defer cfg.deinit(gpa); + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer engine.deinit(); + var r: recall.Reader = .{ .gpa = gpa, .shared = shared }; + defer r.deinit(); + var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = 16, .eos = &.{96}, .proposer = stub.proposer(), .recall = &r }); + defer s.deinit(gpa); + engine.alone = false; + const c = try win.copyProposal(&engine, &s); + return .{ .tokens = try gpa.dupe(u32, c.tokens), .recalled = c.recalled }; +} + +test "a priced proposer's copy is taken as it gives it, and its none stands: no gate, no Recall in its place" { + const p = prompt(40, 5); + const plain = try run(&.{.{ .prompt = &p, .drafts = false }}, null, .{}); + defer plain.deinit(); + var shared = try newShared(); + defer shared.record.deinit(); + try remember(&shared, &p, plain.emitted[0]); + var unpriced: Stub = .{ .out = &.{}, .priced = false }; + const fallback = try copyWith(&unpriced, &shared, &p); + defer gpa.free(fallback.tokens); + try std.testing.expect(fallback.recalled and fallback.tokens.len >= 2); // an unpriced none: Recall's copy instead + var none: Stub = .{ .out = &.{}, .priced = true }; + const kept_none = try copyWith(&none, &shared, &p); + defer gpa.free(kept_none.tokens); + try std.testing.expect(!kept_none.recalled and kept_none.tokens.len == 0); + var one: Stub = .{ .out = &.{plain.emitted[0][0]}, .priced = true }; // one token, no match: under every gate + const kept_one = try copyWith(&one, &shared, &p); + defer gpa.free(kept_one.tokens); + try std.testing.expect(!kept_one.recalled); + try std.testing.expectEqualSlices(u32, one.out, kept_one.tokens); +} diff --git a/zig/src/core/lanes/trim.zig b/zig/src/core/lanes/trim.zig new file mode 100644 index 0000000000..426577030f --- /dev/null +++ b/zig/src/core/lanes/trim.zig @@ -0,0 +1,72 @@ +//! An optional row policy in place of windows.allocate: a family's own depth rule trims drafts. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Stream = @import("stream.zig").Stream; + +/// One planned window as the policy sees it. +pub const Window = struct { + stream: *Stream, + drafted: bool, // head drafts the policy may trim; every other window runs whole + rows: u32, // rows as planned, the pending row included + parents: ?[]const i32, // the drafts' parents for a tree (null: a chain) +}; + +/// What a window keeps: its first `count` drafts, or the listed drafts (ascending, each one's parent listed). +pub const Choice = struct { count: u32, nodes: ?[]const u32 = null }; + +/// A committed window: its rows, the accepted path (row 0 first) and the target's choice after the path's last row. +pub const Commit = struct { + stream: *Stream, + drafted: bool, + rows: u32, + parents: ?[]const i32, // each row's parent row (row 0: -1) for a tree; null for a chain + path: []const u32, + bonus: u32, +}; + +pub const Trim = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// Each window's drafts for this round; `alone`: the round's only window. Slices in `out` live in `arena`. + choose: *const fn (ptr: *anyopaque, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) anyerror!void, + /// Every window's commit, in plan order (a shared round's windows one after the other). + commit: *const fn (ptr: *anyopaque, c: Commit) anyerror!void, + }; + + pub fn choose(t: Trim, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) !void { + return t.vtable.choose(t.ptr, windows, alone, arena, out); + } + pub fn commit(t: Trim, c: Commit) !void { + return t.vtable.commit(t.ptr, c); + } +}; + +pub const Sub = struct { tokens: []u32, parents: []i32 }; + +/// The listed drafts of a tree re-indexed among themselves (parents -1: the pending row); arena-owned. +pub fn select(arena: Allocator, tokens: []const u32, parents: []const i32, nodes: []const u32) !Sub { + const at = try arena.alloc(i32, tokens.len); + @memset(at, -1); + const out_t = try arena.alloc(u32, nodes.len); + const out_p = try arena.alloc(i32, nodes.len); + for (nodes, out_t, out_p, 0..) |n, *t, *p, i| { + const q = parents[n]; + if (q >= 0 and at[@intCast(q)] < 0) return error.ParentNotKept; + at[n] = @intCast(i); + t.* = tokens[n]; + p.* = if (q < 0) -1 else at[@intCast(q)]; + } + return .{ .tokens = out_t, .parents = out_p }; +} + +test "select keeps listed drafts with their parents re-indexed" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + // chain 10 11 12, sibling 20 (on the pending row) with child 21 + const s = try select(arena.allocator(), &.{ 10, 11, 12, 20, 21 }, &.{ -1, 0, 1, -1, 3 }, &.{ 0, 3, 4 }); + try std.testing.expectEqualSlices(u32, &.{ 10, 20, 21 }, s.tokens); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 1 }, s.parents); + try std.testing.expectError(error.ParentNotKept, select(arena.allocator(), &.{ 10, 11 }, &.{ -1, 0 }, &.{1})); +} diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index 36200960d2..169dd4e50c 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -4,6 +4,7 @@ const accept = @import("accept.zig"); const alloc = @import("allocate.zig"); const be = @import("backend.zig"); const depth = @import("depth.zig"); +const trim = @import("trim.zig"); const fill = @import("fill.zig"); const shape = @import("shape.zig"); const ev = @import("events.zig"); @@ -172,6 +173,8 @@ pub fn copyProposal(e: *Engine, s: *Stream) !Copy { const width: u32 = if (e.alone) (s.copy_width orelse e.cfg.first_copy) else e.cfg.first_copy; const most = @min(@as(i64, width), s.draftRoom() - 1); if (s.proposer) |p| if (!e.alone or e.cfg.lone_copies) { + // a priced proposer's drafts are as it gives them, its none too: no width, length or match gate, no Recall + if (p.vtable.priced) return .{ .tokens = p.propose(s.context.items, @min(@as(i64, e.cfg.max_copy), s.draftRoom() - 1)) catch &.{} }; const copied = p.propose(s.context.items, most) catch &.{}; if (copied.len >= 2 and p.lastMatch() >= e.cfg.enter_match) return .{ .tokens = copied }; }; @@ -204,6 +207,7 @@ pub fn anySampled(plans: []const Plan) bool { /// Trim draft prefixes by landing probability to the shared width with the most expected tokens a ms. pub fn allocate(e: *Engine, plans: []Plan) !void { + if (e.trim) |t| return trimmed(e, plans, t); const a = e.arena.allocator(); const fixed = try a.alloc(u32, plans.len); const probs = try a.alloc([]const f64, plans.len); @@ -247,6 +251,37 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { try trail.event(e, &.{ f("ev", trail.str("alloc")), f("plans", .{ .list = rows }) }); } +/// A family's row policy decides each head window's drafts (lanes/trim.zig); other windows run whole. +fn trimmed(e: *Engine, plans: []Plan, t: trim.Trim) !void { + const a = e.arena.allocator(); + const asks = try a.alloc(trim.Window, plans.len); + for (asks, plans) |*w, p| w.* = .{ .stream = p.stream, .drafted = p.kind == .head, .rows = @intCast(1 + p.count()), .parents = p.parents }; + const out = try a.alloc(trim.Choice, plans.len); + for (out, plans) |*o, p| o.* = .{ .count = @intCast(p.count()) }; + try t.choose(asks, e.alone and plans.len == 1, a, out); + for (plans, out) |*p, c| { + if (p.kind != .head) continue; + if (c.nodes) |nodes| { + if (p.parents == null or p.held > 0) return error.TrimNodesNeedTree; + const sub = try trim.select(a, p.tokens, p.parents.?, nodes); + p.tokens = sub.tokens; + p.parents = if (accept.isChain(try accept.rowParents(a, 1 + sub.tokens.len, sub.parents))) null else sub.parents; + } else if (c.count == 0) { + p.kind = .none; + p.held = 0; + p.tokens = &.{}; + p.parents = null; + } else if (c.count < p.count()) { + if (p.parents) |parents| { + const tree = try accept.sanitizeTree(a, p.tokens, parents, c.count); + p.tokens = tree.tokens; + p.parents = tree.parents; + } else if (p.held > 0) p.held = c.count else p.tokens = p.tokens[0..c.count]; + } + if (p.count() == 0) p.kind = .none; + } +} + /// The head's chance for each held draft, else chain chances from the stream's per-depth acceptance. pub fn nodeChances(e: *Engine, s: *Stream, n: usize) ![]const f64 { const a = e.arena.allocator();