lanes: a lone stream holds drafts only while its drafted rounds measure faster

zig/src/core/lanes/config.zig

@@ -35,7 +35,12 @@ pub const Model = struct {
     max_streams: u32 = 0, // a shared forward's streams at most (0: 32)
     shared_costs: []const Cost = &.{}, // shared forwards' ms by total rows (timed at load)
     draft_probabilities: bool = false, // the head gives each drafted node's chance of landing
+    granted_depth: bool = false, // with draft_probabilities: every stream drafts its last grant plus a slack
+    draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own)
     draft_streams: bool = false, // the head drafts every stream of a shared round in one batch
+    shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig)
+    copy_first: bool = false, // a copy ahead takes the round: the head holds no drafts for it (a drafter's own pass)
+    hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // a drafter's own hold pass, timed at load
     head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them
     join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join)
     piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them)
@@ -64,10 +69,12 @@ pub const Config = struct {
     batch_rows: u32,
     batch_streams: u32,
     node_probabilities: bool,
+    granted_depth: bool = false, // Model.granted_depth
     fill_lanes: u32 = 0, // rows a window may grow to with suffix-match branches (0: off; lanes/fill.zig)
     fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases)
     fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round
     draft_streams: bool,
+    shared_guard: bool = false,
     head_trees: bool,
     /// a prompt in pieces leaves its last token as the pending row of its first shared round's window
     join_tail: bool = false,
@@ -146,7 +153,9 @@ pub const Config = struct {
             .batch_rows = batch_rows,
             .batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1,
             .node_probabilities = m.draft_probabilities or m.draft_confidence != null,
+            .granted_depth = m.granted_depth and m.draft_probabilities,
             .draft_streams = m.draft_streams,
+            .shared_guard = m.shared_guard,
             .head_trees = m.head_trees,
             .join_tail = m.join_tail,
             .piece_runs = m.piece_runs,
@@ -157,7 +166,14 @@ pub const Config = struct {
             .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows),
             .lane_costs = try alloc.extendCosts(gpa, m.lane_costs, 64),
             .mtp_step_ms = m.mtp_step_ms,
-            .k = .{ .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate },
+            .k = .{
+                .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate,
+                .draft_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).draft_slack,
+                .grant_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).grant_slack,
+                .short_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).short_slack,
+                .copy_first = m.copy_first and !m.context_copies,
+                .hold_costs = m.hold_costs,
+            },
             .draft_step_ms = m.draft_step_ms,
         };
     }

zig/src/core/lanes/depth.zig

@@ -11,15 +11,24 @@ pub const Constants = struct {
     plain_margin: f64 = 1.05, // a drafted depth must beat plain rounds' tokens a ms by this much
     plain_wait_most: u64 = 128, // plain rounds between probes at most
     draft_slack: i64 = 2, // drafts a stream offers past what its last shared round granted it
+    grant_slack: i64 = 1, // granted depths: drafts past a grant that verified every draft the stream held
+    short_slack: i64 = 0, // granted depths: drafts past a grant that left some of them unverified
+    copy_first: bool = false, // a copy ahead takes the next round whole: the head drafts none for it (else one)
+    hold_costs: []const HoldCost = &.{}, // a drafter's own hold pass by blocks and depth: granted holds must pay for it
 };
 
