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/4] 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/4] 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/4] 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 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 4/4] 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)); }