lanes: a shared round drafts only while drafting measures faster (Model.shared_guard)

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,

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);

zig/src/core/lanes/drafted_guard_test.zig

@@ -0,0 +1,155 @@
+//! A drafter that asks for the shared-round guard (Facts.shared_guard): which shared rounds draft, replies unchanged.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Model = @import("config.zig").Model;
+const Engine = @import("engine.zig").Engine;
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const Drafted = @import("drafted.zig").Drafted;
+const dt = @import("drafted_test.zig");
+
+const gpa = std.testing.allocator;
+
+test "a stream joining shared rounds that the guard runs plain holds no first drafts; a lone one does" {
+    var target: fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = true } };
+    defer head.deinit();
+    var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+    defer wrapped.deinit();
+    const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+    var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+    defer cfg.deinit(gpa);
+    var clock: fake.FixedClock = .{};
+    var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+    defer engine.deinit();
+    var streams: [2]sm.Stream = undefined;
+    for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} });
+    defer for (&streams) |*s| s.deinit(gpa);
+    try engine.addStream(&streams[0]);
+    try std.testing.expectEqual(@as(u32, 1), head.holds); // alone: its first drafts are held
+    for (0..3) |_| try engine.step();
+    const before = head.holds;
+    try std.testing.expect(!try engine.guard.drafts(2)); // the guard measures plain shared rounds first
+    try engine.addStream(&streams[1]);
+    try std.testing.expectEqual(before, head.holds); // joining: none held
+    while (engine.activeCount() > 0) try engine.step();
+}
+
+test "without the guard's flag a stream joining shared rounds holds its first drafts as before" {
+    var target: fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true } };
+    defer head.deinit();
+    var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+    defer wrapped.deinit();
+    const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+    var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+    defer cfg.deinit(gpa);
+    try std.testing.expect(!cfg.shared_guard);
+    var clock: fake.FixedClock = .{};
+    var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+    defer engine.deinit();
+    var streams: [2]sm.Stream = undefined;
+    for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} });
+    defer for (&streams) |*s| s.deinit(gpa);
+    try engine.addStream(&streams[0]);
+    for (0..3) |_| try engine.step();
+    const before = head.holds;
+    try engine.addStream(&streams[1]);
+    try std.testing.expectEqual(before + 1, head.holds); // joining: its first drafts held
+    while (engine.activeCount() > 0) try engine.step();
+}
+
+const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0 };
+
+/// Three streams decoded together: drafts off, or the fake drafter with or without the guard.
+fn run(drafts: bool, guard: bool) !Run {
+    var target: fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } };
+    defer head.deinit();
+    var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+    defer wrapped.deinit();
+    const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+    var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15);
+    defer cfg.deinit(gpa);
+    var clock: fake.FixedClock = .{};
+    var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock());
+    defer engine.deinit();
+    var streams: [3]sm.Stream = undefined;
+    for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = drafts });
+    defer for (&streams) |*s| s.deinit(gpa);
+    for (&streams) |*s| try engine.addStream(s);
+    while (engine.activeCount() > 0) try engine.step();
+    var out: Run = .{ .holds = head.holds };
+    for (&streams, 0..) |*s, i| {
+        out.lens[i] = s.emitted().len;
+        @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+    }
+    return out;
+}
+
+test "the guard changes only which shared rounds draft: replies equal one-token decode with it on and off" {
+    const plain = try run(false, false);
+    const off = try run(true, false);
+    const on = try run(true, true);
+    for (0..3) |i| {
+        try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], off.tokens[i][0..off.lens[i]]);
+        try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], on.tokens[i][0..on.lens[i]]);
+    }
+    try std.testing.expect(on.holds < off.holds); // its plain stretch holds no drafts
+}
+
+/// A clock that reads the same ms for every span: two of them with different ms stand for a fast and a slow machine.
+const FlatClock = struct {
+    ms: f64,
+
+    fn clock(c: *FlatClock) @import("backend.zig").Clock {
+        return .{ .ptr = c, .vtable = &.{ .start = start, .elapsed_ms = elapsed } };
+    }
+
+    fn start(_: *anyopaque) void {}
+
+    fn elapsed(ptr: *anyopaque, _: @import("backend.zig").Mark) f64 {
+        const c: *FlatClock = @ptrCast(@alignCast(ptr));
+        return c.ms;
+    }
+};
+
+/// Three streams, the guard asked for or not, timed or not, on a clock reading `ms` for every span.
+fn runClock(timed: bool, guard: bool, ms: f64) !Run {
+    var target: fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } };
+    defer head.deinit();
+    var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+    defer wrapped.deinit();
+    const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+    var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+    defer cfg.deinit(gpa);
+    cfg.timed = timed;
+    var clock: FlatClock = .{ .ms = ms };
+    var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+    defer engine.deinit();
+    var streams: [3]sm.Stream = undefined;
+    for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = true });
+    defer for (&streams) |*s| s.deinit(gpa);
+    for (&streams) |*s| try engine.addStream(s);
+    while (engine.activeCount() > 0) try engine.step();
+    var out: Run = .{ .holds = head.holds };
+    for (&streams, 0..) |*s, i| {
+        out.lens[i] = s.emitted().len;
+        @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+    }
+    return out;
+}
+
+test "with timed off the shared guard keeps no plain stretch: the same schedule on a fast and a slow clock" {
+    const fast = try runClock(false, true, 0.5);
+    const slow = try runClock(false, true, 400.0);
+    try std.testing.expectEqual(fast.holds, slow.holds);
+    for (0..3) |i| try std.testing.expectEqualSlices(u32, fast.tokens[i][0..fast.lens[i]], slow.tokens[i][0..slow.lens[i]]);
+    const unasked = try runClock(false, false, 0.5);
+    try std.testing.expectEqual(unasked.holds, fast.holds); // as if the guard were not asked for
+}

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.

