diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index cb2159b03a..a7eac1ea7d 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -36,6 +36,7 @@ pub const Model = struct { 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 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) 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) @@ -68,6 +69,7 @@ pub const Config = struct { 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, @@ -147,6 +149,7 @@ pub const Config = struct { .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, .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, diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 999aaeab6a..e814988934 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -103,6 +103,7 @@ pub const Drafted = struct { m.draft_streams = d.batched; m.draft_probabilities = false; m.head_trees = d.trees; + m.shared_guard = d.shared_guard; 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_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..88492dd650 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -18,6 +18,7 @@ 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) }; /// 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..664447ceaf 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 @@ -74,6 +78,7 @@ pub const Engine = struct { .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, }, + .guard = .{ .gpa = gpa }, }; } @@ -81,6 +86,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 +130,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 +153,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 +201,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; @@ -255,6 +261,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 +292,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.drafts(plans.len); + 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 +323,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 +503,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/lanes.zig b/zig/src/core/lanes/lanes.zig index ff09b58fb2..9f9ab516a0 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -59,6 +59,8 @@ test { _ = @import("windows_grammar_test.zig"); _ = @import("pricing_test.zig"); _ = @import("depth_probe_test.zig"); + _ = @import("shared_guard.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/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..8adaee2646 --- /dev/null +++ b/zig/src/core/lanes/shared_guard.zig @@ -0,0 +1,145 @@ +//! 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; + +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; + } + + /// 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 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 }; + for (0..200) |i| for ([_]usize{ 1, 4, 8, 16 }) |streams| { + _ = try g.drafts(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/windows.zig b/zig/src/core/lanes/windows.zig index e8e48be2ad..b70c2f634d 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) {