From a8461606ac02e3c9865b03c65cdd761fa1a5b6d4 Mon Sep 17 00:00:00 2001
From: chaog992 <338238273+chaog992@users.noreply.github.com>
Date: Fri, 16 Oct 2026 18:02:01 +0800
Subject: [PATCH] fix(lanes): allocate trims a grafted window's held count and
 branch start with its rows

allocate trimmed a window with suffix-match grafts as a plain tree. On drafts the head holds,
sanitizeTree paired the grafts' tokens with the held drafts' parents and the held count stayed,
so the window claimed rows its parents did not cover; and a count inside a copied trunk left
branch_from past the window's last row. It now sanitizes every draft, the held ones as
placeholders, keeps the first min(held, count) of them held, and clears branch_from when the
trimmed tree keeps no graft.
---
 zig/src/core/lanes/fill_test.zig | 117 +++++++++++++++++++++++++++++++
 zig/src/core/lanes/lanes.zig     |   1 +
 zig/src/core/lanes/windows.zig   |  11 ++-
 3 files changed, 127 insertions(+), 2 deletions(-)
 create mode 100644 zig/src/core/lanes/fill_test.zig

diff --git a/zig/src/core/lanes/fill_test.zig b/zig/src/core/lanes/fill_test.zig
new file mode 100644
index 0000000000..57768e3d53
--- /dev/null
+++ b/zig/src/core/lanes/fill_test.zig
@@ -0,0 +1,117 @@
+//! Suffix-match grafts (lanes/fill.zig) on the fake target: a shared round that trims a grafted window stays exact.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Cost = @import("config.zig").Cost;
+const Engine = @import("engine.zig").Engine;
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const win = @import("windows.zig");
+const SuffixLookup = @import("proposer.zig").SuffixLookup;
+
+const gpa = std.testing.allocator;
+
+/// The drafts the grafts follow: the head's, held by the target (no proposer, as the Metal CLI's streams), or copies.
+const Trunk = enum { held, copied };
+
+const Run = struct {
+    replies: [][]u32,
+    reasons: []sm.Reason,
+    drafted: u64, // draft rows verified, all streams
+    branch_rows: u64, // of those, grafted rows
+    branch_kept: u64, // of those, rows on the accepted path
+
+    fn deinit(r: Run) void {
+        for (r.replies) |x| gpa.free(x);
+        gpa.free(r.replies);
+        gpa.free(r.reasons);
+    }
+};
+
+fn run(trunk: Trunk, drafts: bool, n: usize, prompt: []const u32) !Run {
+    // the cost of each row, not the 32-row cap, is what makes the allocator cut the held set's windows
+    var costs: [16]Cost = undefined;
+    for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) };
+    const head = drafts and trunk == .held;
+    // a copied trunk shares 4 rows a round, so the allocator cuts into it
+    const rows: u32 = if (trunk == .held) 32 else 4;
+    var cfg = try Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = head, .speculate = head, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = rows, .max_streams = 8, .draft_streams = true }, 16, 15);
+    defer cfg.deinit(gpa);
+    cfg.fill_lanes = 9; // the Metal CLI's --lanes 8
+    // the reply repeats from its first token, so suffix matches find branches; the head ranks its held chain
+    var target: fake.Fake = .{ .gpa = gpa, .cycle_after = 0, .probe_prompt = prompt.len, .ranked = true };
+    defer target.deinit();
+    var clock: fake.FixedClock = .{};
+    var engine = Engine.init(gpa, &cfg, target.backend(), clock.clock());
+    defer engine.deinit();
+    const streams = try gpa.alloc(sm.Stream, n);
+    defer gpa.free(streams);
+    const proposers = try gpa.alloc(SuffixLookup, n);
+    defer gpa.free(proposers);
+    for (streams, proposers) |*s, *p| {
+        p.* = try SuffixLookup.init(gpa, .{ .min_match = 4 });
+        const proposer = if (trunk == .copied) p.proposer() else null;
+        s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = prompt, .max_new = 120, .eos = &.{96}, .drafts = drafts, .proposer = proposer });
+    }
+    defer for (streams, proposers) |*s, *p| {
+        s.deinit(gpa);
+        p.deinit();
+    };
+    for (streams) |*s| try engine.addStream(s);
+    while (engine.activeCount() > 0) try engine.step();
+    var out: Run = .{ .replies = try gpa.alloc([]u32, n), .reasons = try gpa.alloc(sm.Reason, n), .drafted = 0, .branch_rows = 0, .branch_kept = 0 };
+    for (out.replies, out.reasons, streams) |*r, *why, *s| {
+        r.* = try gpa.dupe(u32, s.emitted());
+        why.* = s.reason;
+        out.drafted += s.drafted;
+        out.branch_rows += s.branch_rows;
+        out.branch_kept += s.branch_accepted;
+    }
+    return out;
+}
+
+test "a shared round that trims a window with suffix-match grafts replies as the one-token decode" {
+    const prompt = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 };
+    for ([_]Trunk{ .held, .copied }) |trunk| {
+        const want = try run(trunk, false, 1, &prompt);
+        defer want.deinit();
+        const got = try run(trunk, true, 2, &prompt);
+        defer got.deinit();
+        try std.testing.expect(got.branch_kept > 0 and got.branch_kept <= got.branch_rows and got.branch_rows <= got.drafted);
+        for (got.replies, got.reasons) |r, why| {
+            try std.testing.expectEqualSlices(u32, want.replies[0], r);
+            try std.testing.expectEqual(want.reasons[0], why);
+        }
+    }
+}
+
+test "allocate keeps a grafted window's first drafts at every count, the trunk's before the grafts" {
+    // flat costs, so a round fills its batch rows and the window's count is batch_rows - 1
+    var costs: [16]Cost = undefined;
+    for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 };
+    const Window = struct { kind: win.Kind, held: u32, tokens: []const u32, parents: []const i32 };
+    const windows = [_]Window{
+        // three held drafts, then grafts: one after the last held draft, one off the first held draft and its child
+        .{ .kind = .head, .held = 3, .tokens = &.{ 7, 8, 9 }, .parents = &.{ -1, 0, 1, 2, 0, 4 } },
+        // three copied drafts, then grafts past their end: one after the last, its child and that child's sibling
+        .{ .kind = .copy, .held = 0, .tokens = &.{ 4, 1, 5, 7, 8, 9 }, .parents = &.{ -1, 0, 1, 2, 3, 3 } },
+    };
+    for (windows) |t| for (1..1 + t.parents.len) |c| {
+        var cfg = try Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = @intCast(c + 1), .max_streams = 8, .draft_streams = true }, 16, 15);
+        defer cfg.deinit(gpa);
+        var target: fake.Fake = .{ .gpa = gpa };
+        defer target.deinit();
+        var clock: fake.FixedClock = .{};
+        var engine = Engine.init(gpa, &cfg, target.backend(), clock.clock());
+        defer engine.deinit();
+        var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &.{ 3, 1, 4 }, .max_new = 64, .drafts = true });
+        defer s.deinit(gpa);
+        var plans = [_]win.Plan{.{ .stream = &s, .position = 3, .kind = t.kind, .held = t.held, .tokens = t.tokens, .parents = t.parents, .branch_from = 4 }};
+        try win.allocate(&engine, &plans);
+        const p = plans[0];
+        try std.testing.expectEqual(c, p.count());
+        try std.testing.expectEqual(@min(t.held, c), p.held);
+        try std.testing.expectEqualSlices(u32, t.tokens[0 .. c - @min(t.held, c)], p.tokens);
+        try std.testing.expectEqualSlices(i32, t.parents[0..c], p.parents.?);
+        try std.testing.expectEqual(@as(u32, if (c > 3) 4 else 0), p.branch_from);
+    };
+}
diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig
index e943bf9d46..48c1c61242 100644
--- a/zig/src/core/lanes/lanes.zig
+++ b/zig/src/core/lanes/lanes.zig
@@ -45,6 +45,7 @@ test {
     _ = @import("gpu_full_test.zig");
     _ = @import("drafted_test.zig");
     _ = @import("pricing_test.zig");
+    _ = @import("fill_test.zig");
 }
 
 pub const learned_disk = @import("learned_disk.zig");
diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig
index 36200960d2..11a029d171 100644
--- a/zig/src/core/lanes/windows.zig
+++ b/zig/src/core/lanes/windows.zig
@@ -230,9 +230,16 @@ pub fn allocate(e: *Engine, plans: []Plan) !void {
             p.tokens = &.{};
             p.parents = null;
         } else if (p.parents) |parents| {
-            const tree = try accept.sanitizeTree(a, p.tokens, parents, c);
-            p.tokens = tree.tokens;
+            // the held drafts lead the window (placeholders here), so the first kept rows stay held
+            const drafts = try a.alloc(u32, p.count());
+            @memset(drafts[0..p.held], 0);
+            @memcpy(drafts[p.held..], p.tokens);
+            const tree = try accept.sanitizeTree(a, drafts, parents, c);
+            const held = @min(p.held, c);
+            p.held = held;
+            p.tokens = tree.tokens[held..];
             p.parents = tree.parents;
+            if (c < p.branch_from) p.branch_from = 0; // the grafts follow the trunk's drafts: this count keeps none
         } else if (p.held > 0) {
             p.held = c;
         } else p.tokens = p.tokens[0..c];