zig/src/core/lanes/engine.zig

@@ -3,6 +3,7 @@ const std = @import("std");
 const Allocator = std.mem.Allocator;
 const Config = @import("config.zig").Config;
 const depth = @import("depth.zig");
+const shared_guard = @import("shared_guard.zig");
 const accept = @import("accept.zig");
 const be = @import("backend.zig");
 const sm = @import("stream.zig");
@@ -16,6 +17,7 @@ const lone_costs = @import("lone_costs.zig");
 const open = @import("open.zig");
 const handover = @import("handover.zig");
 const trim = @import("trim.zig");
+const pipelined = @import("pipelined.zig");
 const Stream = sm.Stream;
 const Feed = be.Feed;
 const LogRow = @import("logprob.zig").Row;
@@ -24,7 +26,7 @@ const f = ev.f;
 const str = trail.str;
 const int = trail.int;
 const Outcome = struct { got: []const u32, path: []const u32, cut: ?usize, follow: []const u32 };
-const Result = struct { rows: u32, keep: u32, got: []const u32 };
+pub const Result = struct { rows: u32, keep: u32, got: []const u32 };
 const Unread = struct { event: usize, handle: u64 };
 pub const Engine = struct {
     const pieces_ = @import("pieces.zig");
@@ -42,11 +44,13 @@ pub const Engine = struct {
     log: ?*ev.Log = null,
     trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig)
     rule: depth.Rule,
+    guard: shared_guard.Guard,
     live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`)
     drafted: u64 = 0,
     accepted: u64 = 0,
     shared_rounds: u64 = 0,
     alone: bool = true, // this step has one live stream: its copies may take their ramp
+    plain_round: bool = false, // a shared round the guard keeps plain: no stream copies or drafts, as with drafts off
     lone_kept: [2][lone_costs.bands]?lone_costs.Kept = @splat(@splat(null)), // the newest lone calibrations, greedy then sampled, a band each
     steps: u64 = 0,
     arena: std.heap.ArenaAllocator, // a round's temporaries
@@ -74,13 +78,15 @@ pub const Engine = struct {
                 .batch_rows = cfg.batch_rows,
                 .batched_heads = cfg.draft_streams,
             },
+            .guard = .{ .gpa = gpa },
         };
     }
 
     pub fn deinit(e: *Engine) void {
         e.live.deinit(e.gpa);
         e.unread.deinit(e.gpa);
         e.rule.deinit();
+        e.guard.deinit();
         for (&e.lone_kept) |*mode| for (mode) |*k| if (k.*) |*x| x.deinit(e.gpa);
         e.arena.deinit();
     }
@@ -124,7 +130,7 @@ pub const Engine = struct {
         if (live.items.len > 1 and e.cfg.family_streams) {
             // a shared round is synchronous: a stream that ran a step ahead lands its queued token first
             for (live.items) |s| {
-                if (s.inflight != null) try e.landInflight(s);
+                if (s.inflight != null) try pipelined.land(e, s);
             }
             var kept: usize = 0;
             for (live.items) |s| {
@@ -147,7 +153,7 @@ pub const Engine = struct {
             for (live.items) |s| {
                 // logprobs and grammars take verify windows: rows' logits are read, or bits made, before a forward
                 const windowed = (e.cfg.family_mtp and s.drafts) or !e.cfg.pipelined or s.logprobs != null or s.grammar != null;
-                const r = if (windowed) try e.familyRound(s, null) else try e.pipelinedRound(s);
+                const r = if (windowed) try e.familyRound(s, null) else try pipelined.round(e, s);
                 s.min_rows = if (s.min_rows == 0) r.rows else @min(s.min_rows, r.rows);
             }
         }
@@ -195,7 +201,7 @@ pub const Engine = struct {
     }
 
     /// One stream's round: verify pending and drafted tokens together, keep drafts to the first mismatch.
-    fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result {
+    pub fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result {
         e.clock.start();
         const priced = s.priced;
         s.priced = null;
@@ -255,6 +261,8 @@ pub const Engine = struct {
         e.clock.start();
         const a = e.arena.allocator();
         const plans = try a.alloc(Plan, entries.len);
+        e.plain_round = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(entries.len);
+        defer e.plain_round = false;
         for (entries, plans) |s, *p| p.* = try win.plan(e, s, null);
         try win.allocate(e, plans);
         for (plans) |*p| try win.constrain(e, p);
@@ -284,7 +292,9 @@ pub const Engine = struct {
         for (whos, heads.items) |*w, i| w.* = win.who(plans[i].stream);
         const budgets = try a.alloc(i64, heads.items.len);
         const sampled = win.anySampled(plans);
-        try e.rule.budgets(whos, win.probe(e), sampled, budgets);
+        // the guard's plain stretches: no stream drafts, so the round costs what it would without a head
+        const guarded = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(plans.len);
+        if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets);
         const ids = try a.alloc(ev.Value, heads.items.len);
         for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id);
         try trail.event(e, &.{ f("ev", str("budgets")), f("streams", .{ .list = ids }), f("depths", .{ .i64s = budgets }) });
@@ -313,6 +323,7 @@ pub const Engine = struct {
             try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) });
             try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets));
         }
+        if (e.cfg.shared_guard and e.cfg.timed) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead));
         if (e.backend.round_observer != null) e.backend.observed(paths);
     }
 
@@ -492,64 +503,6 @@ pub const Engine = struct {
         }
     }
 
-    /// One token, the next forward queued first; a copied continuation ahead switches to verify windows.
-    fn pipelinedRound(e: *Engine, s: *Stream) !Result {
-        const before = s.mode;
-        // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one
-        const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe;
-        var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} };
-        if (mode == .verify or mode == .exit) {
-            const verifying = mode == .verify and e.cfg.family_width >= 2;
-            const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{};
-            if (proposal.tokens.len > 0) {
-                const copied = try e.arena.allocator().dupe(u32, proposal.tokens);
-                result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled });
-                if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit;
-            } else {
-                // nothing to copy: back to steps queued ahead (this one lands next round)
-                s.mode = .pipe;
-                try e.queueNext(s, .{ .value = s.pending.? });
-            }
-            try trail.pipe(e, s, before, result.got);
-            return result;
-        }
-        var current: Feed = .{ .handle = s.inflight.? };
-        s.inflight = null;
-        if (try e.forcedNext(s)) |t| current = .{ .value = t };
-        var token: u32 = undefined;
-        if (mode == .drain) {
-            token = try e.readFeed(current); // the last queued step: no new one
-            s.mode = .verify;
-        } else {
-            try e.queueNext(s, current); // the GPU starts the next step first
-            token = try e.readFeed(current);
-        }
-        s.rounds += 1;
-        const landed = try s.commit(e.gpa, &.{token}, &.{});
-        s.pending = token;
-        if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat
-        if (e.cfg.root_copies) s.root_copy_miss = false;
-        const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts;
-        if (piping and (try win.copyProposal(e, s)).tokens.len > 0)
-            s.mode = .drain; // a copy window is ahead: land the queued step
-        result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]);
-        try trail.pipe(e, s, before, result.got);
-        return result;
-    }
-
-    /// Land the queued token without queuing another (a shared round is synchronous).
-    fn landInflight(e: *Engine, s: *Stream) !void {
-        var current: Feed = .{ .handle = s.inflight.? };
-        s.inflight = null;
-        if (try e.forcedNext(s)) |t| current = .{ .value = t };
-        const token = try e.readFeed(current);
-        s.rounds += 1;
-        const landed = try s.commit(e.gpa, &.{token}, &.{});
-        s.pending = token;
-        s.mode = .verify;
-        try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) });
-    }
-
     /// The first token's draw: among the grammar's allowed tokens when the stream has one that constrains it.
     pub fn drawFirst(e: *Engine, s: *Stream, position: u64) !u64 {
         const g = s.grammar orelse return e.backend.first(s, position);