+pub const HoldCost = @import("hold_price.zig").HoldCost;
+const holdMs = @import("hold_price.zig").holdMs;
+const lone_guard = @import("lone_guard.zig");
+
 /// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`).
 pub const State = struct {
     p: []f64,
     rounds: u64 = 0,
     plain: ?u64 = null,
     wait: ?u64 = null,
     ms: ?Table = null, // plain_guard: the stream's own round times
+    lone: lone_guard.Lone = .{}, // lone_guard: whether its lone holds draft
 
     pub fn deinit(st: *State, gpa: Allocator) void {
         gpa.free(st.p);
@@ -34,8 +43,11 @@ pub const Who = struct {
     finished: bool,
     forced: bool,
     granted: ?u32 = null,
+    grant_short: bool = false, // the last grant left some held drafts unverified (the width did not bind)
+    gain: ?f64 = null, // the drafts its last grant expected to keep (null: none priced)
     stream: ?*anyopaque = null,
     sampled: bool = false, // keyed draws: host work a row that a greedy round does not have
+    shared: bool = false, // a shared round's budget: the lone guard stays out
 };
 
 /// Whether a copied continuation is ahead of a stream (the round loop proposes; Python `_copy_proposal`).
@@ -58,6 +70,10 @@ pub const Rule = struct {
     draft_step_ms: f64 = 0, // a shared round's draft past a stream's first (budgets price it)
     plain_guard: bool,
     node_probabilities: bool,
+    granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too
+    lone_guard: bool = false, // granted_depth: a lone stream holds drafts only while its drafted rounds measure faster
+    lone_k: lone_guard.Constants = .{},
+    lone_plain: lone_guard.Plain = .{}, // plain rounds' tokens a ms, every lone stream's
     confidence_floor: ?f64 = null,
     batch_rows: u32,
     batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's)
@@ -191,17 +207,123 @@ pub const Rule = struct {
         return true;
     }
 
-    /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead).
+    /// Drafts past a grant: none when it stopped short of the held drafts (the width did not bind), else grant_slack.
+    fn slack(r: *const Rule, who: Who) i64 {
+        return if (who.grant_short) r.k.short_slack else r.k.grant_slack;
+    }
+
+    /// The last grant plus its slack (the deepest chain before any); after a grant of none, a probe of draft_slack.
+    fn grantDepth(r: *Rule, who: Who) !i64 {
+        const most = @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1));
+        if (r.lone_guard and !who.shared) {
+            const l = &(try r.state(who)).lone;
+            if (!l.drafts()) return 0;
+            if (l.fresh) {
+                l.fresh = false;
+                return @min(most, 1 + r.k.grant_slack);
+            }
+        }
+        const granted = who.granted orelse return most;
+        const st = try r.state(who);
+        if (granted == 0) return if (!r.plain_guard or r.probePlain(st)) @min(most, r.k.draft_slack) else 0;
+        st.plain = 0;
+        st.wait = r.k.depth_probe_every;
+        return @min(most, @as(i64, granted) + r.slack(who));
+    }
+
+    fn state(r: *Rule, who: Who) !*State {
+        _ = try r.rates(who);
+        return &who.state.*.?;
+    }
+
+    /// A lone round's committed tokens and wall ms at the stream's context band, for the lone guard.
+    pub fn observeLone(r: *Rule, who: Who, band: usize, tokens: usize, ms: f64) !void {
+        if (!r.lone_guard) return;
+        const memo = r.lone_plain.memo(who.sampled, band);
+        (try r.state(who)).lone.observeBand(r.lone_k, r.lone_plain.at(who.sampled, band), memo, tokens, ms);
+    }
+
+    /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead).
     pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 {
-        const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.node_probabilities)
+        const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth)
+            try r.grantDepth(who.*)
+        else if (r.node_probabilities)
             @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1))
         else
             try r.depth(who.*);
         if (who.forced) return 0;
-        if (d != 0 and probe.copy(probe.ptr, who)) return 1;
+        if (d != 0 and probe.copy(probe.ptr, who)) return @intFromBool(!r.k.copy_first);
         return d;
     }
 
