diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index cb2159b03a..a20d219c55 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -35,7 +35,12 @@ pub const Model = struct { max_streams: u32 = 0, // a shared forward's streams at most (0: 32) shared_costs: []const Cost = &.{}, // shared forwards' ms by total rows (timed at load) draft_probabilities: bool = false, // the head gives each drafted node's chance of landing + granted_depth: bool = false, // with draft_probabilities: every stream drafts its last grant plus a slack + draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own) draft_streams: bool = false, // the head drafts every stream of a shared round in one batch + shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig) + copy_first: bool = false, // a copy ahead takes the round: the head holds no drafts for it (a drafter's own pass) + hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // a drafter's own hold pass, timed at load 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) @@ -64,10 +69,12 @@ pub const Config = struct { batch_rows: u32, batch_streams: u32, node_probabilities: bool, + granted_depth: bool = false, // Model.granted_depth fill_lanes: u32 = 0, // rows a window may grow to with suffix-match branches (0: off; lanes/fill.zig) fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases) 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, + shared_guard: bool = false, 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, @@ -146,7 +153,9 @@ pub const Config = struct { .batch_rows = batch_rows, .batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1, .node_probabilities = m.draft_probabilities or m.draft_confidence != null, + .granted_depth = m.granted_depth and m.draft_probabilities, .draft_streams = m.draft_streams, + .shared_guard = m.shared_guard, .head_trees = m.head_trees, .join_tail = m.join_tail, .piece_runs = m.piece_runs, @@ -157,7 +166,14 @@ pub const Config = struct { .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows), .lane_costs = try alloc.extendCosts(gpa, m.lane_costs, 64), .mtp_step_ms = m.mtp_step_ms, - .k = .{ .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate }, + .k = .{ + .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate, + .draft_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).draft_slack, + .grant_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).grant_slack, + .short_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).short_slack, + .copy_first = m.copy_first and !m.context_copies, + .hold_costs = m.hold_costs, + }, .draft_step_ms = m.draft_step_ms, }; } diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 9f6c711ac5..9e2d3a2f7a 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -11,8 +11,16 @@ pub const Constants = struct { plain_margin: f64 = 1.05, // a drafted depth must beat plain rounds' tokens a ms by this much plain_wait_most: u64 = 128, // plain rounds between probes at most draft_slack: i64 = 2, // drafts a stream offers past what its last shared round granted it + grant_slack: i64 = 1, // granted depths: drafts past a grant that verified every draft the stream held + short_slack: i64 = 0, // granted depths: drafts past a grant that left some of them unverified + copy_first: bool = false, // a copy ahead takes the next round whole: the head drafts none for it (else one) + hold_costs: []const HoldCost = &.{}, // a drafter's own hold pass by blocks and depth: granted holds must pay for it }; +pub const HoldCost = @import("hold_price.zig").HoldCost; +const holdMs = @import("hold_price.zig").holdMs; +const lone_guard = @import("lone_guard.zig"); + /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). pub const State = struct { p: []f64, @@ -20,6 +28,7 @@ pub const State = struct { plain: ?u64 = null, wait: ?u64 = null, ms: ?Table = null, // plain_guard: the stream's own round times + lone: lone_guard.Lone = .{}, // lone_guard: whether its lone holds draft pub fn deinit(st: *State, gpa: Allocator) void { gpa.free(st.p); @@ -34,8 +43,11 @@ pub const Who = struct { finished: bool, forced: bool, granted: ?u32 = null, + grant_short: bool = false, // the last grant left some held drafts unverified (the width did not bind) + gain: ?f64 = null, // the drafts its last grant expected to keep (null: none priced) stream: ?*anyopaque = null, sampled: bool = false, // keyed draws: host work a row that a greedy round does not have + shared: bool = false, // a shared round's budget: the lone guard stays out }; /// Whether a copied continuation is ahead of a stream (the round loop proposes; Python `_copy_proposal`). @@ -58,6 +70,10 @@ pub const Rule = struct { draft_step_ms: f64 = 0, // a shared round's draft past a stream's first (budgets price it) plain_guard: bool, node_probabilities: bool, + granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too + lone_guard: bool = false, // granted_depth: a lone stream holds drafts only while its drafted rounds measure faster + lone_k: lone_guard.Constants = .{}, + lone_plain: lone_guard.Plain = .{}, // plain rounds' tokens a ms, every lone stream's confidence_floor: ?f64 = null, batch_rows: u32, batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's) @@ -191,17 +207,123 @@ pub const Rule = struct { return true; } - /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead). + /// Drafts past a grant: none when it stopped short of the held drafts (the width did not bind), else grant_slack. + fn slack(r: *const Rule, who: Who) i64 { + return if (who.grant_short) r.k.short_slack else r.k.grant_slack; + } + + /// The last grant plus its slack (the deepest chain before any); after a grant of none, a probe of draft_slack. + fn grantDepth(r: *Rule, who: Who) !i64 { + const most = @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1)); + if (r.lone_guard and !who.shared) { + const l = &(try r.state(who)).lone; + if (!l.drafts()) return 0; + if (l.fresh) { + l.fresh = false; + return @min(most, 1 + r.k.grant_slack); + } + } + const granted = who.granted orelse return most; + const st = try r.state(who); + if (granted == 0) return if (!r.plain_guard or r.probePlain(st)) @min(most, r.k.draft_slack) else 0; + st.plain = 0; + st.wait = r.k.depth_probe_every; + return @min(most, @as(i64, granted) + r.slack(who)); + } + + fn state(r: *Rule, who: Who) !*State { + _ = try r.rates(who); + return &who.state.*.?; + } + + /// A lone round's committed tokens and wall ms at the stream's context band, for the lone guard. + pub fn observeLone(r: *Rule, who: Who, band: usize, tokens: usize, ms: f64) !void { + if (!r.lone_guard) return; + const memo = r.lone_plain.memo(who.sampled, band); + (try r.state(who)).lone.observeBand(r.lone_k, r.lone_plain.at(who.sampled, band), memo, tokens, ms); + } + + /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead). pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 { - const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.node_probabilities) + const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth) + try r.grantDepth(who.*) + else if (r.node_probabilities) @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1)) else try r.depth(who.*); if (who.forced) return 0; - if (d != 0 and probe.copy(probe.ptr, who)) return 1; + if (d != 0 and probe.copy(probe.ptr, who)) return @intFromBool(!r.k.copy_first); return d; } + /// A shared hold pads every block to its deepest: depths cut so the drafting streams' blocks fit `rows` (one pass). + fn capPass(out: []i64, rows: u32) void { + var n: i64 = 0; + for (out) |d| n += @intFromBool(d > 0); + if (n == 0) return; + const most = @max(1, @divFloor(@as(i64, rows), n)); + for (out) |*d| d.* = @min(d.*, most); + } + + /// Granted holds kept, most expected drafts a row first, while they raise the next round's expected tokens a ms. + fn payHolds(r: *Rule, whos: []const Who, sampled: bool, out: []i64) !void { + const order = try r.gpa.alloc(usize, whos.len); + defer r.gpa.free(order); + var m: usize = 0; + var probes: usize = 0; // a probe after a grant of none keeps its hold: it measures the stream again + var deepest: i64 = 0; + for (whos, out, 0..) |w, d, i| { + if (d <= 0) continue; + if (w.granted == @as(?u32, 0)) { + probes += 1; + deepest = @max(deepest, d); + } else { + order[m] = i; + m += 1; + } + } + if (m == 0) return; + const Ctx = struct { + r: *Rule, + whos: []const Who, + fn less(c: @This(), x: usize, y: usize) bool { + const gx = c.r.gainOf(c.whos[x]) catch 0; + const gy = c.r.gainOf(c.whos[y]) catch 0; + return gx * rowsOf(c.whos[y]) > gy * rowsOf(c.whos[x]) or (gx * rowsOf(c.whos[y]) == gy * rowsOf(c.whos[x]) and x < y); + } + }; + std.mem.sort(usize, order[0..m], Ctx{ .r = r, .whos = whos }, Ctx.less); + const fixed = r.fixedMs(whos.len, sampled); + var rows: f64 = @floatFromInt(whos.len); + var tokens = rows; + const plain = r.shared_costs.get(@intCast(whos.len)) orelse return; + const probe_ms = if (probes > 0) holdMs(r.k.hold_costs, probes, deepest) else 0; + var best = tokens / (plain + fixed + probe_ms); + var keep: usize = 0; + for (order[0..m], 1..) |i, k| { + rows += rowsOf(whos[i]); + tokens += try r.gainOf(whos[i]); + deepest = @max(deepest, out[i]); + const verify = r.shared_costs.get(@intFromFloat(@ceil(rows))) orelse break; + const now = tokens / (verify + fixed + holdMs(r.k.hold_costs, k + probes, deepest)); + if (now > best) { + best = now; + keep = k; + } + } + for (order[keep..m]) |i| out[i] = 0; + } + + /// The drafts a granted stream's next round is expected to keep: its last grant's, or its first draft's rate. + fn gainOf(r: *Rule, w: Who) !f64 { + if (w.gain) |g| if (w.granted != null) return g; + return (try r.rates(w))[0]; + } + + fn rowsOf(w: Who) f64 { + return @floatFromInt(@max(1, w.granted orelse 1)); + } + /// Head drafts for the streams of a shared round, sized by allocation over their chain chances. pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void { if (r.confidence_floor != null) { @@ -210,10 +332,13 @@ pub const Rule = struct { } if (r.node_probabilities) { for (whos, out) |*w, *o| { + w.shared = true; const inner = try r.headDepth(w, null, probe); const granted: i64 = if (w.granted) |g| g else r.most_drafts; - o.* = try r.headDepth(w, @min(inner, granted + r.k.draft_slack), probe); + o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe); } + if (r.granted_depth) capPass(out, r.batch_rows); + if (r.granted_depth and r.k.hold_costs.len > 0) try r.payHolds(whos, sampled, out); return; } const most = try r.gpa.alloc(i64, whos.len); @@ -380,3 +505,81 @@ test "a shared round's drafts come out the same whichever depths its overhead wa } try std.testing.expectEqualSlices(i64, &got[0], &got[1]); } + +test "granted depths in a shared round fit one hold pass: every block is as deep as the deepest" { + const empty: Table = .{}; + var rule: Rule = .{ .gpa = std.testing.allocator, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + var states: [8]?State = @splat(null); + defer for (&states) |*st| if (st.*) |*x| x.deinit(std.testing.allocator); + var whos: [8]Who = undefined; + for (&whos, &states, 0..) |*w, *st, i| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = @intCast(i) }; + const No = struct { + fn copy(_: *anyopaque, _: *const Who) bool { + return false; + } + }; + var unused: u8 = 0; + var out: [8]i64 = undefined; + try rule.budgets(&whos, .{ .ptr = &unused, .copy = No.copy }, false, &out); + // a grant of none probes two; grants 1..7 that took every held draft, plus one; cut to 32 rows over eight blocks + try std.testing.expectEqualSlices(i64, &.{ 2, 2, 3, 4, 4, 4, 4, 4 }, &out); + try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]); + try std.testing.expectEqualSlices(i64, &.{ 2, 2 }, out[0..2]); // two blocks fit whole + for (whos[0..2]) |*w| w.grant_short = true; + try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]); + try std.testing.expectEqualSlices(i64, &.{ 2, 1 }, out[0..2]); // a grant short of the held drafts: no slack +} + +test "a copy ahead: one head draft beside it, or none when the head's drafts cost a pass of their own" { + const empty: Table = .{}; + const Yes = struct { + fn copy(_: *anyopaque, _: *const Who) bool { + return true; + } + }; + var unused: u8 = 0; + const probe: Probe = .{ .ptr = &unused, .copy = Yes.copy }; + var st: ?State = null; + defer if (st) |*x| x.deinit(std.testing.allocator); + const who: Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 3 }; + for ([_]bool{ false, true }) |first| { + var rule: Rule = .{ .gpa = std.testing.allocator, .k = .{ .copy_first = first }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + try std.testing.expectEqual(@as(i64, @intFromBool(!first)), try rule.headDepth(&who, null, probe)); + var out: [1]i64 = undefined; + var whos = [_]Who{who}; + try rule.budgets(&whos, probe, false, &out); + try std.testing.expectEqual(@as(i64, @intFromBool(!first)), out[0]); + } +} + +test "a granted hold stays only while its stream's expected drafts pay for the verify rows and the hold pass" { + const gpa = std.testing.allocator; + var shared: Table = .{}; + defer shared.deinit(gpa); + for (1..33) |w| try shared.put(gpa, @intCast(w), 10.0 + @as(f64, @floatFromInt(w))); + const empty: Table = .{}; + const cheap = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 8, .depth = 1, .ms = 4.5 } }; + var states: [4]?State = @splat(null); + defer for (&states) |*st| if (st.*) |*x| x.deinit(gpa); + const granted = [_]?u32{ 2, 2, 0, 1 }; + const gains = [_]?f64{ 1.8, 0.2, null, 0.9 }; + var whos: [4]Who = undefined; + for (&whos, &states, granted, gains) |*w, *st, g, x| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = g, .gain = x }; + var rule: Rule = .{ .gpa = gpa, .k = .{ .hold_costs = &cheap }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &shared, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + // 8 ms overhead: plain 4 / 22; stream 0 then 3 raise the rate (5.8 / 25, 6.7 / 26.5), stream 1 would lower it + var out = [_]i64{ 4, 4, 0, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 4, 0, 0, 3 }, &out); + const dear = [_]HoldCost{.{ .streams = 1, .depth = 1, .ms = 20.0 }}; + rule.k.hold_costs = &dear; + out = .{ 4, 4, 0, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 0, 0, 0, 0 }, &out); // no gain pays a 20 ms pass + whos[2].granted = 0; + out = .{ 4, 4, 2, 3 }; + try rule.payHolds(&whos, false, &out); + try std.testing.expectEqualSlices(i64, &.{ 0, 0, 2, 0 }, &out); // a probe after a grant of none keeps its hold +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 999aaeab6a..211f30112c 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -103,6 +103,16 @@ pub const Drafted = struct { m.draft_streams = d.batched; m.draft_probabilities = false; m.head_trees = d.trees; + m.shared_guard = d.shared_guard; + m.copy_first = d.copy_first; + m.hold_costs = d.hold_costs; + if (d.chances and x.drafter.vtable.chances != null) { + // a chain's drafts priced by the drafter's own chances; the next hold as deep as the grant and its slack + m.draft_probabilities = d.confidence_floor == null; + m.draft_confidence = d.confidence_floor; + m.granted_depth = d.confidence_floor == null; + m.draft_slack = d.slack; + } if (d.tree_nodes > 0) { // the drafter grows its own trees, best first, and prices their nodes m.drafts = @min(d.tree_nodes, if (target.exact_width > 1) target.exact_width - 1 else 0); diff --git a/zig/src/core/lanes/drafted_chances_test.zig b/zig/src/core/lanes/drafted_chances_test.zig new file mode 100644 index 0000000000..d3f248e8da --- /dev/null +++ b/zig/src/core/lanes/drafted_chances_test.zig @@ -0,0 +1,266 @@ +//! A drafter that prices its chains (Facts.chances): verify rows go by its chances and replies stay one-token decode. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Model = @import("config.zig").Model; +const Cost = @import("config.zig").Cost; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const depth = @import("depth.zig"); +const Table = @import("table.zig").Table; +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; +const most_streams = 8; + +/// drafted_test's drafter with a chance for each held draft: `sure` where its guess is right, `unsure` where wrong. +pub const Priced = struct { + head: dt.FakeDrafter, + sure: f64 = 0.7, + unsure: f64 = 0.7, + vtable: dr.Drafter.VTable = undefined, + inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined, + at: std.AutoHashMapUnmanaged(*sm.Stream, u64) = .empty, // each stream's last hold's position + asked: u32 = 0, + held: u64 = 0, // drafts its holds asked for + + fn of(ptr: *anyopaque) *Priced { + return @fieldParentPtr("head", @as(*dt.FakeDrafter, @ptrCast(@alignCast(ptr)))); + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + const p = of(ptr); + for (items) |it| { + try p.at.put(gpa, it.stream, it.position); + p.held += it.depth; + } + return p.inner_hold(ptr, items); + } + + /// The running product along the chain, as the DSpark confidence head's survivals are. + fn chances(ptr: *anyopaque, s: *sm.Stream, out: []f64) anyerror!bool { + const p = of(ptr); + const lane = p.head.lanes.getPtr(s) orelse return false; + if (out.len > lane.held.items.len) return error.TooFewDrafts; + const position = p.at.get(s) orelse return error.NoHold; + p.asked += 1; + var reach: f64 = 1; + for (out, 0..) |*o, j| { + const at = position + j; + reach *= if ((at *% 2654435761 + j) % 5 == 0) p.unsure else p.sure; // drafted_test's wrong guesses + o.* = reach; + } + return true; + } + + pub fn drafter(p: *Priced) dr.Drafter { + const base_ = p.head.drafter(); + p.vtable = base_.vtable.*; + p.inner_hold = base_.vtable.hold; + p.vtable.hold = hold; + p.vtable.chances = chances; + return .{ .ptr = &p.head, .vtable = &p.vtable }; + } + + pub fn deinit(p: *Priced) void { + p.at.deinit(gpa); + p.head.deinit(); + } +}; + +/// Every shared verify's window rows: whether one round gave its streams unequal rows (a ragged verify). +const Seen = struct { + var inner: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var ragged: u32 = 0; + var rows: u64 = 0; + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + for (windows) |w| { + rows += w.rows(); + if (w.rows() != windows[0].rows()) { + ragged += 1; + break; + } + } + return inner(ptr, windows, out); + } +}; + +var costs: [16]Cost = undefined; + +pub fn base() Model { + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + return .{ .exact_width = 16, .gpu_tokens = true, .window_costs = &costs, .hidden_rows = true, .batch_rows = 32, .max_streams = most_streams }; +} + +const Run = struct { tokens: [most_streams][64]u32 = undefined, lens: [most_streams]usize = @splat(0), drafted: u64 = 0, holds: u32 = 0, asked: u32 = 0, ragged: u32 = 0, rows: u64 = 0, held: u64 = 0 }; + +/// The drafter's Facts.slack for the next runs (0: the depth rule's own; 9: every hold the whole block). +var slack: u32 = 0; + +/// The drafter's chances: equal everywhere, or knowing its wrong guesses (a confidence head that ranks well). +const Head = enum { flat, knows }; + +/// Each prompt decoded together (one stream: alone); `floor` null: the chances price rows, else a fixed window. +fn run(prompts: []const []const u32, drafts: bool, floor: ?f64, head: Head) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .confidence_floor = floor, .slack = slack } } }; + if (head == .knows) { + p.sure = 0.95; + p.unsure = 0.05; + } + defer p.deinit(); + var strict: dt.Strict = .{ .inner = target.backend() }; + defer strict.deinit(); + var seen = strict.backend(); + var table = seen.vtable.*; + Seen.inner = table.verify; + Seen.ragged = 0; + Seen.rows = 0; + table.verify = Seen.verify; + seen.vtable = &table; + var wrapped = try Drafted.init(gpa, seen, p.drafter()); + defer wrapped.deinit(); + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(base()) else base(), 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock()); + defer engine.deinit(); + var streams: [most_streams]sm.Stream = undefined; + const n = prompts.len; + for (prompts, streams[0..n]) |prompt, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = prompt, .max_new = 48, .eos = &.{96}, .drafts = drafts }); + defer for (streams[0..n]) |*s| s.deinit(gpa); + for (streams[0..n]) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = p.head.holds, .asked = p.asked, .ragged = Seen.ragged, .rows = Seen.rows, .held = p.held }; + for (streams[0..n], 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + out.drafted += s.drafted; + } + return out; +} + +fn expectSame(a: Run, b: Run, n: usize) !void { + for (0..n) |i| try std.testing.expectEqualSlices(u32, a.tokens[i][0..a.lens[i]], b.tokens[i][0..b.lens[i]]); +} + +/// `n` streams of unequal prompts, each a prefix of drafted_test's long prompt. +fn unequal(buf: *[most_streams][]const u32, n: usize) []const []const u32 { + for (buf[0..n], 0..) |*p, i| p.* = dt.long[0 .. 11 + i]; + return buf[0..n]; +} + +test "a drafter's chances price its chain's rows, and drafted replies equal one-token decode alone and shared" { + const sets = [_][]const []const u32{ &.{&dt.p1}, &.{ &dt.p1, &dt.long, &dt.p1 } }; + for (sets) |ps| { + const plain = try run(ps, false, null, .flat); + const priced = try run(ps, true, null, .flat); + try expectSame(plain, priced, ps.len); + try std.testing.expect(priced.asked > 0 and priced.drafted > 0); // the window allocator read the chances + const fixed = try run(ps, true, 0, .flat); + try expectSame(plain, fixed, ps.len); + try std.testing.expect(fixed.drafted > priced.drafted); // a fixed window verifies every draft the rows fit + } +} + +test "at 1, 4 and 8 streams the scheduled and the fixed verify equal one-token decode; shared rounds are ragged" { + var buf: [most_streams][]const u32 = undefined; + for ([_]usize{ 1, 4, 8 }) |n| { + const ps = unequal(&buf, n); + const plain = try run(ps, false, null, .flat); + const scheduled = try run(ps, true, null, .knows); + try expectSame(plain, scheduled, n); + const fixed = try run(ps, true, 0, .knows); + try expectSame(plain, fixed, n); + try std.testing.expect(scheduled.asked > 0 and scheduled.drafted > 0); + // a confident head's survivals cut each stream's chain where its guess goes wrong: fewer rows verified + try std.testing.expect(scheduled.rows < fixed.rows); + if (n > 1) try std.testing.expect(scheduled.ragged > 0); // streams' verify lengths differ within a round + } +} + +test "the drafter's chances reach the lane facts: rows by chances and holds by grants, or a fixed window" { + var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .chances = true } } }; + defer p.deinit(); + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var wrapped: Drafted = .{ .gpa = gpa, .target = target.backend(), .drafter = p.drafter() }; + defer wrapped.deinit(); + const m = wrapped.facts(base()); + try std.testing.expect(m.draft_probabilities and m.granted_depth and m.draft_confidence == null); + var cfg = try Config.init(gpa, m, 16, 15); + defer cfg.deinit(gpa); + try std.testing.expectEqual(@as(i64, 2), cfg.k.draft_slack); // the depth rule's own + try std.testing.expect(cfg.k.grant_slack == 1 and cfg.k.short_slack == 0); // a granted width's own + p.head.facts_.slack = 1; + var one = try Config.init(gpa, wrapped.facts(base()), 16, 15); + defer one.deinit(gpa); + try std.testing.expectEqual(@as(i64, 1), one.k.draft_slack); // a drafter's own slack past its grant + try std.testing.expect(one.k.grant_slack == 1 and one.k.short_slack == 1); // short or not + p.head.facts_.confidence_floor = 0; + const f = wrapped.facts(base()); + try std.testing.expect(!f.draft_probabilities and !f.granted_depth and f.draft_confidence.? == 0); + p.head.facts_.chances = false; + const none = wrapped.facts(base()); + try std.testing.expect(!none.draft_probabilities and none.draft_confidence == null); +} + +test "a granted depth: the deepest chain first, then the grant plus its slack, and after a grant of none a probe" { + const empty: Table = .{}; + var rule: depth.Rule = .{ .gpa = gpa, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 }; + defer rule.deinit(); + var st: ?depth.State = null; + defer if (st) |*s| s.deinit(gpa); + const No = struct { + fn copy(_: *anyopaque, _: *const depth.Who) bool { + return false; + } + }; + var unused: u8 = 0; + const probe: depth.Probe = .{ .ptr = &unused, .copy = No.copy }; + var who: depth.Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false }; + try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe)); + who.granted = 3; + try std.testing.expectEqual(@as(i64, 4), try rule.headDepth(&who, null, probe)); // every held draft granted: one more + who.grant_short = true; + try std.testing.expectEqual(@as(i64, 3), try rule.headDepth(&who, null, probe)); // drafts left unverified: the grant + who.grant_short = false; + who.granted = 8; + try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe)); + rule.k.grant_slack = 2; // a drafter's slack sets both + rule.k.short_slack = 2; + who.grant_short = true; + who.granted = 3; + try std.testing.expectEqual(@as(i64, 5), try rule.headDepth(&who, null, probe)); + who.granted = 0; + var drafted: usize = 0; + for (0..8) |_| drafted += @intFromBool(try rule.headDepth(&who, null, probe) > 0); + try std.testing.expectEqual(@as(usize, 1), drafted); // one probe of draft_slack drafts in the first eight rounds +} + +test "holds as wide as the grants: replies equal one-token decode at 1, 4 and 8 streams, fewer drafts than whole blocks" { + var buf: [most_streams][]const u32 = undefined; + defer slack = 0; + for ([_]usize{ 1, 4, 8 }) |n| { + const ps = unequal(&buf, n); + const plain = try run(ps, false, null, .flat); + for ([_]Head{ .flat, .knows }) |head| { + slack = 0; + const narrow = try run(ps, true, null, head); + slack = 2; + const two = try run(ps, true, null, head); + slack = 9; + const whole = try run(ps, true, null, head); + try expectSame(plain, narrow, n); + try expectSame(plain, two, n); + try expectSame(plain, whole, n); + try std.testing.expect(narrow.drafted > 0); + try std.testing.expect(narrow.held < two.held and two.held < whole.held); + } + } +} diff --git a/zig/src/core/lanes/drafted_copy_test.zig b/zig/src/core/lanes/drafted_copy_test.zig new file mode 100644 index 0000000000..329731f69c --- /dev/null +++ b/zig/src/core/lanes/drafted_copy_test.zig @@ -0,0 +1,135 @@ +//! A drafter whose hold costs a pass of its own: no hold for a stream whose next round is a copy. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const Proposer = @import("proposer.zig").Proposer; +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; + +/// Copies the fake target's own continuation (it lands whole) whenever the context's length is a multiple of three. +const Copier = struct { + out: [3]u32 = undefined, + match: i64 = 0, + + fn propose(ptr: *anyopaque, context: []const u32, most: i64) anyerror![]const u32 { + const c: *Copier = @ptrCast(@alignCast(ptr)); + c.match = 0; + if (context.len % 3 != 0 or most < 2) return &.{}; + var history: [128]u32 = undefined; + @memcpy(history[0..context.len], context); + const n: usize = @intCast(@min(3, most)); + for (c.out[0..n], 0..) |*t, k| { + t.* = fake.next(history[0 .. context.len + k], null, 0); + history[context.len + k] = t.*; + } + c.match = 100; + return c.out[0..n]; + } + + fn lastMatch(ptr: *anyopaque) i64 { + return @as(*Copier, @ptrCast(@alignCast(ptr))).match; + } + + const vtable: Proposer.VTable = .{ .propose = propose, .last_match = lastMatch }; +}; + +/// Each stream's last hold position, and the copy windows verified at a position a hold was asked for. +const Watch = struct { + var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined; + var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var at: [4]?u64 = @splat(null); + var streams: [4]?*sm.Stream = @splat(null); + var wasted: u32 = 0; + var copies: u32 = 0; + var blocks: u32 = 0; // the streams' blocks held, over every hold + + fn slot(s: *sm.Stream) usize { + for (&streams, 0..) |*x, i| if (x.* == null or x.* == s) { + x.* = s; + return i; + }; + unreachable; + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + blocks += @intCast(items.len); + // a joining stream's first hold comes before its first token is read: no copy can be foreseen for it + for (items) |it| at[slot(it.stream)] = if (it.stream.emitted().len > 0) it.position else null; + return inner_hold(ptr, items); + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + for (windows) |w| if (w.held == 0 and w.tokens.len > 0) { + copies += 1; + if (at[slot(w.stream)] == w.positions[0]) wasted += 1; + }; + return inner_verify(ptr, windows, out); + } +}; + +const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0, wasted: u32 = 0, copies: u32 = 0 }; + +fn run(n: usize, drafts: bool, copy_first: bool) !Run { + Watch.at = @splat(null); + Watch.streams = @splat(null); + Watch.wasted = 0; + Watch.copies = 0; + Watch.blocks = 0; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .copy_first = copy_first } }; + defer head.deinit(); + var hooked = head.drafter(); + var table = hooked.vtable.*; + Watch.inner_hold = table.hold; + table.hold = Watch.hold; + hooked.vtable = &table; + var wrapped = try Drafted.init(gpa, target.backend(), hooked); + defer wrapped.deinit(); + // the round loop's windows, before the wrapper turns held drafts into host tokens + var seen = if (drafts) wrapped.backend() else target.backend(); + var verified = seen.vtable.*; + Watch.inner_verify = verified.verify; + verified.verify = Watch.verify; + seen.vtable = &verified; + const m: @import("config.zig").Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, seen, clock.clock()); + defer engine.deinit(); + var copiers: [3]Copier = @splat(.{}); + var streams: [3]sm.Stream = undefined; + for (streams[0..n], copiers[0..n], 0..) |*s, *c, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 40, .eos = &.{96}, .drafts = drafts, .proposer = .{ .ptr = c, .vtable = &Copier.vtable } }); + defer for (streams[0..n]) |*s| s.deinit(gpa); + for (streams[0..n]) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = Watch.blocks, .wasted = Watch.wasted, .copies = Watch.copies }; + for (streams[0..n], 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "a copy ahead takes the round: with copy_first no hold is asked for it, and replies equal one-token decode" { + for ([_]usize{ 1, 3 }) |n| { + const plain = try run(n, false, false); + const before = try run(n, true, false); + const after = try run(n, true, true); + for (0..n) |i| { + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], before.tokens[i][0..before.lens[i]]); + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], after.tokens[i][0..after.lens[i]]); + } + try std.testing.expect(before.copies > 0 and after.copies > 0); + try std.testing.expect(before.wasted > 0); // a hold of one draft beside each copy, dropped by the copy round + try std.testing.expectEqual(@as(u32, 0), after.wasted); + try std.testing.expect(after.holds < before.holds); + } +} diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig new file mode 100644 index 0000000000..3715bec357 --- /dev/null +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -0,0 +1,155 @@ +//! A drafter that asks for the shared-round guard (Facts.shared_guard): which shared rounds draft, replies unchanged. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Model = @import("config.zig").Model; +const Engine = @import("engine.zig").Engine; +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); + +const gpa = std.testing.allocator; + +test "a stream joining shared rounds that the guard runs plain holds no first drafts; a lone one does" { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = true } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [2]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} }); + defer for (&streams) |*s| s.deinit(gpa); + try engine.addStream(&streams[0]); + try std.testing.expectEqual(@as(u32, 1), head.holds); // alone: its first drafts are held + for (0..3) |_| try engine.step(); + const before = head.holds; + try std.testing.expect(!try engine.guard.drafts(2)); // the guard measures plain shared rounds first + try engine.addStream(&streams[1]); + try std.testing.expectEqual(before, head.holds); // joining: none held + while (engine.activeCount() > 0) try engine.step(); +} + +test "without the guard's flag a stream joining shared rounds holds its first drafts as before" { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + try std.testing.expect(!cfg.shared_guard); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [2]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} }); + defer for (&streams) |*s| s.deinit(gpa); + try engine.addStream(&streams[0]); + for (0..3) |_| try engine.step(); + const before = head.holds; + try engine.addStream(&streams[1]); + try std.testing.expectEqual(before + 1, head.holds); // joining: its first drafts held + while (engine.activeCount() > 0) try engine.step(); +} + +const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0 }; + +/// Three streams decoded together: drafts off, or the fake drafter with or without the guard. +fn run(drafts: bool, guard: bool) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15); + defer cfg.deinit(gpa); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock()); + defer engine.deinit(); + var streams: [3]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = drafts }); + defer for (&streams) |*s| s.deinit(gpa); + for (&streams) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = head.holds }; + for (&streams, 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "the guard changes only which shared rounds draft: replies equal one-token decode with it on and off" { + const plain = try run(false, false); + const off = try run(true, false); + const on = try run(true, true); + for (0..3) |i| { + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], off.tokens[i][0..off.lens[i]]); + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], on.tokens[i][0..on.lens[i]]); + } + try std.testing.expect(on.holds < off.holds); // its plain stretch holds no drafts +} + +/// A clock that reads the same ms for every span: two of them with different ms stand for a fast and a slow machine. +const FlatClock = struct { + ms: f64, + + fn clock(c: *FlatClock) @import("backend.zig").Clock { + return .{ .ptr = c, .vtable = &.{ .start = start, .elapsed_ms = elapsed } }; + } + + fn start(_: *anyopaque) void {} + + fn elapsed(ptr: *anyopaque, _: @import("backend.zig").Mark) f64 { + const c: *FlatClock = @ptrCast(@alignCast(ptr)); + return c.ms; + } +}; + +/// Three streams, the guard asked for or not, timed or not, on a clock reading `ms` for every span. +fn runClock(timed: bool, guard: bool, ms: f64) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } }; + defer head.deinit(); + var wrapped = try Drafted.init(gpa, target.backend(), head.drafter()); + defer wrapped.deinit(); + const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15); + defer cfg.deinit(gpa); + cfg.timed = timed; + var clock: FlatClock = .{ .ms = ms }; + var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock()); + defer engine.deinit(); + var streams: [3]sm.Stream = undefined; + for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = true }); + defer for (&streams) |*s| s.deinit(gpa); + for (&streams) |*s| try engine.addStream(s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .holds = head.holds }; + for (&streams, 0..) |*s, i| { + out.lens[i] = s.emitted().len; + @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); + } + return out; +} + +test "with timed off the shared guard keeps no plain stretch: the same schedule on a fast and a slow clock" { + const fast = try runClock(false, true, 0.5); + const slow = try runClock(false, true, 400.0); + try std.testing.expectEqual(fast.holds, slow.holds); + for (0..3) |i| try std.testing.expectEqualSlices(u32, fast.tokens[i][0..fast.lens[i]], slow.tokens[i][0..slow.lens[i]]); + const unasked = try runClock(false, false, 0.5); + try std.testing.expectEqual(unasked.holds, fast.holds); // as if the guard were not asked for +} diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig index 80577e0edf..3ba9feb671 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -18,6 +18,12 @@ pub const Facts = struct { window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs) tree_nodes: u32 = 0, // hold() without lanes grows its own tree of up to `depth` nodes, at most this many (tree, chances) + shared_guard: bool = false, // shared rounds draft only while they measure faster than plain ones (Model.shared_guard) + chances: bool = false, // chances() gives each held chain draft's chance: verify rows go by them (Model.granted_depth) + confidence_floor: ?f64 = null, // with chances: every round drafts the deepest chain and verifies drafts this likely + slack: u32 = 0, // with chances: drafts a hold offers past the stream's last grant (0: the depth rule's own) + copy_first: bool = false, // a hold costs a pass of its own: none is asked for a stream with a copy ahead + hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // that pass's ms by blocks and depth (timed at load) }; /// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a92b27ae7f..5fbb009101 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -3,6 +3,7 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const Config = @import("config.zig").Config; const depth = @import("depth.zig"); +const shared_guard = @import("shared_guard.zig"); const accept = @import("accept.zig"); const be = @import("backend.zig"); const sm = @import("stream.zig"); @@ -16,6 +17,7 @@ const lone_costs = @import("lone_costs.zig"); const open = @import("open.zig"); const handover = @import("handover.zig"); const trim = @import("trim.zig"); +const pipelined = @import("pipelined.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -24,7 +26,7 @@ const f = ev.f; const str = trail.str; const int = trail.int; const Outcome = struct { got: []const u32, path: []const u32, cut: ?usize, follow: []const u32 }; -const Result = struct { rows: u32, keep: u32, got: []const u32 }; +pub 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"); @@ -42,11 +44,13 @@ pub const Engine = struct { 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, + guard: shared_guard.Guard, live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`) drafted: u64 = 0, accepted: u64 = 0, shared_rounds: u64 = 0, alone: bool = true, // this step has one live stream: its copies may take their ramp + plain_round: bool = false, // a shared round the guard keeps plain: no stream copies or drafts, as with drafts off lone_kept: [2][lone_costs.bands]?lone_costs.Kept = @splat(@splat(null)), // the newest lone calibrations, greedy then sampled, a band each steps: u64 = 0, arena: std.heap.ArenaAllocator, // a round's temporaries @@ -70,10 +74,13 @@ pub const Engine = struct { .draft_step_ms = cfg.draft_step_ms, .plain_guard = cfg.plain_guard, .node_probabilities = cfg.node_probabilities, + .granted_depth = cfg.granted_depth, + .lone_guard = cfg.granted_depth and cfg.shared_guard and cfg.timed, .confidence_floor = cfg.draft_confidence, .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, }, + .guard = .{ .gpa = gpa }, }; } @@ -81,6 +88,7 @@ pub const Engine = struct { e.live.deinit(e.gpa); e.unread.deinit(e.gpa); e.rule.deinit(); + e.guard.deinit(); for (&e.lone_kept) |*mode| for (mode) |*k| if (k.*) |*x| x.deinit(e.gpa); e.arena.deinit(); } @@ -124,7 +132,7 @@ pub const Engine = struct { if (live.items.len > 1 and e.cfg.family_streams) { // a shared round is synchronous: a stream that ran a step ahead lands its queued token first for (live.items) |s| { - if (s.inflight != null) try e.landInflight(s); + if (s.inflight != null) try pipelined.land(e, s); } var kept: usize = 0; for (live.items) |s| { @@ -147,7 +155,7 @@ pub const Engine = struct { for (live.items) |s| { // logprobs and grammars take verify windows: rows' logits are read, or bits made, before a forward const windowed = (e.cfg.family_mtp and s.drafts) or !e.cfg.pipelined or s.logprobs != null or s.grammar != null; - const r = if (windowed) try e.familyRound(s, null) else try e.pipelinedRound(s); + const r = if (windowed) try e.familyRound(s, null) else try pipelined.round(e, s); s.min_rows = if (s.min_rows == 0) r.rows else @min(s.min_rows, r.rows); } } @@ -195,7 +203,7 @@ pub const Engine = struct { } /// One stream's round: verify pending and drafted tokens together, keep drafts to the first mismatch. - fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result { + pub fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result { e.clock.start(); const priced = s.priced; s.priced = null; @@ -231,6 +239,7 @@ pub const Engine = struct { lone_costs.observe(s, plan.cost_probe_rows, rows, ms); } } + if (e.rule.lone_guard and e.alone and s.drafts and (plan.kind == .head or plan.kind == .none)) try e.rule.observeLone(win.who(s), lone_costs.band(s.context.items.len), o.got.len, e.clock.elapsedMs(.cost)); 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| { @@ -255,6 +264,8 @@ pub const Engine = struct { e.clock.start(); const a = e.arena.allocator(); const plans = try a.alloc(Plan, entries.len); + e.plain_round = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(entries.len); + defer e.plain_round = false; for (entries, plans) |s, *p| p.* = try win.plan(e, s, null); try win.allocate(e, plans); for (plans) |*p| try win.constrain(e, p); @@ -284,7 +295,9 @@ pub const Engine = struct { for (whos, heads.items) |*w, i| w.* = win.who(plans[i].stream); const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); - try e.rule.budgets(whos, win.probe(e), sampled, budgets); + // the guard's plain stretches: no stream drafts, so the round costs what it would without a head + const guarded = e.cfg.shared_guard and e.cfg.timed and !try e.guard.draftsFresh(plans.len, whos); + if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets); const ids = try a.alloc(ev.Value, heads.items.len); for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id); try trail.event(e, &.{ f("ev", str("budgets")), f("streams", .{ .list = ids }), f("depths", .{ .i64s = budgets }) }); @@ -313,6 +326,7 @@ pub const Engine = struct { 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.shared_guard and e.cfg.timed) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); if (e.backend.round_observer != null) e.backend.observed(paths); } @@ -492,64 +506,6 @@ pub const Engine = struct { } } - /// One token, the next forward queued first; a copied continuation ahead switches to verify windows. - fn pipelinedRound(e: *Engine, s: *Stream) !Result { - const before = s.mode; - // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one - const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe; - var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} }; - if (mode == .verify or mode == .exit) { - const verifying = mode == .verify and e.cfg.family_width >= 2; - const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{}; - if (proposal.tokens.len > 0) { - const copied = try e.arena.allocator().dupe(u32, proposal.tokens); - result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled }); - if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit; - } else { - // nothing to copy: back to steps queued ahead (this one lands next round) - s.mode = .pipe; - try e.queueNext(s, .{ .value = s.pending.? }); - } - try trail.pipe(e, s, before, result.got); - return result; - } - var current: Feed = .{ .handle = s.inflight.? }; - s.inflight = null; - if (try e.forcedNext(s)) |t| current = .{ .value = t }; - var token: u32 = undefined; - if (mode == .drain) { - token = try e.readFeed(current); // the last queued step: no new one - s.mode = .verify; - } else { - try e.queueNext(s, current); // the GPU starts the next step first - token = try e.readFeed(current); - } - s.rounds += 1; - const landed = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat - if (e.cfg.root_copies) s.root_copy_miss = false; - const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts; - if (piping and (try win.copyProposal(e, s)).tokens.len > 0) - s.mode = .drain; // a copy window is ahead: land the queued step - result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]); - try trail.pipe(e, s, before, result.got); - return result; - } - - /// Land the queued token without queuing another (a shared round is synchronous). - fn landInflight(e: *Engine, s: *Stream) !void { - var current: Feed = .{ .handle = s.inflight.? }; - s.inflight = null; - if (try e.forcedNext(s)) |t| current = .{ .value = t }; - const token = try e.readFeed(current); - s.rounds += 1; - const landed = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - s.mode = .verify; - try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) }); - } - /// The first token's draw: among the grammar's allowed tokens when the stream has one that constrains it. pub fn drawFirst(e: *Engine, s: *Stream, position: u64) !u64 { const g = s.grammar orelse return e.backend.first(s, position); diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig index 36167e30d6..6995de7dbe 100644 --- a/zig/src/core/lanes/engine_test.zig +++ b/zig/src/core/lanes/engine_test.zig @@ -463,6 +463,9 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea try std.testing.expect((try win.copyProposal(&engine, &s)).tokens.len >= 2); engine.alone = true; try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).tokens.len >= 2); + engine.alone = false; + engine.plain_round = true; // a shared round the guard keeps plain: no copies either, as with drafts off + try std.testing.expectEqual(@as(usize, 0), (try win.copyProposal(&engine, &s)).tokens.len); } } diff --git a/zig/src/core/lanes/hold_price.zig b/zig/src/core/lanes/hold_price.zig new file mode 100644 index 0000000000..e4c854410e --- /dev/null +++ b/zig/src/core/lanes/hold_price.zig @@ -0,0 +1,58 @@ +//! A drafter's hold pass priced from its timed passes: by the blocks it holds and their depth. +const std = @import("std"); + +/// A drafter's hold pass of `streams` blocks, each `depth` drafts deep, in ms (timed at load). +pub const HoldCost = struct { streams: u32, depth: u32, ms: f64 }; + +/// A hold pass's ms: linear in depth at each timed stream count, then in streams (past the widest, by its slope). +pub fn holdMs(costs: []const HoldCost, streams: usize, depth: i64) f64 { + var below: ?u32 = null; + var above: ?u32 = null; + for (costs) |c| { + if (c.streams <= streams and (below == null or c.streams > below.?)) below = c.streams; + if (c.streams >= streams and (above == null or c.streams < above.?)) above = c.streams; + } + const lo = below orelse above.?; + const hi = above orelse lo; + const at_lo = atDepth(costs, lo, depth); + if (hi == lo) { + if (above != null) return at_lo; + // past the widest timed count: the slope from the count below it + var under: ?u32 = null; + for (costs) |c| if (c.streams < lo and (under == null or c.streams > under.?)) { + under = c.streams; + }; + const u = under orelse return at_lo * @as(f64, @floatFromInt(streams)) / @as(f64, @floatFromInt(lo)); + const slope = (at_lo - atDepth(costs, u, depth)) / @as(f64, @floatFromInt(lo - u)); + return at_lo + slope * @as(f64, @floatFromInt(streams - lo)); + } + const t = @as(f64, @floatFromInt(streams - lo)) / @as(f64, @floatFromInt(hi - lo)); + return at_lo + t * (atDepth(costs, hi, depth) - at_lo); +} + +/// One stream count's hold ms at `depth`: linear between its two nearest timed depths (one alone: as timed). +fn atDepth(costs: []const HoldCost, streams: u32, depth: i64) f64 { + var a: ?HoldCost = null; + var b: ?HoldCost = null; + for (costs) |c| { + if (c.streams != streams) continue; + if (a == null or c.depth < a.?.depth) { + b = a; + a = c; + } else if (b == null or c.depth < b.?.depth) b = c; + } + const x = a.?; + const y = b orelse return x.ms; + const t = (@as(f64, @floatFromInt(depth)) - @as(f64, @floatFromInt(x.depth))) / @as(f64, @floatFromInt(y.depth - x.depth)); + return x.ms + t * (y.ms - x.ms); +} + +test "a hold's ms between and past the timed passes" { + const t = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 1, .depth = 4, .ms = 2.5 }, .{ .streams = 4, .depth = 1, .ms = 3.0 }, .{ .streams = 4, .depth = 4, .ms = 4.5 }, .{ .streams = 8, .depth = 1, .ms = 4.0 } }; + try std.testing.expectApproxEqAbs(@as(f64, 1.0), holdMs(&t, 1, 1), 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 2.0), holdMs(&t, 1, 3), 1e-12); // linear in depth + try std.testing.expectApproxEqAbs(@as(f64, 3.0), holdMs(&t, 1, 5), 1e-12); // and past the deepest + try std.testing.expectApproxEqAbs(@as(f64, 1.0 + 2.0 / 3.0), holdMs(&t, 2, 1), 1e-12); // linear in streams + try std.testing.expectApproxEqAbs(@as(f64, 4.0), holdMs(&t, 8, 3), 1e-12); // one depth timed: as timed + try std.testing.expectApproxEqAbs(@as(f64, 5.0), holdMs(&t, 12, 1), 1e-12); // past the widest: the slope below it +} diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index ff09b58fb2..a3b2587704 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -18,6 +18,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 hold_price = @import("hold_price.zig"); pub const mtp_policy = @import("mtp_policy.zig"); pub const trim = @import("trim.zig"); pub const config = @import("config.zig"); @@ -54,11 +55,17 @@ test { _ = @import("recorder_test.zig"); _ = @import("gpu_full_test.zig"); _ = @import("drafted_test.zig"); + _ = @import("drafted_chances_test.zig"); + _ = @import("drafted_copy_test.zig"); _ = @import("root_copy_test.zig"); _ = @import("drafted_grammar_test.zig"); _ = @import("windows_grammar_test.zig"); _ = @import("pricing_test.zig"); _ = @import("depth_probe_test.zig"); + _ = @import("shared_guard.zig"); + _ = @import("lone_guard.zig"); + _ = @import("lone_guard_test.zig"); + _ = @import("drafted_guard_test.zig"); _ = @import("proposal_external_test.zig"); _ = @import("lone_costs_test.zig"); _ = @import("lone_backoff_test.zig"); diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig new file mode 100644 index 0000000000..4657ffe829 --- /dev/null +++ b/zig/src/core/lanes/lone_guard.zig @@ -0,0 +1,242 @@ +//! Whether a lone stream holds drafts: its drafted rounds' tokens a ms against plain rounds', the loser re-measured. +const std = @import("std"); + +pub const Constants = struct { + drafted: u32 = 16, // measured drafted rounds a verdict takes, again and again while drafting wins + plain: u32 = 16, // measured plain rounds that take the plain rate when there is none yet + wait: u32 = 16, // plain rounds before drafting is measured again; doubles each time plain wins again + wait_most: u32 = 512, + margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much + probe_cost: f64 = 0.03, // after two losses a probe costs at most this share of plain time + confirm: u32 = 2, // losing verdicts in a row (a remembered one counts) before the wait follows probe_cost + blend: f64 = 0.5, // the newest stretch's weight in a rate +}; + +pub const Mode = enum { drafted, plain }; + +/// Plain rounds' tokens a ms by sampling mode and context band, kept from stream to stream (no draft changes them). +pub const Plain = struct { + rates: [2][bands]?f64 = @splat(@splat(null)), + verdicts: [2][bands]?Verdict = @splat(@splat(null)), // each band's last losing verdict, until drafting wins there + + pub const Verdict = struct { drafted: f64, wait: u32 }; + + pub fn memo(p: *Plain, sampled: bool, band: usize) *?Verdict { + return &p.verdicts[@intFromBool(sampled)][@min(band, bands - 1)]; + } + + pub const bands = 6; + + pub fn at(p: *Plain, sampled: bool, band: usize) *?f64 { + return &p.rates[@intFromBool(sampled)][@min(band, bands - 1)]; + } +}; + +/// One lone stream's guard: the mode its holds follow, the stretch measured and its drafted rate. +pub const Lone = struct { + mode: Mode = .drafted, + left: u32 = 0, + tokens: f64 = 0, + ms: f64 = 0, + skip: bool = true, // a mode's first round still verifies the other mode's hold (or none): not measured + started: bool = false, + pending: bool = false, // a drafted verdict waits for the plain rate this stretch takes + fresh: bool = false, // a drafted stretch's first hold: one draft and its slack, whatever the last grant + drafted: ?f64 = null, + drafted_ms: f64 = 0, // the last drafted stretch's wall ms + drafted_now: f64 = 0, // the last drafted stretch's own tokens a ms (unblended) + losses: u32 = 0, // losing verdicts in a row + wait: u32 = 0, + + /// Whether the next hold may draft. + pub fn drafts(l: *const Lone) bool { + return l.mode == .drafted; + } + + /// A lone round of held drafts or none (copies left out): its tokens and wall ms; `plain` the band's rate. + pub fn observe(l: *Lone, k: Constants, plain: *?f64, tokens: usize, ms: f64) void { + l.observeBand(k, plain, null, tokens, ms); + } + + /// observe, with the band's last verdict: a stream starting where drafting lost starts plain for its wait. + pub fn observeBand(l: *Lone, k: Constants, plain: *?f64, memo: ?*?Plain.Verdict, tokens: usize, ms: f64) void { + var switched = false; + if (!l.started) { + l.started = true; + l.left = k.drafted; + l.wait = k.wait; + if (memo) |m| if (m.*) |v| if (plain.*) |p| if (v.drafted < p * k.margin) { + l.drafted = v.drafted; + l.losses = k.confirm - 1; // one more losing probe confirms the band's loss again + l.switchTo(.plain, v.wait); + l.wait = @min(v.wait *| 2, k.wait_most); + switched = true; + }; + } + if (l.skip) { + l.skip = switched; // the next round still verifies this round's drafted hold + return; + } + l.tokens += @floatFromInt(tokens); + l.ms += ms; + l.left -|= 1; + if (l.left > 0) return; + const now = if (l.ms > 0) l.tokens / l.ms else 0; + const stretch_ms = l.ms; + l.tokens = 0; + l.ms = 0; + switch (l.mode) { + .drafted => { + l.drafted = blend(k, l.drafted, now); + l.drafted_ms = stretch_ms; + l.drafted_now = now; + if (plain.* == null) { + l.pending = true; + l.switchTo(.plain, k.plain); + } else l.verdict(k, plain.*.?, memo); + }, + .plain => { + plain.* = blend(k, plain.*, now); + if (l.pending) { + l.pending = false; + l.verdict(k, plain.*.?, memo); + } else l.switchTo(.drafted, k.drafted); + }, + } + } + + /// Drafting goes on while it beats plain by the margin, else plain rounds for a wait that doubles. + fn verdict(l: *Lone, k: Constants, plain: f64, memo: ?*?Plain.Verdict) void { + if (l.drafted.? >= plain * k.margin) { + l.wait = k.wait; + l.losses = 0; + l.switchTo(.drafted, k.drafted); + if (memo) |m| m.* = null; + return; + } + var wait = l.wait; + if (l.drafted_now >= plain * k.margin) { + // the blend lost but this stretch won: no loss counted, the band forgets its loss + l.losses = 0; + if (memo) |m| m.* = null; + } else { + l.losses += 1; + if (l.losses >= k.confirm) { + const lost_ms = l.drafted_ms * @max(0, 1 - l.drafted_now / plain); + const need: f64 = @ceil(lost_ms * plain / k.probe_cost); // plain rounds of about one token each + wait = @min(k.wait_most, @max(wait, @as(u32, @intFromFloat(@min(need, 1e6))))); + } + if (memo) |m| m.* = .{ .drafted = l.drafted.?, .wait = wait }; + } + l.switchTo(.plain, wait); + l.wait = @min(wait *| 2, k.wait_most); + } + + fn switchTo(l: *Lone, mode: Mode, rounds: u32) void { + if (mode != l.mode) { + l.skip = true; + l.fresh = mode == .drafted; + } + l.mode = mode; + l.left = @max(1, rounds); + } +}; + +fn blend(k: Constants, old: ?f64, now: f64) f64 { + return if (old) |r| r + k.blend * (now - r) else now; +} + +/// Rounds of one stream whose drafted rounds give `drafted` and plain ones `plain` tokens a ms: how many drafted. +fn run(l: *Lone, k: Constants, rate: *?f64, rounds: usize, plain: f64, drafted: f64) usize { + var n: usize = 0; + for (0..rounds) |_| { + const d = l.drafts(); + n += @intFromBool(d); + l.observe(k, rate, if (d) 2 else 1, if (d) 2 / drafted else 1 / plain); + } + return n; +} + +test "a lone stream whose drafted rounds lose runs plain but for probes that grow rarer" { + var l: Lone = .{}; + var rate: ?f64 = null; + const k: Constants = .{}; + const drafted = run(&l, k, &rate, 1000, 0.128, 0.122); + // the first drafted stretch, the plain rate's own, then probes after 16, 32, 64, 128, 256 and 512 plain rounds + try std.testing.expect(drafted <= 7 * (k.drafted + 1)); + try std.testing.expect(drafted > k.drafted); + try std.testing.expect(!l.drafts() or l.left <= k.drafted); + try std.testing.expectApproxEqAbs(@as(f64, 0.128), rate.?, 1e-9); +} + +test "a lone stream whose drafted rounds win keeps drafting; a later stream reuses the plain rate" { + var rate: ?f64 = null; + const k: Constants = .{}; + var first: Lone = .{}; + const n = run(&first, k, &rate, 1000, 0.128, 0.15); + try std.testing.expect(n > 1000 - k.plain - 2); // one plain stretch takes the plain rate + var second: Lone = .{}; + try std.testing.expectEqual(@as(usize, 300), run(&second, k, &rate, 300, 0.128, 0.15)); // no plain stretch again + var third: Lone = .{}; + const lost = run(&third, k, &rate, 300, 0.128, 0.11); + try std.testing.expect(lost < 300 / 3); // plain from its first verdict, without measuring plain again +} + +test "a drafted stretch's first round is not measured and its first hold is fresh" { + var l: Lone = .{}; + var rate: ?f64 = 0.2; + const k: Constants = .{ .drafted = 2, .wait = 2 }; + l.observe(k, &rate, 9, 1); // skipped: the stream's first round + l.observe(k, &rate, 1, 10); + l.observe(k, &rate, 1, 10); + try std.testing.expect(!l.drafts()); // 0.1 against 0.2 + try std.testing.expect(!l.fresh); + l.observe(k, &rate, 9, 1); // skipped: the last drafted round's hold + l.observe(k, &rate, 1, 5); + l.observe(k, &rate, 1, 5); + try std.testing.expect(l.drafts() and l.fresh); // the probe after the wait + try std.testing.expectApproxEqAbs(@as(f64, 0.2), rate.?, 1e-12); +} + +test "a probe that loses a third waits long enough that probes cost at most probe_cost; a band's verdict carries over" { + const k: Constants = .{}; + var rate: ?f64 = null; + var memo: ?Plain.Verdict = null; + var l: Lone = .{}; + // 130k on the M5: plain 8.67 ms a token, drafted 15.7 ms for 1.2 tokens + var drafted: usize = 0; + for (0..200) |_| { + const d = l.drafts(); + drafted += @intFromBool(d); + l.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7); + } + try std.testing.expect(drafted <= 2 * (1 + k.drafted)); // two losing stretches, then the probe_cost wait + try std.testing.expect(memo != null and memo.?.wait > 200); + var next: Lone = .{}; + var later: usize = 0; + for (0..200) |_| { + const d = next.drafts(); + later += @intFromBool(d); + next.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7); + } + try std.testing.expect(later <= 1); +} + +test "a winning stretch the blend still calls a loss counts no loss: no probe_cost wait" { + const k: Constants = .{ .drafted = 2, .wait = 2 }; + var rate: ?f64 = 0.2; + var memo: ?Plain.Verdict = null; + var l: Lone = .{}; + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 1, 10); + l.observeBand(k, &rate, &memo, 1, 10); // 0.1: a loss, wait 2 + try std.testing.expect(!l.drafts() and l.losses == 1 and memo != null); + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 1, 5); + l.observeBand(k, &rate, &memo, 1, 5); + try std.testing.expect(l.drafts()); + l.observeBand(k, &rate, &memo, 9, 1); // skipped + l.observeBand(k, &rate, &memo, 2, 8); + l.observeBand(k, &rate, &memo, 2, 8); // 0.25 alone wins; blended 0.175 loses + try std.testing.expect(!l.drafts() and l.losses == 0 and memo == null and l.left == 4); +} diff --git a/zig/src/core/lanes/lone_guard_test.zig b/zig/src/core/lanes/lone_guard_test.zig new file mode 100644 index 0000000000..2becd90158 --- /dev/null +++ b/zig/src/core/lanes/lone_guard_test.zig @@ -0,0 +1,115 @@ +//! The lone guard through the round loop: a lone stream drafts while drafting measures faster, else runs plain. +const std = @import("std"); +const Config = @import("config.zig").Config; +const Engine = @import("engine.zig").Engine; +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const fake = @import("fake.zig"); +const dr = @import("drafter.zig"); +const Drafted = @import("drafted.zig").Drafted; +const dt = @import("drafted_test.zig"); +const ct = @import("drafted_chances_test.zig"); + +const gpa = std.testing.allocator; + +/// A round's ms from its verify's rows and whether a hold ran before it, as a GPU that runs both in turn. +const Costs = struct { row: f64, extra: f64, hold: f64 }; + +const Cost = struct { + var costs: Costs = undefined; + var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined; + var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined; + var rows: usize = 1; + var held: bool = false; // a hold since the last verify + var before: bool = false; // the last verify ran after a hold + var holds: u32 = 0; + var verifies: u32 = 0; + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + rows = windows[0].rows(); + before = held; + held = false; + verifies += 1; + return inner_verify(ptr, windows, out); + } + + fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + for (items) |it| held = held or it.depth > 0; + holds += @intFromBool(held); + return inner_hold(ptr, items); + } + + fn clock() be.Clock { + return .{ .ptr = @ptrCast(&costs), .vtable = &.{ .start = start, .elapsed_ms = elapsed } }; + } + + fn start(_: *anyopaque) void {} + + fn elapsed(_: *anyopaque, _: be.Mark) f64 { + const extra: f64 = @floatFromInt(rows - 1); + return costs.row + costs.extra * extra + (if (before) costs.hold else 0); + } +}; + +const Run = struct { tokens: [400]u32 = undefined, len: usize = 0, holds: u32 = 0, rounds: u32 = 0, drafted: u64 = 0 }; + +/// One lone stream of 300 tokens; `guard` the drafter's guard switch, `costs` null: drafts off. +fn run(costs: ?Costs, guard: bool) !Run { + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var p: ct.Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .shared_guard = guard } }, .sure = 0.95, .unsure = 0.05 }; + defer p.deinit(); + var strict: dt.Strict = .{ .inner = target.backend() }; + defer strict.deinit(); + var seen = strict.backend(); + var table = seen.vtable.*; + Cost.inner_verify = table.verify; + table.verify = Cost.verify; + seen.vtable = &table; + var drafter = p.drafter(); + var hold_table = drafter.vtable.*; + Cost.inner_hold = hold_table.hold; + hold_table.hold = Cost.hold; + drafter.vtable = &hold_table; + Cost.costs = costs orelse .{ .row = 8, .extra = 0, .hold = 0 }; + Cost.holds = 0; + Cost.verifies = 0; + Cost.held = false; + var wrapped = try Drafted.init(gpa, seen, drafter); + defer wrapped.deinit(); + const drafts = costs != null; + var cfg = try Config.init(gpa, if (drafts) wrapped.facts(ct.base()) else ct.base(), 16, 15); + defer cfg.deinit(gpa); + var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), Cost.clock()); + defer engine.deinit(); + var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &dt.p1, .max_new = 300, .eos = &.{}, .drafts = drafts }); + defer s.deinit(gpa); + try engine.addStream(&s); + while (engine.activeCount() > 0) try engine.step(); + var out: Run = .{ .len = s.emitted().len, .holds = Cost.holds, .rounds = Cost.verifies, .drafted = s.drafted }; + @memcpy(out.tokens[0..out.len], s.emitted()); + return out; +} + +const slow: Costs = .{ .row = 8, .extra = 4, .hold = 9 }; // a hold and a verify row cost more than drafts save +const fast: Costs = .{ .row = 8, .extra = 1, .hold = 1 }; + +test "a lone stream whose drafted rounds lose runs plain rounds and probes now and then, replies unchanged" { + const plain = try run(null, true); + const unguarded = try run(slow, false); + const guarded = try run(slow, true); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], unguarded.tokens[0..unguarded.len]); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]); + try std.testing.expect(unguarded.holds * 10 > unguarded.rounds * 9); // without the guard nearly every round holds + try std.testing.expect(guarded.holds * 3 < guarded.rounds); // with it most rounds are plain + try std.testing.expect(guarded.holds > 16); // the first drafted stretch and probes +} + +test "a lone stream whose drafted rounds win keeps drafting with the guard on" { + const plain = try run(null, true); + const unguarded = try run(fast, false); + const guarded = try run(fast, true); + try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]); + try std.testing.expect(guarded.holds * 10 > guarded.rounds * 8); // drafting but for the plain rate's one stretch + try std.testing.expect(guarded.rounds * 10 < unguarded.rounds * 12); +} diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 4c4cb585e5..153da72e53 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -135,6 +135,9 @@ fn openedCommitted(e: *Engine, s: *Stream, reach: Reach, hook: ?be.Committed) !v /// The first drafts' depth for `reach`. pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { var who = win.who(s); + // a stream joining shared rounds the guard runs plain holds no drafts, as the guard measured it faster + const others = e.activeCount(); + if (others > 0 and e.cfg.family_streams and e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(others + 1)) return 0; 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)); } diff --git a/zig/src/core/lanes/pipelined.zig b/zig/src/core/lanes/pipelined.zig new file mode 100644 index 0000000000..99de33304b --- /dev/null +++ b/zig/src/core/lanes/pipelined.zig @@ -0,0 +1,70 @@ +//! A lone stream's pipelined steps: the next forward queued before the last one's token is read. +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const ev = @import("events.zig"); +const win = @import("windows.zig"); +const trail = @import("trail.zig"); +const Engine = @import("engine.zig").Engine; +const Result = @import("engine.zig").Result; +const Stream = sm.Stream; +const Feed = be.Feed; +const f = ev.f; +const str = trail.str; + +/// One token, the next forward queued first; a copied continuation ahead switches to verify windows. +pub fn round(e: *Engine, s: *Stream) !Result { + const before = s.mode; + // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one + const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe; + var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} }; + if (mode == .verify or mode == .exit) { + const verifying = mode == .verify and e.cfg.family_width >= 2; + const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{}; + if (proposal.tokens.len > 0) { + const copied = try e.arena.allocator().dupe(u32, proposal.tokens); + result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled }); + if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit; + } else { + // nothing to copy: back to steps queued ahead (this one lands next round) + s.mode = .pipe; + try e.queueNext(s, .{ .value = s.pending.? }); + } + try trail.pipe(e, s, before, result.got); + return result; + } + var current: Feed = .{ .handle = s.inflight.? }; + s.inflight = null; + if (try e.forcedNext(s)) |t| current = .{ .value = t }; + var token: u32 = undefined; + if (mode == .drain) { + token = try e.readFeed(current); // the last queued step: no new one + s.mode = .verify; + } else { + try e.queueNext(s, current); // the GPU starts the next step first + token = try e.readFeed(current); + } + s.rounds += 1; + const landed = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat + if (e.cfg.root_copies) s.root_copy_miss = false; + const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts; + if (piping and (try win.copyProposal(e, s)).tokens.len > 0) + s.mode = .drain; // a copy window is ahead: land the queued step + result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]); + try trail.pipe(e, s, before, result.got); + return result; +} + +/// Land the queued token without queuing another (a shared round is synchronous). +pub fn land(e: *Engine, s: *Stream) !void { + var current: Feed = .{ .handle = s.inflight.? }; + s.inflight = null; + if (try e.forcedNext(s)) |t| current = .{ .value = t }; + const token = try e.readFeed(current); + s.rounds += 1; + const landed = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + s.mode = .verify; + try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) }); +} diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig new file mode 100644 index 0000000000..52dc848b95 --- /dev/null +++ b/zig/src/core/lanes/shared_guard.zig @@ -0,0 +1,176 @@ +//! Whether a shared round drafts at all: drafted and plain rounds' measured tokens a ms, by stream count. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const depth = @import("depth.zig"); + +pub const Constants = struct { + explore: [2]u32 = .{ 16, 4 }, // measured rounds plain and drafted rounds run when their rate is (re)taken + exploit: u32 = 512, // measured rounds the faster mode runs before the slower one is measured again + longest: u5 = 6, // each time the same mode wins again its run doubles, up to exploit << longest + margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much + blend: f64 = 0.5, // the newest stretch's weight in a mode's rate +}; + +/// One stream count's state: each mode's rate (0 plain, 1 drafted), the mode running and its stretch. +const Entry = struct { + streams: u32, + rate: [2]?f64 = .{ null, null }, + mode: u1 = 0, // plain first: drafted rounds that lose cost more than plain ones, so they are measured short + exploring: bool = true, + left: u32, + tokens: f64 = 0, + ms: f64 = 0, + skip: bool = true, // a mode's first round still carries the other mode's drafts or holds: not measured + best: ?u1 = null, // the mode the last measurement chose + streak: u5 = 0, // how many times in a row it chose that mode +}; + +pub const Guard = struct { + gpa: Allocator, + k: Constants = .{}, + list: std.ArrayList(Entry) = .empty, + + pub fn deinit(g: *Guard) void { + g.list.deinit(g.gpa); + } + + fn entry(g: *Guard, streams: usize) !*Entry { + const n: u32 = @intCast(streams); + for (g.list.items) |*e| if (e.streams == n) return e; + try g.list.append(g.gpa, .{ .streams = n, .left = g.k.explore[0] }); + return &g.list.items[g.list.items.len - 1]; + } + + /// Whether a shared round of `streams` streams may draft (false: every stream runs its row alone). + pub fn drafts(g: *Guard, streams: usize) !bool { + return (try g.entry(streams)).mode == 1; + } + + /// drafts(); a drafted stretch's first round grants every stream one draft, so it measures drafted rounds. + pub fn draftsFresh(g: *Guard, streams: usize, whos: []depth.Who) !bool { + const e = try g.entry(streams); + if (e.mode == 1 and e.skip) for (whos) |*w| { + w.granted = 1; + w.grant_short = false; + w.gain = null; // priced by its first draft's rate, not a grant of none + }; + return e.mode == 1; + } + + /// A finished shared round: its committed tokens and wall ms, in the mode drafts() gave it. + pub fn observe(g: *Guard, streams: usize, tokens: u64, ms: f64) !void { + const e = try g.entry(streams); + if (e.skip) { + e.skip = false; + return; + } + e.tokens += @floatFromInt(tokens); + e.ms += ms; + e.left -|= 1; + if (e.left > 0) return; + if (e.ms > 0) { + const now = e.tokens / e.ms; + e.rate[e.mode] = if (e.rate[e.mode]) |r| r + g.k.blend * (now - r) else now; + } + e.tokens = 0; + e.ms = 0; + const before = e.mode; + if (e.exploring and e.rate[e.mode ^ 1] == null) { + e.mode ^= 1; // the other mode has no rate yet: measure it next + } else if (e.exploring) { + e.mode = if (e.rate[1].? >= e.rate[0].? * g.k.margin) 1 else 0; + e.streak = if (e.best == e.mode) @min(e.streak + 1, g.k.longest) else 0; + e.best = e.mode; + e.exploring = false; + } else { + e.mode ^= 1; // the slower mode is measured again, in case the streams' text changed + e.exploring = true; + } + e.left = if (e.exploring) g.k.explore[e.mode] else g.k.exploit << e.streak; + e.skip = e.mode != before; + } +}; + +/// The tokens a round's outcomes committed (each outcome's `got`). +pub fn committed(outcomes: anytype) u64 { + var n: u64 = 0; + for (outcomes) |o| n += o.got.len; + return n; +} + +/// Rounds of a guard at `streams` with rates `plain` and `drafted` tokens a ms: how many it drafted. +fn run(g: *Guard, streams: usize, rounds: usize, plain: f64, drafted: f64) !usize { + var n: usize = 0; + for (0..rounds) |_| { + const d = try g.drafts(streams); + n += @intFromBool(d); + try g.observe(streams, if (d) 3 else 1, if (d) 3 / drafted else 1 / plain); + } + return n; +} + +test "a shared round drafts only while drafted rounds beat plain ones, measuring the slower mode again now and then" { + var g: Guard = .{ .gpa = std.testing.allocator }; + defer g.deinit(); + // drafted rounds slower at four streams: after one stretch of each, plain but for a re-measure a period + const slow = try run(&g, 4, 1000, 0.25, 0.20); + try std.testing.expect(slow <= 4 * (g.k.explore[1] + 1)); // the first measure and three re-measures, short ones + try std.testing.expect(slow > 0); + // the same verdict again and again: the slower mode is measured ever more rarely, under 1% of rounds + const later = try run(&g, 4, 20000, 0.25, 0.20); + try std.testing.expect(later * 500 < 20000); + // faster at one stream: drafted but for the plain re-measures + const fast = try run(&g, 1, 1000, 0.13, 0.15); + try std.testing.expect(fast > 800); + // each stream count keeps its own rates + try std.testing.expect(!try g.drafts(4) or (try g.entry(4)).exploring); + try std.testing.expect(try g.drafts(1) or (try g.entry(1)).exploring); +} + +test "a mode's first round is not measured: it carries the other mode's drafts" { + var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + try std.testing.expect(!try g.drafts(2)); // plain first + try g.observe(2, 0, 1); // skipped: the stream count's first round + try g.observe(2, 1, 4); + try g.observe(2, 1, 4); + try std.testing.expect(try g.drafts(2)); // plain measured at 0.25 a ms; drafted next + try g.observe(2, 100, 1); // skipped: plain's last held rows, a fantasy rate that would keep drafting + try g.observe(2, 1, 10); + try g.observe(2, 1, 10); + try std.testing.expect(!try g.drafts(2)); // plain at 0.25 beats drafted at 0.1 + try std.testing.expectEqual(@as(?f64, 0.1), (try g.entry(2)).rate[1]); +} + +test "a drafted stretch starts from a grant of one: a plain stretch's grants of none would hold nothing to measure" { + var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + var st: ?depth.State = null; + var whos = [_]depth.Who{.{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 0 }}; + try std.testing.expect(!try g.draftsFresh(2, &whos)); // plain first: grants as they are + try std.testing.expectEqual(@as(?u32, 0), whos[0].granted); + for (0..3) |_| try g.observe(2, 1, 4); + try std.testing.expect(try g.draftsFresh(2, &whos)); // the drafted stretch's first round + try std.testing.expectEqual(@as(?u32, 1), whos[0].granted); + try g.observe(2, 1, 4); // its first round, skipped + whos[0].granted = 0; + try std.testing.expect(try g.draftsFresh(2, &whos)); // later rounds keep the allocation's grants + try std.testing.expectEqual(@as(?u32, 0), whos[0].granted); +} + +test "a guard allocates once a stream count: its rounds after that take no allocation" { + var counted = std.testing.FailingAllocator.init(std.testing.allocator, .{}); + var g: Guard = .{ .gpa = counted.allocator(), .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } }; + defer g.deinit(); + for ([_]usize{ 1, 4, 8, 16 }) |streams| _ = try run(&g, streams, 1, 0.25, 0.2); + const first = .{ counted.allocations, counted.resize_index }; + var st: [16]?depth.State = @splat(null); + var whos: [16]depth.Who = undefined; + for (&whos, &st) |*w, *s| w.* = .{ .state = s, .draft_room = 100, .finished = false, .forced = false, .granted = 1 }; + for (0..200) |i| for ([_]usize{ 1, 4, 8, 16 }) |streams| { + _ = try g.draftsFresh(streams, whos[0..streams]); + try g.observe(streams, i % 3 + 1, 4); + }; + try std.testing.expectEqual(first[0], counted.allocations); + try std.testing.expectEqual(first[1], counted.resize_index); +} diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 8c348ff95f..41ffc9054c 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -225,6 +225,8 @@ pub const Stream = struct { served: ?u64 = null, proposal_depth: ?u32 = null, // null: legacy; otherwise the last routed request, zero until it succeeds granted: ?u32 = null, + grant_short: bool = false, // the last grant verified fewer drafts than the stream held + grant_gain: ?f64 = null, // the drafts its last grant expected to keep (the chances it was priced by) copy_width: ?u32 = null, pub fn init(gpa: Allocator, spec: Spec) !Stream { @@ -287,6 +289,8 @@ pub const Stream = struct { s.inflight = null; s.served = null; s.granted = null; + s.grant_short = false; + s.grant_gain = null; s.copy_width = null; s.proposal_depth = null; } diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index e8e48be2ad..41f87b769c 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -197,7 +197,7 @@ fn headTrunk(e: *Engine, s: *Stream, p: Plan, levels: usize) !?fill.Trunk { /// A copy from the context backed by `enter_match` ids (Python `_copy_proposal`), else Recall's when it should pay. pub fn copyProposal(e: *Engine, s: *Stream) !Copy { - if (e.cfg.max_copy == 0 or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{}; + if (e.cfg.max_copy == 0 or e.plain_round or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{}; 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) { @@ -226,7 +226,7 @@ pub fn probe(e: *Engine) depth.Probe { } pub fn who(s: *Stream) depth.Who { - return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .stream = s, .sampled = s.sampling != null }; + return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .gain = s.grant_gain, .stream = s, .sampled = s.sampling != null }; } /// Whether any of a round's streams draws keyed samples (its rows carry their host work). @@ -281,8 +281,16 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { const ms = own.get(1 + @as(i64, counts[0])); try trail.event(e, &.{ f("ev", trail.str("lone-price")), f("stream", trail.str(plans[0].stream.id)), f("rows", trail.int(counts[0] + 1)), f("expected", ev.bits(expected)), f("ms", if (ms) |value| ev.bits(value) else .null), f("probe", .{ .bool = trial != null }) }); } - for (plans, counts) |*p, c| { - if (p.kind == .head) p.stream.granted = c; + for (plans, counts, probs) |*p, c, pr| { + if (p.kind == .head) { + p.stream.granted = c; + if (e.cfg.granted_depth) { + p.stream.grant_short = c < p.count(); + var gain: f64 = 0; + for (pr[0..@min(c, pr.len)]) |chance| gain += chance; + p.stream.grant_gain = gain; + } + } if (p.kind == .forced or p.count() == 0 or c == p.count()) continue; try cut.prefix(a, p, c); }