From 47eecb523aa3902616fdcda85d353f7aea7c5acb Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:02:12 +0100 Subject: [PATCH 1/8] lanes: a lone stream's pipelined steps move to pipelined.zig engine.zig is 592 lines and the shared-round guard needs room in it. pipelinedRound and landInflight move unchanged to pipelined.zig as round and land; Engine.familyRound and Result become public for them. No behaviour change. --- zig/src/core/lanes/engine.zig | 67 +++--------------------------- zig/src/core/lanes/pipelined.zig | 70 ++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 62 deletions(-) create mode 100644 zig/src/core/lanes/pipelined.zig diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a92b27ae7f..1664587c6d 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -16,6 +16,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 +25,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"); @@ -124,7 +125,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 +148,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 +196,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; @@ -492,64 +493,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/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 ..] }) }); +} From 5d9fcf68f87dd46ff653c17a2fd937fd09074ce1 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:03:43 +0100 Subject: [PATCH 2/8] lanes: a shared round drafts only while drafting measures faster (Model.shared_guard) A drafter whose hold is a pass of its own costs a shared round more than the depth rule's head steps price: the rule treats the rest of a round's time as overhead that plain rounds pay too. Facts.shared_guard (and Model.shared_guard for a backend) lets the engine compare the two modes directly. shared_guard.zig keeps, for each stream count, committed tokens over the round's wall time in plain stretches (no budgets, so no holds) and in drafted stretches. It starts plain, measures plain rounds for 16 rounds and drafted ones for 4, runs the faster mode for 512 rounds and doubles that, up to 64 times, while the same mode wins again; then it measures the other mode again. Drafted rounds must beat plain ones by 5%. A mode's first round still carries the other mode's drafts or holds and is not measured. An entry is allocated once a stream count; later rounds take no allocation. A plain stretch changes only how many drafts a window carries, never what the target draws, so replies stay the same tokens. Without the flag the engine never consults the guard. Tests: the guard's choice and re-measures, ever rarer re-measures of a losing mode, its skipped first rounds, no allocation after a stream count's first round, and replies equal to one-token decode with the guard on and off. --- zig/src/core/lanes/config.zig | 3 + zig/src/core/lanes/drafted.zig | 1 + zig/src/core/lanes/drafted_guard_test.zig | 51 ++++++++ zig/src/core/lanes/drafter.zig | 1 + zig/src/core/lanes/engine.zig | 9 +- zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/shared_guard.zig | 145 ++++++++++++++++++++++ 7 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 zig/src/core/lanes/drafted_guard_test.zig create mode 100644 zig/src/core/lanes/shared_guard.zig 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..e0e38868a4 --- /dev/null +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -0,0 +1,51 @@ +//! 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; + +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 +} 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 1664587c6d..ba545f18c4 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"); @@ -43,6 +44,7 @@ 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, @@ -75,6 +77,7 @@ pub const Engine = struct { .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, }, + .guard = .{ .gpa = gpa }, }; } @@ -82,6 +85,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(); } @@ -285,7 +289,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 !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 }) }); @@ -314,6 +320,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) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); if (e.backend.round_observer != null) e.backend.observed(paths); } 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/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); +} From 1b57e583855b552aa003004e80db702ac1e6ae4e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:04:57 +0100 Subject: [PATCH 3/8] lanes: a shared round the guard keeps plain proposes no copies, and a stream joining one holds no drafts With the shared-round guard on, a plain round still verified the lane core's context and Recall copies, which are drafts too, priced from a prior; and a stream joining other streams held its first drafts even while the guard ran plain rounds at that stream count, a hold pass for each join. The engine now marks a shared round plain when the guard runs that stream count plain, and copyProposal proposes nothing for it, so the guard's plain mode is the drafts-off round. Open.firstDepth gives no drafts to a stream that joins live streams while the guard runs their count plus one plain. Both are behind Model.shared_guard; without it plain_round stays false and firstDepth is unchanged. Tests: a guarded plain round proposes no copies; a stream joining plain-guarded rounds holds no first drafts and, without the flag, holds them as before. --- zig/src/core/lanes/drafted_guard_test.zig | 51 +++++++++++++++++++++++ zig/src/core/lanes/engine.zig | 3 ++ zig/src/core/lanes/engine_test.zig | 3 ++ zig/src/core/lanes/open.zig | 3 ++ zig/src/core/lanes/windows.zig | 2 +- 5 files changed, 61 insertions(+), 1 deletion(-) diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig index e0e38868a4..ca6d6bc568 100644 --- a/zig/src/core/lanes/drafted_guard_test.zig +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -10,6 +10,57 @@ 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. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index ba545f18c4..5349b38904 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -50,6 +50,7 @@ pub const Engine = struct { 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 @@ -260,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 !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); 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/open.zig b/zig/src/core/lanes/open.zig index 4c4cb585e5..733d3bb2f8 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 !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/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) { From 47eaf78bd7278d5dcbd58d63d6d59559f19523dc Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:06:32 +0100 Subject: [PATCH 4/8] lanes: an external drafter's chain chances price verify rows, and its holds follow the grants A drafter that gives each held chain draft's chance (Drafter.VTable.chances) may set Facts.chances. Drafted.facts then sets Model.draft_probabilities and Model.granted_depth, so the window allocator prices each stream's verify rows by those chances, as it does a head's node chances. With Facts.confidence_floor it instead verifies the drafts at least that likely (a fixed window when 0). With granted_depth a stream's next hold is as deep as its last grant plus a slack: one more draft when the grant took every held draft, none when it stopped short of them (Stream.grant_short), and after a grant of none a probe of draft_slack drafts at the depth rule's probe interval. Facts.slack sets one slack for all three. A shared round's granted depths are cut so the drafting streams' blocks fit one hold pass of batch_rows. When the shared-round guard starts a drafted stretch, its first round grants every stream one draft, so the stretch measures drafted rounds rather than plain ones left by grants of none. All of it is behind Facts.chances and Model.granted_depth, off by default; windows.allocate records grant_short only for granted depths. Tests: chances price the chain's rows and drafted replies equal one-token decode alone and at 1, 4 and 8 shared streams, with the scheduled verify ragged across streams and verifying fewer rows than a fixed window; the facts reach the lane config; granted depths, their slacks and the one-pass cut; holds as wide as the grants. --- zig/src/core/lanes/config.zig | 11 +- zig/src/core/lanes/depth.zig | 62 ++++- zig/src/core/lanes/drafted.zig | 7 + zig/src/core/lanes/drafted_chances_test.zig | 266 ++++++++++++++++++++ zig/src/core/lanes/drafter.zig | 3 + zig/src/core/lanes/engine.zig | 3 +- zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/shared_guard.zig | 32 ++- zig/src/core/lanes/stream.zig | 2 + zig/src/core/lanes/windows.zig | 7 +- 10 files changed, 387 insertions(+), 7 deletions(-) create mode 100644 zig/src/core/lanes/drafted_chances_test.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index a7eac1ea7d..340788ccfe 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -35,6 +35,8 @@ 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) head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them @@ -65,6 +67,7 @@ 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 @@ -148,6 +151,7 @@ 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, @@ -160,7 +164,12 @@ 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, + }, .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..3c45e6ee6d 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -11,6 +11,8 @@ 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 }; /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). @@ -34,6 +36,7 @@ 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) stream: ?*anyopaque = null, sampled: bool = false, // keyed draws: host work a row that a greedy round does not have }; @@ -58,6 +61,7 @@ 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 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,9 +195,28 @@ pub const Rule = struct { return true; } + /// 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)); + const granted = who.granted orelse return most; + _ = try r.rates(who); + const st = &who.state.*.?; + 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)); + } + /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 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.*); @@ -202,6 +225,15 @@ pub const Rule = struct { 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); + } + /// 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) { @@ -212,8 +244,9 @@ pub const Rule = struct { for (whos, out) |*w, *o| { 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); return; } const most = try r.gpa.alloc(i64, whos.len); @@ -380,3 +413,28 @@ 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 +} diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index e814988934..50f5739913 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -104,6 +104,13 @@ pub const Drafted = struct { m.draft_probabilities = false; m.head_trees = d.trees; m.shared_guard = d.shared_guard; + 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..02c5eb4cb1 --- /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. +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; + } + + 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 }; + } + + 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; + +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/drafter.zig b/zig/src/core/lanes/drafter.zig index 88492dd650..35a79e2e1c 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -19,6 +19,9 @@ pub const Facts = struct { 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) }; /// 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 5349b38904..9ff27b18d5 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -74,6 +74,7 @@ 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, .confidence_floor = cfg.draft_confidence, .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, @@ -293,7 +294,7 @@ pub const Engine = struct { const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); // 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 !try e.guard.drafts(plans.len); + const guarded = e.cfg.shared_guard 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); diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 9f9ab516a0..8f8adf1e1b 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -54,6 +54,7 @@ test { _ = @import("recorder_test.zig"); _ = @import("gpu_full_test.zig"); _ = @import("drafted_test.zig"); + _ = @import("drafted_chances_test.zig"); _ = @import("root_copy_test.zig"); _ = @import("drafted_grammar_test.zig"); _ = @import("windows_grammar_test.zig"); diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig index 8adaee2646..cdaf870736 100644 --- a/zig/src/core/lanes/shared_guard.zig +++ b/zig/src/core/lanes/shared_guard.zig @@ -1,6 +1,7 @@ //! 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 @@ -45,6 +46,16 @@ pub const Guard = struct { 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; + }; + 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); @@ -130,14 +141,33 @@ test "a mode's first round is not measured: it carries the other mode's drafts" 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.drafts(streams); + _ = try g.draftsFresh(streams, whos[0..streams]); try g.observe(streams, i % 3 + 1, 4); }; try std.testing.expectEqual(first[0], counted.allocations); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 8c348ff95f..a35c8aaddc 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -225,6 +225,7 @@ 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 copy_width: ?u32 = null, pub fn init(gpa: Allocator, spec: Spec) !Stream { @@ -287,6 +288,7 @@ pub const Stream = struct { s.inflight = null; s.served = null; s.granted = null; + s.grant_short = false; 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 b70c2f634d..278b761325 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -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, .stream = s, .sampled = s.sampling != null }; } /// Whether any of a round's streams draws keyed samples (its rows carry their host work). @@ -282,7 +282,10 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { 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; + if (p.kind == .head) { + p.stream.granted = c; + if (e.cfg.granted_depth) p.stream.grant_short = c < p.count(); + } if (p.kind == .forced or p.count() == 0 or c == p.count()) continue; try cut.prefix(a, p, c); } From 27089360acf7ff50ce12840577559445cb02d20e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:07:53 +0100 Subject: [PATCH 5/8] lanes: a drafter whose hold is a pass of its own: no hold beside a copy, and holds kept only while they pay The depth rule asks one head draft for a stream with a copy ahead, and a shared round's granted holds were priced only by their verify rows. A drafter whose hold runs as a pass of its own (Facts.copy_first) pays that pass for the one draft, and the copy round then drops it; Model.copy_first now asks it for none. A drafter may also give its hold pass's timed cost by blocks and depth (Facts.hold_costs, priced by hold_price.zig). With granted depths a shared round then keeps granted holds, most expected drafts a row first, only while each raises the next round's expected tokens a ms over its verify rows and the pass; a probe after a grant of none keeps its hold. A stream's last grant records the drafts it expected to keep (Stream.grant_gain). Both are off by default (copy_first false, hold_costs empty). Tests: with copy_first no hold is asked for a copy round and replies equal one-token decode at 1 and 3 streams; a hold's cost between and past its timed passes; granted holds kept while they pay, none for a dear pass, and a probe's hold kept. --- zig/src/core/lanes/config.zig | 4 + zig/src/core/lanes/depth.zig | 123 ++++++++++++++++++++- zig/src/core/lanes/drafted.zig | 2 + zig/src/core/lanes/drafted_copy_test.zig | 135 +++++++++++++++++++++++ zig/src/core/lanes/drafter.zig | 2 + zig/src/core/lanes/hold_price.zig | 58 ++++++++++ zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/shared_guard.zig | 1 + zig/src/core/lanes/stream.zig | 2 + zig/src/core/lanes/windows.zig | 11 +- 10 files changed, 335 insertions(+), 5 deletions(-) create mode 100644 zig/src/core/lanes/drafted_copy_test.zig create mode 100644 zig/src/core/lanes/hold_price.zig diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 340788ccfe..a20d219c55 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -39,6 +39,8 @@ pub const Model = struct { 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) @@ -169,6 +171,8 @@ pub const Config = struct { .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 3c45e6ee6d..02d096a07c 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -13,8 +13,13 @@ pub const Constants = struct { 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; + /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`). pub const State = struct { p: []f64, @@ -37,6 +42,7 @@ pub const Who = struct { 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 }; @@ -212,7 +218,7 @@ pub const Rule = struct { return @min(most, @as(i64, granted) + r.slack(who)); } - /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead). + /// 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.granted_depth) try r.grantDepth(who.*) @@ -221,7 +227,7 @@ pub const Rule = struct { 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; } @@ -234,6 +240,65 @@ pub const Rule = struct { 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) { @@ -247,6 +312,7 @@ pub const Rule = struct { 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); @@ -438,3 +504,56 @@ test "granted depths in a shared round fit one hold pass: every block is as deep 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 50f5739913..211f30112c 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -104,6 +104,8 @@ pub const Drafted = struct { 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; 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/drafter.zig b/zig/src/core/lanes/drafter.zig index 35a79e2e1c..3ba9feb671 100644 --- a/zig/src/core/lanes/drafter.zig +++ b/zig/src/core/lanes/drafter.zig @@ -22,6 +22,8 @@ pub const Facts = struct { 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/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 8f8adf1e1b..c5cd5f5e24 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"); @@ -55,6 +56,7 @@ test { _ = @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"); diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig index cdaf870736..52dc848b95 100644 --- a/zig/src/core/lanes/shared_guard.zig +++ b/zig/src/core/lanes/shared_guard.zig @@ -52,6 +52,7 @@ pub const Guard = struct { 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; } diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index a35c8aaddc..41ffc9054c 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -226,6 +226,7 @@ pub const Stream = struct { 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 { @@ -289,6 +290,7 @@ pub const Stream = struct { 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 278b761325..41f87b769c 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -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, .grant_short = s.grant_short, .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,10 +281,15 @@ 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| { + 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(); + 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); From b66a7629e60d7283d5a837f7b3d2632f03000896 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:09:02 +0100 Subject: [PATCH 6/8] lanes: a lone stream holds drafts only while its drafted rounds measure faster With granted depths a lone stream's verify rows are priced over the shared costs and the one-stream overhead, which is measured on drafted rounds and so holds the drafter's hold time; nothing compared a lone drafted round's tokens a ms with a plain round's. lone_guard.zig measures each lone stream's rounds of held drafts in stretches of 16 against plain rounds' tokens a ms, kept from stream to stream by sampling mode and context band and taken once in a plain stretch when unknown; copy rounds are left out of both. Drafting goes on while it beats plain by 5%; else the stream runs plain rounds and measures drafting again after 16 plain rounds, doubling to 512 while plain keeps winning. A drafted stretch's first hold drafts one and its slack, and a mode's first round is not measured. It runs only with granted depths and the shared-round guard on a timed engine (Facts.chances and Facts.shared_guard); shared rounds keep their own guard. Off by default. Tests: the guard's verdicts, rarer probes and kept plain rate; through the round loop a lone stream whose drafted rounds lose runs mostly plain and one whose drafted rounds win keeps drafting, replies equal to one-token decode. --- zig/src/core/lanes/depth.zig | 29 +++- zig/src/core/lanes/drafted_chances_test.zig | 8 +- zig/src/core/lanes/engine.zig | 2 + zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/lone_guard.zig | 155 ++++++++++++++++++++ zig/src/core/lanes/lone_guard_test.zig | 115 +++++++++++++++ 6 files changed, 305 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/lanes/lone_guard.zig create mode 100644 zig/src/core/lanes/lone_guard_test.zig diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 02d096a07c..e02c1e6240 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -19,6 +19,7 @@ pub const Constants = struct { 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 { @@ -27,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); @@ -45,6 +47,7 @@ pub const Who = struct { 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`). @@ -68,6 +71,9 @@ pub const Rule = struct { 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) @@ -209,15 +215,33 @@ pub const Rule = struct { /// 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; - _ = try r.rates(who); - const st = &who.state.*.?; + 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; + (try r.state(who)).lone.observe(r.lone_k, r.lone_plain.at(who.sampled, band), 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.granted_depth) @@ -307,6 +331,7 @@ 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, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe); diff --git a/zig/src/core/lanes/drafted_chances_test.zig b/zig/src/core/lanes/drafted_chances_test.zig index 02c5eb4cb1..d3f248e8da 100644 --- a/zig/src/core/lanes/drafted_chances_test.zig +++ b/zig/src/core/lanes/drafted_chances_test.zig @@ -17,7 +17,7 @@ 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. -const Priced = struct { +pub const Priced = struct { head: dt.FakeDrafter, sure: f64 = 0.7, unsure: f64 = 0.7, @@ -56,7 +56,7 @@ const Priced = struct { return true; } - fn drafter(p: *Priced) dr.Drafter { + pub fn drafter(p: *Priced) dr.Drafter { const base_ = p.head.drafter(); p.vtable = base_.vtable.*; p.inner_hold = base_.vtable.hold; @@ -65,7 +65,7 @@ const Priced = struct { return .{ .ptr = &p.head, .vtable = &p.vtable }; } - fn deinit(p: *Priced) void { + pub fn deinit(p: *Priced) void { p.at.deinit(gpa); p.head.deinit(); } @@ -91,7 +91,7 @@ const Seen = struct { var costs: [16]Cost = undefined; -fn base() Model { +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 }; } diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 9ff27b18d5..32e46c02e8 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -75,6 +75,7 @@ pub const Engine = struct { .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, @@ -238,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| { diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index c5cd5f5e24..a3b2587704 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -63,6 +63,8 @@ test { _ = @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"); diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig new file mode 100644 index 0000000000..d861174dc3 --- /dev/null +++ b/zig/src/core/lanes/lone_guard.zig @@ -0,0 +1,155 @@ +//! 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 + 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)), + + 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, + 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 { + if (!l.started) { + l.started = true; + l.left = k.drafted; + l.wait = k.wait; + } + if (l.skip) { + l.skip = false; + 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; + l.tokens = 0; + l.ms = 0; + switch (l.mode) { + .drafted => { + l.drafted = blend(k, l.drafted, now); + if (plain.* == null) { + l.pending = true; + l.switchTo(.plain, k.plain); + } else l.verdict(k, plain.*.?); + }, + .plain => { + plain.* = blend(k, plain.*, now); + if (l.pending) { + l.pending = false; + l.verdict(k, plain.*.?); + } else l.switchTo(.drafted, k.drafted); + }, + } + } + + /// Drafting goes on while its rate beats plain's by the margin, else plain rounds for a wait that doubles. + fn verdict(l: *Lone, k: Constants, plain: f64) void { + if (l.drafted.? >= plain * k.margin) { + l.wait = k.wait; + l.switchTo(.drafted, k.drafted); + return; + } + l.switchTo(.plain, l.wait); + l.wait = @min(l.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); +} 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); +} From 4fa02928a7484ccf7bc1a529936618e9489eb24e Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:14:47 +0100 Subject: [PATCH 7/8] lanes: the shared guard is off when Config.timed is false, so a schedule does not depend on the clock The guard planned, observed and chose with the wall clock even when timed was false. Round planning, joining a shared round and observation now skip it in that mode. A test runs the same three streams on a fast and a slow fake clock and checks equal holds and replies. --- zig/src/core/lanes/drafted_guard_test.zig | 53 +++++++++++++++++++++++ zig/src/core/lanes/engine.zig | 6 +-- zig/src/core/lanes/open.zig | 2 +- 3 files changed, 57 insertions(+), 4 deletions(-) diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig index ca6d6bc568..3715bec357 100644 --- a/zig/src/core/lanes/drafted_guard_test.zig +++ b/zig/src/core/lanes/drafted_guard_test.zig @@ -100,3 +100,56 @@ test "the guard changes only which shared rounds draft: replies equal one-token } 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/engine.zig b/zig/src/core/lanes/engine.zig index 5349b38904..664447ceaf 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -261,7 +261,7 @@ 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 !try e.guard.drafts(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); @@ -293,7 +293,7 @@ pub const Engine = struct { const budgets = try a.alloc(i64, heads.items.len); const sampled = win.anySampled(plans); // 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 !try e.guard.drafts(plans.len); + 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); @@ -323,7 +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) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead)); + 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); } diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 733d3bb2f8..153da72e53 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -137,7 +137,7 @@ 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 !try e.guard.drafts(others + 1)) return 0; + 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)); } From 6539b72beacf5bf911d5a5141f50da88cfcab4c6 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:53:24 +0100 Subject: [PATCH 8/8] lanes: a lone stream that drafting loses on keeps its loss for the next stream, and probes less often A loss is counted on the stretch's own rate. After two losses in a row the plain wait grows until a probe costs at most 3% of the plain time. The band keeps its last losing verdict, so the next stream in that band starts plain. Any winning stretch clears it. --- zig/src/core/lanes/depth.zig | 3 +- zig/src/core/lanes/lone_guard.zig | 101 +++++++++++++++++++++++++++--- 2 files changed, 96 insertions(+), 8 deletions(-) diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index e02c1e6240..9e2d3a2f7a 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -239,7 +239,8 @@ pub const Rule = struct { /// 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; - (try r.state(who)).lone.observe(r.lone_k, r.lone_plain.at(who.sampled, band), tokens, ms); + 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). diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig index d861174dc3..4657ffe829 100644 --- a/zig/src/core/lanes/lone_guard.zig +++ b/zig/src/core/lanes/lone_guard.zig @@ -7,6 +7,8 @@ pub const Constants = struct { 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 }; @@ -15,6 +17,13 @@ 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; @@ -34,6 +43,9 @@ pub const Lone = struct { 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. @@ -43,13 +55,26 @@ pub const Lone = struct { /// 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 = false; + l.skip = switched; // the next round still verifies this round's drafted hold return; } l.tokens += @floatFromInt(tokens); @@ -57,35 +82,54 @@ pub const Lone = struct { 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.*.?); + } else l.verdict(k, plain.*.?, memo); }, .plain => { plain.* = blend(k, plain.*, now); if (l.pending) { l.pending = false; - l.verdict(k, plain.*.?); + l.verdict(k, plain.*.?, memo); } else l.switchTo(.drafted, k.drafted); }, } } - /// Drafting goes on while its rate beats plain's by the margin, else plain rounds for a wait that doubles. - fn verdict(l: *Lone, k: Constants, plain: f64) void { + /// 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; } - l.switchTo(.plain, l.wait); - l.wait = @min(l.wait *| 2, k.wait_most); + 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 { @@ -153,3 +197,46 @@ test "a drafted stretch's first round is not measured and its first hold is fres 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); +}