+    /// A shared hold pads every block to its deepest: depths cut so the drafting streams' blocks fit `rows` (one pass).
+    fn capPass(out: []i64, rows: u32) void {
+        var n: i64 = 0;
+        for (out) |d| n += @intFromBool(d > 0);
+        if (n == 0) return;
+        const most = @max(1, @divFloor(@as(i64, rows), n));
+        for (out) |*d| d.* = @min(d.*, most);
+    }
+
+    /// Granted holds kept, most expected drafts a row first, while they raise the next round's expected tokens a ms.
+    fn payHolds(r: *Rule, whos: []const Who, sampled: bool, out: []i64) !void {
+        const order = try r.gpa.alloc(usize, whos.len);
+        defer r.gpa.free(order);
+        var m: usize = 0;
+        var probes: usize = 0; // a probe after a grant of none keeps its hold: it measures the stream again
+        var deepest: i64 = 0;
+        for (whos, out, 0..) |w, d, i| {
+            if (d <= 0) continue;
+            if (w.granted == @as(?u32, 0)) {
+                probes += 1;
+                deepest = @max(deepest, d);
+            } else {
+                order[m] = i;
+                m += 1;
+            }
+        }
+        if (m == 0) return;
+        const Ctx = struct {
+            r: *Rule,
+            whos: []const Who,
+            fn less(c: @This(), x: usize, y: usize) bool {
+                const gx = c.r.gainOf(c.whos[x]) catch 0;
+                const gy = c.r.gainOf(c.whos[y]) catch 0;
+                return gx * rowsOf(c.whos[y]) > gy * rowsOf(c.whos[x]) or (gx * rowsOf(c.whos[y]) == gy * rowsOf(c.whos[x]) and x < y);
+            }
+        };
+        std.mem.sort(usize, order[0..m], Ctx{ .r = r, .whos = whos }, Ctx.less);
+        const fixed = r.fixedMs(whos.len, sampled);
+        var rows: f64 = @floatFromInt(whos.len);
+        var tokens = rows;
+        const plain = r.shared_costs.get(@intCast(whos.len)) orelse return;
+        const probe_ms = if (probes > 0) holdMs(r.k.hold_costs, probes, deepest) else 0;
+        var best = tokens / (plain + fixed + probe_ms);
+        var keep: usize = 0;
+        for (order[0..m], 1..) |i, k| {
+            rows += rowsOf(whos[i]);
+            tokens += try r.gainOf(whos[i]);
+            deepest = @max(deepest, out[i]);
+            const verify = r.shared_costs.get(@intFromFloat(@ceil(rows))) orelse break;
+            const now = tokens / (verify + fixed + holdMs(r.k.hold_costs, k + probes, deepest));
+            if (now > best) {
+                best = now;
+                keep = k;
+            }
+        }
+        for (order[keep..m]) |i| out[i] = 0;
+    }
+
+    /// The drafts a granted stream's next round is expected to keep: its last grant's, or its first draft's rate.
+    fn gainOf(r: *Rule, w: Who) !f64 {
+        if (w.gain) |g| if (w.granted != null) return g;
+        return (try r.rates(w))[0];
+    }
+
+    fn rowsOf(w: Who) f64 {
+        return @floatFromInt(@max(1, w.granted orelse 1));
+    }
+
     /// Head drafts for the streams of a shared round, sized by allocation over their chain chances.
     pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void {
         if (r.confidence_floor != null) {
@@ -210,10 +332,13 @@ pub const Rule = struct {
         }
         if (r.node_probabilities) {
             for (whos, out) |*w, *o| {
+                w.shared = true;
                 const inner = try r.headDepth(w, null, probe);
                 const granted: i64 = if (w.granted) |g| g else r.most_drafts;
-                o.* = try r.headDepth(w, @min(inner, granted + r.k.draft_slack), probe);
+                o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe);
             }
+            if (r.granted_depth) capPass(out, r.batch_rows);
+            if (r.granted_depth and r.k.hold_costs.len > 0) try r.payHolds(whos, sampled, out);
             return;
         }
         const most = try r.gpa.alloc(i64, whos.len);
@@ -380,3 +505,81 @@ test "a shared round's drafts come out the same whichever depths its overhead wa
     }
     try std.testing.expectEqualSlices(i64, &got[0], &got[1]);
 }
+
+test "granted depths in a shared round fit one hold pass: every block is as deep as the deepest" {
+    const empty: Table = .{};
+    var rule: Rule = .{ .gpa = std.testing.allocator, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+    defer rule.deinit();
+    var states: [8]?State = @splat(null);
+    defer for (&states) |*st| if (st.*) |*x| x.deinit(std.testing.allocator);
+    var whos: [8]Who = undefined;
+    for (&whos, &states, 0..) |*w, *st, i| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = @intCast(i) };
+    const No = struct {
+        fn copy(_: *anyopaque, _: *const Who) bool {
+            return false;
+        }
+    };
+    var unused: u8 = 0;
+    var out: [8]i64 = undefined;
+    try rule.budgets(&whos, .{ .ptr = &unused, .copy = No.copy }, false, &out);
+    // a grant of none probes two; grants 1..7 that took every held draft, plus one; cut to 32 rows over eight blocks
+    try std.testing.expectEqualSlices(i64, &.{ 2, 2, 3, 4, 4, 4, 4, 4 }, &out);
+    try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]);
+    try std.testing.expectEqualSlices(i64, &.{ 2, 2 }, out[0..2]); // two blocks fit whole
+    for (whos[0..2]) |*w| w.grant_short = true;
+    try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]);
+    try std.testing.expectEqualSlices(i64, &.{ 2, 1 }, out[0..2]); // a grant short of the held drafts: no slack
+}
+
+test "a copy ahead: one head draft beside it, or none when the head's drafts cost a pass of their own" {
+    const empty: Table = .{};
+    const Yes = struct {
+        fn copy(_: *anyopaque, _: *const Who) bool {
+            return true;
+        }
+    };
+    var unused: u8 = 0;
+    const probe: Probe = .{ .ptr = &unused, .copy = Yes.copy };
+    var st: ?State = null;
+    defer if (st) |*x| x.deinit(std.testing.allocator);
+    const who: Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 3 };
+    for ([_]bool{ false, true }) |first| {
+        var rule: Rule = .{ .gpa = std.testing.allocator, .k = .{ .copy_first = first }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+        defer rule.deinit();
+        try std.testing.expectEqual(@as(i64, @intFromBool(!first)), try rule.headDepth(&who, null, probe));
+        var out: [1]i64 = undefined;
+        var whos = [_]Who{who};
+        try rule.budgets(&whos, probe, false, &out);
+        try std.testing.expectEqual(@as(i64, @intFromBool(!first)), out[0]);
+    }
+}
+
+test "a granted hold stays only while its stream's expected drafts pay for the verify rows and the hold pass" {
+    const gpa = std.testing.allocator;
+    var shared: Table = .{};
+    defer shared.deinit(gpa);
+    for (1..33) |w| try shared.put(gpa, @intCast(w), 10.0 + @as(f64, @floatFromInt(w)));
+    const empty: Table = .{};
+    const cheap = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 8, .depth = 1, .ms = 4.5 } };
+    var states: [4]?State = @splat(null);
+    defer for (&states) |*st| if (st.*) |*x| x.deinit(gpa);
+    const granted = [_]?u32{ 2, 2, 0, 1 };
+    const gains = [_]?f64{ 1.8, 0.2, null, 0.9 };
+    var whos: [4]Who = undefined;
+    for (&whos, &states, granted, gains) |*w, *st, g, x| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = g, .gain = x };
+    var rule: Rule = .{ .gpa = gpa, .k = .{ .hold_costs = &cheap }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &shared, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+    defer rule.deinit();
+    // 8 ms overhead: plain 4 / 22; stream 0 then 3 raise the rate (5.8 / 25, 6.7 / 26.5), stream 1 would lower it
+    var out = [_]i64{ 4, 4, 0, 3 };
+    try rule.payHolds(&whos, false, &out);
+    try std.testing.expectEqualSlices(i64, &.{ 4, 0, 0, 3 }, &out);
+    const dear = [_]HoldCost{.{ .streams = 1, .depth = 1, .ms = 20.0 }};
+    rule.k.hold_costs = &dear;
+    out = .{ 4, 4, 0, 3 };
+    try rule.payHolds(&whos, false, &out);
+    try std.testing.expectEqualSlices(i64, &.{ 0, 0, 0, 0 }, &out); // no gain pays a 20 ms pass
+    whos[2].granted = 0;
+    out = .{ 4, 4, 2, 3 };
+    try rule.payHolds(&whos, false, &out);
+    try std.testing.expectEqualSlices(i64, &.{ 0, 0, 2, 0 }, &out); // a probe after a grant of none keeps its hold
+}

zig/src/core/lanes/drafted.zig

@@ -103,6 +103,16 @@ pub const Drafted = struct {
         m.draft_streams = d.batched;
         m.draft_probabilities = false;
         m.head_trees = d.trees;
+        m.shared_guard = d.shared_guard;
+        m.copy_first = d.copy_first;
+        m.hold_costs = d.hold_costs;
+        if (d.chances and x.drafter.vtable.chances != null) {
+            // a chain's drafts priced by the drafter's own chances; the next hold as deep as the grant and its slack
+            m.draft_probabilities = d.confidence_floor == null;
+            m.draft_confidence = d.confidence_floor;
+            m.granted_depth = d.confidence_floor == null;
+            m.draft_slack = d.slack;
+        }
         if (d.tree_nodes > 0) {
             // the drafter grows its own trees, best first, and prices their nodes
             m.drafts = @min(d.tree_nodes, if (target.exact_width > 1) target.exact_width - 1 else 0);