From 3ced86d9242e2b7a5bbb5feb5560739c679bc77b Mon Sep 17 00:00:00 2001
From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com>
Date: Mon, 19 Oct 2026 14:53:53 +0200
Subject: [PATCH 1/4] Flash Next CUDA: drafted windows of several streams
 verified in one shared pass, each row's bits its lone pass's

---
 zig/src/families/flashnext/cuda_batch.zig     | 66 ++++++++++--------
 .../families/flashnext/cuda_batch_head.zig    | 10 +++
 .../families/flashnext/cuda_batch_state.zig   | 34 ++++++++--
 .../families/flashnext/cuda_batch_train.zig   |  2 +-
 .../families/flashnext/cuda_images_test.zig   |  2 +-
 zig/src/families/flashnext/cuda_lanes.zig     | 24 ++++---
 .../families/flashnext/cuda_lanes_load.zig    |  3 +-
 .../families/flashnext/cuda_lanes_verify.zig  | 34 ++++++++++
 zig/src/families/flashnext/cuda_media.zig     |  8 ++-
 .../families/flashnext/cuda_media_test.zig    |  2 +-
 zig/src/families/flashnext/cuda_mtp_ops.zig   | 32 +++++++++
 zig/src/families/flashnext/cuda_prompt.zig    |  8 ++-
 .../flashnext/cuda_prompt_windows.zig         | 67 +++++++++++++++++++
 .../families/flashnext/cuda_round_graph.zig   |  2 +
 14 files changed, 246 insertions(+), 48 deletions(-)
 create mode 100644 zig/src/families/flashnext/cuda_lanes_verify.zig
 create mode 100644 zig/src/families/flashnext/cuda_prompt_windows.zig

diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig
index 5e79eee832..272ea9a4df 100644
--- a/zig/src/families/flashnext/cuda_batch.zig
+++ b/zig/src/families/flashnext/cuda_batch.zig
@@ -267,31 +267,43 @@ pub fn Pass(comptime Tri: type) type {
                 try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide);
                 return;
             }
-            // a round: each row's convolution over its own stream's tail, the one-row launch a lone decode makes
-            for (self.lanes, 0..) |l, r| {
-                const g = r * wide * 2;
-                try self.tri.pleConv(s.ple_gated + g, s.ple_pss + r * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, 1, dims, streams, 4, 3);
-                try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, 1, ple_tail_n, wide);
+            // shared rounds: each lane's rows convolved over its own stream's tail, the launch its lone pass makes
+            var row0: usize = 0;
+            for (self.lanes) |l| {
+                const g = row0 * wide * 2;
+                try self.tri.pleConv(s.ple_gated + g, s.ple_pss + row0 * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, l.rows, dims, streams, 4, 3);
+                try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, l.rows, ple_tail_n, wide);
+                row0 += l.rows;
             }
         }
 
-        pub fn gdnBlock(self: *Self, lw: *const Layer, li: usize, keep: LayerKeep) !void {
+        pub fn gdnBlock(self: *Self, lw: *const Layer, layer: usize, li: usize, keep: LayerKeep) !void {
             const s = self.s;
             const n = self.rows;
             try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1);
-            try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2);
-            try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2);
             // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits
             if (self.lanes.len == 0) {
+                try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2);
+                try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2);
                 try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb);
                 try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys);
-            } else for (self.lanes, 0..) |l, r| {
-                // a round: each row's front and recurrence on its own stream's conv window and state
-                const qk = r * gdn.nk * gdn.dk * 4;
-                const v = r * gdn.value_dim * 4;
-                const gb = r * gdn.nv * 2 * 4;
-                try self.gdnFront(lw, s.proj + r * gdn.proj_width * 2, l.seq.conv[li].ptr, 1, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb);
-                try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, 0, 1, s.ys + v);
+            } else {
+                // shared rounds: each lane's rows kept, fronted and chained on its own conv window and state
+                var row0: usize = 0;
+                for (self.lanes) |l| {
+                    const lk: LayerKeep = if (l.keep) |kp| kp.layers[layer] else .{};
+                    const proj = s.proj + row0 * gdn.proj_width * 2;
+                    const qk = row0 * gdn.nk * gdn.dk * 4;
+                    const v = row0 * gdn.value_dim * 4;
+                    const gb = row0 * gdn.nv * 2 * 4;
+                    try self.copy(lk.proj, proj, l.rows * gdn.proj_width * 2);
+                    try self.copy(lk.hist, l.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2);
+                    try self.gdnFront(lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb);
+                    try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, lk.state_out, l.rows, s.ys + v);
+                    // the lane's conv window moves on now: the gated norm below reads none of it
+                    try self.shiftWindow(l.seq.conv[li].ptr, proj, gdn.proj_width, l.rows, gdn.taps - 1, gdn.conv_dim);
+                    row0 += l.rows;
+                }
             }
             var b: cuda.Args = .{};
             b.add(s.ys);
@@ -301,11 +313,7 @@ pub fn Pass(comptime Tri: type) type {
             b.add(s.gout);
             b.add(s.gxs);
             try cuda.launch.launch(self.k.gnorm, .{ .grid = .{ .x = @intCast(n), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, self.stream, &b);
-            if (self.lanes.len == 0) {
-                try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim);
-            } else for (self.lanes, 0..) |l, r| {
-                try self.shiftWindow(l.seq.conv[li].ptr, s.proj + r * gdn.proj_width * 2, gdn.proj_width, 1, gdn.taps - 1, gdn.conv_dim);
-            }
+            if (self.lanes.len == 0) try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim);
             try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk);
         }
 
@@ -354,17 +362,19 @@ pub fn Pass(comptime Tri: type) type {
             try self.copyCols(pa + qw2 * 2, proj_n, keep.pk, kw, kw, n);
             try self.copyCols(pa + (qw2 + kw) * 2, proj_n, keep.pv, kw, kw, n);
             if (self.lanes.len > 0) {
-                // a round: each row's rope, cache writes, selection and attention on its own stream's caches
+                // shared rounds: each lane's rope, cache writes, selection and attention on its own stream's caches
                 const own_seq = self.seq;
                 const own_pos0 = self.pos0;
                 defer {
                     self.seq = own_seq;
                     self.pos0 = own_pos0;
                 }
-                for (self.lanes, 0..) |l, r| {
+                var row0: usize = 0;
+                for (self.lanes, 0..) |l, i| {
                     self.seq = l.seq;
                     self.pos0 = l.pos;
-                    try self.attnRows(lw, ai, pa + r * proj_n * 2, 1, .{ .pos = s.pos + r * 16, .q = s.q + r * q_heads * head_dim * 2, .iq = s.iq + r * index_heads * index_dim * 2, .ao = s.ao + r * out_k * 2 });
+                    try self.attnRows(lw, ai, pa + row0 * proj_n * 2, l.rows, .{ .pos = s.pos + i * 16, .q = s.q + row0 * q_heads * head_dim * 2, .iq = s.iq + row0 * index_heads * index_dim * 2, .ao = s.ao + row0 * out_k * 2 });
+                    row0 += l.rows;
                 }
             } else try self.attnRows(lw, ai, pa, n, .{ .pos = s.pos, .q = s.q, .iq = s.iq, .ao = s.ao });
             const gated = if (keep.gated != 0) keep.gated else s.gated;
@@ -470,7 +480,7 @@ pub fn Pass(comptime Tri: type) type {
             const id32 = s.hostIds(n);
             for (id32, ids) |*o, v| o.* = @intCast(v);
             try self.upload(s.ids, std.mem.sliceAsBytes(id32));
-            const blocks = if (self.lanes.len > 0) n else cdiv(n, att_rows);
+            const blocks = if (self.lanes.len > 0) self.lanes.len else cdiv(n, att_rows);
             // one int a block (a round: a row), 16 bytes apart: _chunks and _merge take POS0 16-byte aligned
             const pos = s.hostPos(blocks * 4);
             @memset(pos, 0);
@@ -490,9 +500,9 @@ pub fn Pass(comptime Tri: type) type {
             };
             if (n == 0 or n > s.cap) return error.UnexpectedTensor;
             if (self.lanes.len > 0) {
-                // a round reads every row's scores and keeps nothing (a pocket's adds are row by row)
-                if (n != self.lanes.len or self.hooks.logits != .all or self.hooks.keep != null or self.hooks.streams != 0 or self.hooks.moe_add != null) return error.UnexpectedTensor;
-                for (self.lanes) |l| if (l.pos + 1 > l.seq.cap) return error.UnexpectedTensor;
+                // shared rounds read every row's scores; a lane keeps its own rollback and hidden rows, nothing else
+                if (self.hooks.logits != .all or self.hooks.keep != null or self.hooks.streams != 0 or self.hooks.moe_add != null) return error.UnexpectedTensor;
+                try parts.checkLanes(self.lanes, n);
             } else if (pos0 + n > self.seq.cap) return error.UnexpectedTensor;
             self.rows = n;
             self.pos0 = pos0;
@@ -528,7 +538,7 @@ pub fn Pass(comptime Tri: type) type {
                 try self.writeback(0, n);
                 try self.readoutKeep(lw.attn_hc, 0, n, keep.attn_hc);
                 if (lw.linear) {
-                    try self.gdnBlock(lw, wt.linearIndex(layer), keep);
+                    try self.gdnBlock(lw, layer, wt.linearIndex(layer), keep);
                 } else {
                     try self.attnBlockKeep(lw, wt.attnIndex(layer), keep);
                 }
diff --git a/zig/src/families/flashnext/cuda_batch_head.zig b/zig/src/families/flashnext/cuda_batch_head.zig
index 9b226226dc..23594a4f33 100644
--- a/zig/src/families/flashnext/cuda_batch_head.zig
+++ b/zig/src/families/flashnext/cuda_batch_head.zig
@@ -24,9 +24,19 @@ pub fn head(self: anytype, logits_host: ?[]u8) !void {
         .all => {
             try self.writeback(0, n);
             try self.copy(self.hooks.streams, s.h, n * wide * 2);
+            var row0: usize = 0;
+            for (self.lanes) |l| {
+                try self.copy(l.streams, s.h + row0 * wide * 2, l.rows * wide * 2);
+                row0 += l.rows;
+            }
             try self.readoutKeep(self.m.mixer, 0, n, mixer_keep);
             const out = if (self.hooks.logits_out != 0) self.hooks.logits_out else return error.UnexpectedTensor;
             try self.projectHead(out, n);
+            row0 = 0;
+            for (self.lanes) |l| {
+                try self.copy(l.logits, out + row0 * vocab * 2, l.rows * vocab * 2);
+                row0 += l.rows;
+            }
             if (logits_host) |lh| {
                 try self.complete();
                 const d = self.stream.d;
diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig
index dc1749c657..7aa01b7cb1 100644
--- a/zig/src/families/flashnext/cuda_batch_state.zig
+++ b/zig/src/families/flashnext/cuda_batch_state.zig
@@ -509,8 +509,32 @@ pub fn scoreBytes(cap: usize, rows: usize) usize {
     return @min(rows, att_rows) * scoreStride(cap) * 4;
 }
 
-/// One row of a decode round: its stream's next token, on that stream's own caches at its own position.
-pub const Lane = struct { seq: SeqState, pos: usize };
+/// One stream's rows in a shared pass from `pos` on its own caches, and its window's rollback, hidden rows and scores.
+pub const Lane = struct { seq: SeqState, pos: usize, rows: usize = 1, keep: ?*const Keep = null, streams: u64 = 0, logits: u64 = 0 };
+
+/// Lanes whose rows add up to `n`, each window within its cache and the attention block, keeping a rollback at most.
+pub fn checkLanes(lanes: []const Lane, n: usize) error{UnexpectedTensor}!void {
+    var total: usize = 0;
+    for (lanes) |l| {
+        if (l.rows == 0 or l.rows > att_rows or l.pos + l.rows > l.seq.cap) return error.UnexpectedTensor;
+        if (l.keep) |kp| if (!rollbackOnly(kp)) return error.UnexpectedTensor;
+        total += l.rows;
+    }
+    if (total != n) return error.UnexpectedTensor;
+}
+
+/// A drafted window's rollback (the DeltaNet rows, history and end states alone), not a training pass's tensors.
+pub fn rollbackOnly(k: *const Keep) bool {
+    if (!std.meta.eql(k.mixer, HcKeep{})) return false;
+    for (k.layers) |lk| {
+        var rest = lk;
+        rest.proj = 0;
+        rest.hist = 0;
+        rest.state_out = 0;
+        if (!std.meta.eql(rest, LayerKeep{})) return false;
+    }
+    return true;
+}
 
 /// The engine's resident weights the pass reads.
 pub const Model = struct {
@@ -549,11 +573,13 @@ pub const Gather = struct {
             };
             return;
         }
-        for (g.lanes, 0..) |l, r| {
-            g.ngram.rows(l.seq.history[0..], g.ids[r..][0..1], g.emb[r * dims ..][0..dims], g.xs[r * (dims / 32) ..][0 .. dims / 32]) catch |err| {
+        var row0: usize = 0;
+        for (g.lanes) |l| {
+            g.ngram.rows(l.seq.history[0..], g.ids[row0..][0..l.rows], g.emb[row0 * dims ..][0 .. l.rows * dims], g.xs[row0 * (dims / 32) ..][0 .. l.rows * (dims / 32)]) catch |err| {
                 g.failed = err;
                 return;
             };
+            row0 += l.rows;
         }
     }
 };
diff --git a/zig/src/families/flashnext/cuda_batch_train.zig b/zig/src/families/flashnext/cuda_batch_train.zig
index 20747b763c..5e6daf45bc 100644
--- a/zig/src/families/flashnext/cuda_batch_train.zig
+++ b/zig/src/families/flashnext/cuda_batch_train.zig
@@ -88,7 +88,7 @@ pub fn spans(comptime Tri: type, p: *batch.Pass(Tri), ids: []const u32, list: []
                 const li = wt.linearIndex(layer);
                 try zero(p, p.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim / 2);
                 try zero(p, p.seq.state[li].ptr, state_n);
-                try p.gdnBlock(lw, li, k);
+                try p.gdnBlock(lw, layer, li, k);
             } else {
                 try p.attnBlockKeep(lw, wt.attnIndex(layer), k);
             }
diff --git a/zig/src/families/flashnext/cuda_images_test.zig b/zig/src/families/flashnext/cuda_images_test.zig
index 38fe41a616..624b90d982 100644
--- a/zig/src/families/flashnext/cuda_images_test.zig
+++ b/zig/src/families/flashnext/cuda_images_test.zig
@@ -92,7 +92,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, draf
     try view.image.inject(stream, 0x123400, tokens.len, 1, 4);
     try std.testing.expectEqual(@as(usize, 1), Fake.launches);
     const media = @import("cuda_media.zig");
-    const Round = struct { seq: struct { media: ?media.View }, pos: usize };
+    const Round = struct { seq: struct { media: ?media.View }, pos: usize, rows: usize = 1 };
     const rows = [_]Round{ .{ .seq = .{ .media = view }, .pos = tokens.len }, .{ .seq = .{ .media = view }, .pos = tokens.len + 1 } };
     try media.inject(view, &rows, stream, 0x123400, 0, 2);
     try std.testing.expectEqual(@as(usize, 1), Fake.launches);
diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig
index 226c988b55..7d036518d2 100644
--- a/zig/src/families/flashnext/cuda_lanes.zig
+++ b/zig/src/families/flashnext/cuda_lanes.zig
@@ -18,6 +18,8 @@ const be = lanes.backend;
 const ring = 1024;
 /// Streams a round takes at most (a round with more live streams takes them in turns).
 const round_streams = 32;
+/// Drafted windows' rows a shared verify takes at most: a drafting round's (TF_FLASHNEXT_WINDOW_ROWS; 0: each alone).
+const window_rows_default = 4 * mtp.max_rows;
 
 /// The Triton entry points the pass launches and the head's two more: a captured set replaces our glue with all.
 pub const captured_entries = [_][]const u8{ "_embed", "_hc_writeback", "_qmm_hcdown", "_hc_reduce_act", "_qmm_upmix", "_ple_gate", "_ple_conv", "_attn_prep", "_attn_gate", "_pool", "_scores", "_chunks", "_merge", "_router", "_topk_rows" };
@@ -174,7 +176,7 @@ pub fn Namespace(comptime Tri: type) type {
                 return .{ .set = if (self.set) |*x| x else null, .own = &self.glue, .s = self.stream };
             }
 
-            fn ops(self: *Cuda) Ops {
+            pub fn ops(self: *Cuda) Ops {
                 return .{ .e = self.engine, .tri = self.tri(), .h = self.head.?, .main = if (self.ranks != null) .{
                     .ptr = self,
                     .forward = mainForward,
@@ -255,7 +257,7 @@ pub fn Namespace(comptime Tri: type) type {
                 return mtp.drawRow(self.gpa, row, null, position, s.sampling);
             }
 
-            fn lane(self: *Cuda, s: *lanes.Stream) !*Lane {
+            pub fn lane(self: *Cuda, s: *lanes.Stream) !*Lane {
                 return self.lanes.getPtr(s) orelse return error.UnknownStream;
             }
 
@@ -429,19 +431,23 @@ pub fn Namespace(comptime Tri: type) type {
             defer self.gpa.free(ids);
             const at_window = try self.gpa.alloc(usize, windows.len);
             defer self.gpa.free(at_window);
+            const row_ids = try self.gpa.alloc([mtp.max_rows]u32, windows.len);
+            defer self.gpa.free(row_ids);
+            const drafted = try self.gpa.alloc(usize, windows.len);
+            defer self.gpa.free(drafted);
             var plain: usize = 0;
+            var held: usize = 0;
             for (windows, out, 0..) |w, *o, wi| {
                 if (w.held != 0 or w.parents != null) return error.HeldDraftsNotBuilt;
                 const l = try self.lane(w.stream);
                 const rows = w.rows();
-                if (l.draft) |*d| {
+                if (l.draft != null) {
                     if (rows > mtp.max_rows) return error.WindowTooWide;
                     for (w.positions, 0..) |p, r| if (p != l.seq.fed + 1 + r) return error.PositionMismatch;
-                    var row_ids: [mtp.max_rows]u32 = undefined;
-                    row_ids[0] = w.pending;
-                    @memcpy(row_ids[1..rows], w.tokens);
-                    try self.ops().verify(&l.seq, d, row_ids[0..rows]);
-                    try self.ops().draws(d, w.stream.sampling, o.sampled[0..rows]);
+                    row_ids[wi][0] = w.pending;
+                    @memcpy(row_ids[wi][1..rows], w.tokens);
+                    drafted[held] = wi;
+                    held += 1;
                 } else {
                     if (rows != 1) return error.CannotRewind;
                     if (w.positions[0] != l.seq.fed + 1) return error.PositionMismatch;
@@ -452,6 +458,8 @@ pub fn Namespace(comptime Tri: type) type {
                 }
                 if (w.tokens.len > 0) @memcpy(o.drafts[0..w.tokens.len], w.tokens);
             }
+            const cap = @min(@min(envUsize("TF_FLASHNEXT_WINDOW_ROWS", window_rows_default), window_rows_default), self.engine.chunk);
+            try @import("cuda_lanes_verify.zig").windows(self, windows, out, drafted[0..held], row_ids, cap);
             var at: usize = 0;
             while (at < plain) : (at += round_streams) {
                 const n = @min(round_streams, plain - at);
diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig
index 69c300ef69..f9437c242f 100644
--- a/zig/src/families/flashnext/cuda_lanes_load.zig
+++ b/zig/src/families/flashnext/cuda_lanes_load.zig
@@ -69,7 +69,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu
     const window = @min(o.context + mtp.max_rows, prompt.max_tokens);
     // and an attention block's block scores at the window (64 MiB at 262,144 tokens)
     _ = try own.engine.reserveScores(window, @min(o.context, own.engine.chunk));
-    own.logits = try cuda.DeviceBuffer.alloc(driver, round_streams * prompt.vocabulary * 2);
+    // a shared round's or a shared verify's scores: up to a drafting round's rows (facts' batch_rows)
+    own.logits = try cuda.DeviceBuffer.alloc(driver, @max(round_streams, 4 * mtp.max_rows) * prompt.vocabulary * 2);
     errdefer own.logits.free();
     own.lanes = .empty;
     own.next = 0;
diff --git a/zig/src/families/flashnext/cuda_lanes_verify.zig b/zig/src/families/flashnext/cuda_lanes_verify.zig
new file mode 100644
index 0000000000..0b4d25ab5d
--- /dev/null
+++ b/zig/src/families/flashnext/cuda_lanes_verify.zig
@@ -0,0 +1,34 @@
+//! A round's drafted windows grouped into shared verify passes (cuda_lanes.zig's verifyFn).
+const lanes = @import("lanes");
+const prompt = @import("cuda_prompt.zig");
+const mtp = @import("cuda_mtp.zig");
+const be = lanes.backend;
+
+/// Windows a shared pass takes at most (a round's streams).
+const round_streams = 32;
+
+/// The drafted windows verified in shared passes of up to `cap` rows (alone on two ranks), then drawn.
+pub fn windows(self: anytype, wins: []const be.Window, out: []be.Verified, drafted: []const usize, row_ids: []const [mtp.max_rows]u32, cap: usize) !void {
+    var seqs: [round_streams]*prompt.Seq = undefined;
+    var ds: [round_streams]*mtp.Draft = undefined;
+    var idss: [round_streams][]const u32 = undefined;
+    var at: usize = 0;
+    while (at < drafted.len) {
+        var n: usize = 0;
+        var rows: usize = 0;
+        while (at + n < drafted.len and n < round_streams) : (n += 1) {
+            const w = wins[drafted[at + n]];
+            if (n > 0 and rows + w.rows() > cap) break;
+            const l = try self.lane(w.stream);
+            seqs[n] = &l.seq;
+            ds[n] = &l.draft.?;
+            idss[n] = row_ids[drafted[at + n]][0..w.rows()];
+            rows += w.rows();
+        }
+        if (n == 1 or self.ranks != null or cap <= 1) {
+            for (seqs[0..n], ds[0..n], idss[0..n]) |q, d, ids| try self.ops().verify(q, d, ids);
+        } else try self.ops().verifyMany(seqs[0..n], ds[0..n], idss[0..n], self.logits.ptr);
+        for (drafted[at..][0..n], ds[0..n]) |wi, d| try self.ops().draws(d, wins[wi].stream.sampling, out[wi].sampled[0..wins[wi].rows()]);
+        at += n;
+    }
+}
diff --git a/zig/src/families/flashnext/cuda_media.zig b/zig/src/families/flashnext/cuda_media.zig
index 6e66da20a0..395c492516 100644
--- a/zig/src/families/flashnext/cuda_media.zig
+++ b/zig/src/families/flashnext/cuda_media.zig
@@ -42,8 +42,12 @@ pub const View = struct {
 pub fn inject(view: ?View, rows: anytype, stream: cuda.Stream, out: u64, pos: usize, n: usize) !void {
     if (rows.len == 0) {
         if (view) |m| try m.image.inject(stream, out, pos, n, wt.streams);
-    } else for (rows, 0..) |l, r| {
-        if (l.seq.media) |m| try m.image.inject(stream, out + r * wt.wide * 2, l.pos, 1, wt.streams);
+    } else {
+        var row0: usize = 0;
+        for (rows) |l| {
+            if (l.seq.media) |m| try m.image.inject(stream, out + row0 * wt.wide * 2, l.pos, l.rows, wt.streams);
+            row0 += l.rows;
+        }
     }
 }
 
diff --git a/zig/src/families/flashnext/cuda_media_test.zig b/zig/src/families/flashnext/cuda_media_test.zig
index 0f58af79ba..ec92b66351 100644
--- a/zig/src/families/flashnext/cuda_media_test.zig
+++ b/zig/src/families/flashnext/cuda_media_test.zig
@@ -40,7 +40,7 @@ test "the splice, attention and pooled indexer launches get every pointer, scala
     try m.pooled(s, 1, 2, 3, 4, 5, 1e-6, 8);
     try std.testing.expectEqualSlices(u64, &.{ 1, 2, 3, 4, 5, eps, 8, wt.index_dim, wt.half, wt.index_ratio, 3000, 20, delta }, Capture.values[2][0..13]);
     try std.testing.expectEqual([7]u32{ 8 / wt.index_ratio + 2, 1, 1, 128, 1, 1, 0 }, Capture.configs[2]);
-    const Row = struct { seq: struct { media: ?media.View }, pos: usize };
+    const Row = struct { seq: struct { media: ?media.View }, pos: usize, rows: usize = 1 };
     const rows = [_]Row{ .{ .seq = .{ .media = m }, .pos = 2 }, .{ .seq = .{ .media = null }, .pos = 2 }, .{ .seq = .{ .media = view(&d, 1100) }, .pos = 5 } };
     try media.inject(null, &rows, s, 9000, 0, 3);
     try std.testing.expectEqualSlices(u64, &.{ 1000, 2000, 9000, 2, 20, 2560, 4 }, Capture.values[3][0..7]);
diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig
index cfcd9fdd05..3adf0e7169 100644
--- a/zig/src/families/flashnext/cuda_mtp_ops.zig
+++ b/zig/src/families/flashnext/cuda_mtp_ops.zig
@@ -245,6 +245,38 @@ pub fn Ops(comptime Tri: type) type {
             d.open = keeps;
         }
 
+        /// Several streams' windows in one forward, each left as its `verify` leaves it; `out` takes all rows' scores.
+        pub fn verifyMany(o: Self, seqs: []const *prompt.Seq, ds: []const *Draft, idss: []const []const u32, out: u64) !void {
+            if (o.main != null) return error.UnsupportedRankOperation;
+            // every window refused before any stream's state moves
+            for (seqs, idss) |seq, ids| {
+                if (ids.len == 0 or ids.len > max_rows) return error.WindowTooWide;
+                if (seq.fed + ids.len > seq.cap) return error.PromptTooLong;
+            }
+            o.e.round_graph.clear();
+            const wins = try o.e.gpa.alloc(prompt.Engine.Window, seqs.len);
+            defer o.e.gpa.free(wins);
+            for (seqs, ds, idss, wins) |seq, d, ids, *w| {
+                const n = ids.len;
+                d.settle(seq);
+                d.fed0 = seq.fed;
+                d.rows = n;
+                d.history = seq.history;
+                @memcpy(d.ids[0..n], ids);
+                const keeps = n > 1;
+                if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2);
+                w.* = .{ .seq = seq, .ids = ids, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .logits = d.logits.ptr };
+            }
+            try o.e.windows(Tri, o.tri, wins, out);
+            var row0: usize = 0;
+            for (ds, idss) |d, ids| {
+                const keeps = ids.len > 1;
+                if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow + row0 * wide * 2, ids.len * wide * 2);
+                d.open = keeps;
+                row0 += ids.len;
+            }
+        }
+
         /// The window's rows drawn at their keyed positions (seq position fed0 + 1 + r), as serial decoding draws.
         pub fn draws(o: Self, d: *Draft, sm: ?Sampling, out: []u32) !void {
             const n = d.rows;
diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig
index 0ff498c22d..cedd20f2e3 100644
--- a/zig/src/families/flashnext/cuda_prompt.zig
+++ b/zig/src/families/flashnext/cuda_prompt.zig
@@ -279,7 +279,7 @@ pub const Engine = struct {
     }
 
     /// The weights a pass of streams that read the model as shipped reads (untaught): the pockets left out.
-    fn shipped(e: *Engine) batch.Model {
+    pub fn shipped(e: *Engine) batch.Model {
         var m = e.model();
         m.pocket = null;
         return m;
@@ -326,7 +326,7 @@ pub const Engine = struct {
         return e.forward(Tri, tri, seq, ids, .{});
     }
 
-    fn complete(e: *Engine) !void {
+    pub fn complete(e: *Engine) !void {
         if (e.collective) |c| return c.finish(e.stream.*);
         try e.stream.synchronize();
     }
@@ -388,6 +388,10 @@ pub const Engine = struct {
         seq.fed += ids.len;
     }
 
+    /// A stream's window in a shared pass and the pass over several (cuda_prompt_windows.zig).
+    pub const Window = @import("cuda_prompt_windows.zig").Window;
+    pub const windows = @import("cuda_prompt_windows.zig").windows;
+
     /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone.
     pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void {
         if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor;
diff --git a/zig/src/families/flashnext/cuda_prompt_windows.zig b/zig/src/families/flashnext/cuda_prompt_windows.zig
new file mode 100644
index 0000000000..0f04a3f590
--- /dev/null
+++ b/zig/src/families/flashnext/cuda_prompt_windows.zig
@@ -0,0 +1,67 @@
+//! Several streams' drafted windows in one batched pass (Engine.windows), each row's bits its lone pass's.
+const std = @import("std");
+const prompt = @import("cuda_prompt.zig");
+const batch = prompt.batch;
+const pocket = @import("cuda_pocket.zig");
+const vocab = @import("cuda_prompt_wt.zig").vocab;
+const Engine = prompt.Engine;
+const Seq = prompt.Seq;
+
+/// A stream's window in a shared pass: its rows, and where its rollback rows, hidden rows and scores go (0: none).
+pub const Window = struct { seq: *Seq, ids: []const u32, keep: ?*const batch.Keep = null, streams: u64 = 0, logits: u64 = 0 };
+
+/// Several streams' windows in one batched pass, each row's bits its lone pass's; `out` holds every row's scores.
+pub fn windows(e: *Engine, comptime Tri: type, tri: Tri, wins: []const Window, out: u64) !void {
+    if (wins.len == 0) return error.UnexpectedTensor;
+    var total: usize = 0;
+    var cap: usize = 0;
+    for (wins) |w| {
+        if (w.ids.len == 0 or w.seq.fed + w.ids.len > w.seq.cap) return error.UnexpectedTensor;
+        total += w.ids.len;
+        cap = @max(cap, w.seq.cap);
+    }
+    if (total > e.chunk) return error.UnexpectedTensor;
+    const lanes = try e.gpa.alloc(batch.Lane, wins.len);
+    defer e.gpa.free(lanes);
+    const ids = try e.gpa.alloc(u32, total);
+    defer e.gpa.free(ids);
+    // an untaught stream's rows read no pocket, as in a round
+    const untaught = try e.gpa.alloc(bool, total);
+    defer e.gpa.free(untaught);
+    const scales = try e.gpa.alloc(f32, total);
+    defer e.gpa.free(scales);
+    const s = try e.reserve(total);
+    const scores = try e.reserveScores(cap, s.cap);
+    var at: usize = 0;
+    for (wins, lanes) |w, *l| {
+        l.* = .{ .seq = w.seq.view(scores), .pos = w.seq.fed, .rows = w.ids.len, .keep = w.keep, .streams = w.streams, .logits = w.logits };
+        @memcpy(ids[at..][0..w.ids.len], w.ids);
+        @memset(untaught[at..][0..w.ids.len], w.seq.untaught);
+        at += w.ids.len;
+    }
+    const reads = pocket.reads(untaught, scales);
+    defer e.taught_rows = 0;
+    if (reads == .some and e.pockets.mask() != 0) e.taught_rows = try e.pockets.rowScales(e.driver, e.stream, scales);
+    var p: batch.Pass(Tri) = .{
+        .gpa = e.gpa,
+        .io = e.io,
+        .stream = e.stream.*,
+        .tri = tri,
+        .k = &e.batch_k,
+        .s = s,
+        .m = if (reads == .none) e.shipped() else e.model(),
+        .seq = lanes[0].seq,
+        .lanes = lanes,
+        .hooks = .{ .logits = .all, .logits_out = out },
+    };
+    try p.chunk(ids, 0, null);
+    try e.complete();
+    // each stream's last row in its host scores, as its lone pass leaves them
+    const d = e.stream.d;
+    at = 0;
+    for (wins) |w| {
+        at += w.ids.len;
+        try d.check(d.api.cuMemcpyDtoH_v2(w.seq.logits.ptr, out + (at - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH");
+        w.seq.fed += w.ids.len;
+    }
+}
diff --git a/zig/src/families/flashnext/cuda_round_graph.zig b/zig/src/families/flashnext/cuda_round_graph.zig
index e367e31dd9..16977d9833 100644
--- a/zig/src/families/flashnext/cuda_round_graph.zig
+++ b/zig/src/families/flashnext/cuda_round_graph.zig
@@ -43,6 +43,8 @@ pub const Cache = struct {
     /// CPU histories and pinned inputs advance once per round, outside stream capture and replay.
     fn prepare(pass: anytype, ids: []const u32) !void {
         const s = pass.s;
+        // a captured round is one row a lane; windows run the plain pass
+        if (ids.len != pass.lanes.len) return error.UnexpectedTensor;
         try s.ready();
         for (s.hostIds(ids.len), ids) |*out, token| out.* = @intCast(token);
         const positions = s.hostPos(ids.len * 4);

From 71fc0971f069c1e94644814f0190ea762df53c85 Mon Sep 17 00:00:00 2001
From: Mirrdhyn <21200274+Mirrdhyn@users.noreply.github.com>
Date: Mon, 19 Oct 2026 14:53:53 +0200
Subject: [PATCH 2/4] Flash Next CUDA tests: several streams' windows in one
 shared verify against each alone, byte for byte, dense and sparse, up to 64
 rows

---
 zig/tests/cuda/flashnext_windows_tests.zig | 127 +++++++++++++++++++++
 zig/tests/cuda/main.zig                    |   3 +
 2 files changed, 130 insertions(+)
 create mode 100644 zig/tests/cuda/flashnext_windows_tests.zig

diff --git a/zig/tests/cuda/flashnext_windows_tests.zig b/zig/tests/cuda/flashnext_windows_tests.zig
new file mode 100644
index 0000000000..2d264197ee
--- /dev/null
+++ b/zig/tests/cuda/flashnext_windows_tests.zig
@@ -0,0 +1,127 @@
+//! Several streams' drafted windows in one shared verify pass against each window alone, byte for byte.
+const std = @import("std");
+const cuda = @import("cuda");
+const flash = @import("flashnext");
+const check = @import("check.zig");
+const prompt = flash.prompt;
+const mtp = flash.mtp;
+
+/// Window widths in one shared pass: none past 16 rows in all, past 16 (the router's few rows), 32, and 64.
+const shapes = [_][]const usize{ &.{ 1, 1, 1, 1 }, &.{ 5, 5, 5, 5 }, &.{ 1, 16, 3, 7 }, &.{ 16, 16 }, &.{ 16, 16, 16, 16 } };
+
+const Side = struct {
+    seqs: [4]prompt.Seq = undefined,
+    drafts: [4]mtp.Draft = undefined,
+    made: usize = 0,
+
+    fn deinit(s: *Side, gpa: std.mem.Allocator) void {
+        for (s.seqs[0..s.made], s.drafts[0..s.made]) |*q, *d| {
+            d.deinit();
+            q.deinit(gpa);
+        }
+    }
+};
+
+pub fn run(gpu: check.Gpu, dir: []const u8, ids_path: []const u8) !void {
+    const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20));
+    defer gpu.gpa.free(source);
+    var base: std.ArrayList(u32) = .empty;
+    defer base.deinit(gpu.gpa);
+    var words = std.mem.tokenizeAny(u8, source, " \r\n\t");
+    while (words.next()) |word| try base.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10));
+    if (base.items.len < 16) return error.BadPrompt;
+    var stream = try cuda.Stream.init(gpu.d, true);
+    defer stream.deinit();
+    var glue = try flash.Glue.load(gpu.d);
+    defer glue.deinit();
+    const tri: flash.Tri = .{ .s = stream, .own = &glue };
+    const e = try prompt.Engine.init(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir);
+    defer e.deinit();
+    const ops: mtp.Ops(flash.Tri) = .{ .e = e, .tri = tri, .h = undefined };
+    const row_bytes = prompt.vocabulary * 2;
+    var shared = try cuda.DeviceBuffer.alloc(gpu.d, 64 * row_bytes);
+    defer shared.free();
+    // a short prompt (dense attention) and one past 2,048 keys (the block indexer's sparse selection)
+    for ([_]usize{ @min(base.items.len, 100), 2600 }) |length| {
+        const ids = try gpu.gpa.alloc(u32, length);
+        defer gpu.gpa.free(ids);
+        for (ids, 0..) |*t, i| t.* = base.items[i % base.items.len];
+        for (shapes) |widths| try shape(gpu, e, ops, &stream, shared.ptr, ids, widths);
+    }
+}
+
+fn shape(gpu: check.Gpu, e: *prompt.Engine, ops: mtp.Ops(flash.Tri), stream: *cuda.Stream, shared: u64, ids: []const u32, widths: []const usize) !void {
+    const gpa = gpu.gpa;
+    const row_bytes = prompt.vocabulary * 2;
+    const cap = ids.len + 64;
+    var alone: Side = .{};
+    defer alone.deinit(gpa);
+    var together: Side = .{};
+    defer together.deinit(gpa);
+    var windows: [4][16]u32 = undefined;
+    var total: usize = 0;
+    for (widths, 0..) |w, i| {
+        // each stream its own prompt length (so its own positions) and its own forced window
+        const prompt_len = ids.len - 3 * i;
+        for (windows[i][0..w], 0..) |*t, r| t.* = @intCast((37 + 101 * i + 17 * r) % 50_000);
+        inline for (.{ &alone, &together }) |side| {
+            side.seqs[i] = try e.newSeq(cap);
+            side.drafts[i] = mtp.Draft.init(gpa, e.driver, stream.*, cap, 1, e.chunk) catch |err| {
+                side.seqs[i].deinit(gpa);
+                return err;
+            };
+            side.made += 1;
+            try e.forward(flash.Tri, ops.tri, &side.seqs[i], ids[0..prompt_len], .{});
+        }
+        total += w;
+    }
+    const hidden = alone.drafts[0].streams.len / alone.drafts[0].rows_cap;
+    // each window alone, its rows' scores and hidden rows read back
+    const want_logits = try gpa.alloc(u8, total * row_bytes);
+    defer gpa.free(want_logits);
+    const want_hidden = try gpa.alloc(u8, total * hidden);
+    defer gpa.free(want_hidden);
+    var row0: usize = 0;
+    for (widths, 0..) |w, i| {
+        try ops.verify(&alone.seqs[i], &alone.drafts[i], windows[i][0..w]);
+        try stream.synchronize();
+        try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(want_logits[row0 * row_bytes ..].ptr, alone.drafts[i].logits.ptr, w * row_bytes), "cuMemcpyDtoH");
+        try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(want_hidden[row0 * hidden ..].ptr, alone.drafts[i].streams.ptr, w * hidden), "cuMemcpyDtoH");
+        row0 += w;
+    }
+    // the same windows in one shared pass
+    var seqs: [4]*prompt.Seq = undefined;
+    var ds: [4]*mtp.Draft = undefined;
+    var idss: [4][]const u32 = undefined;
+    for (widths, 0..) |w, i| {
+        seqs[i] = &together.seqs[i];
+        ds[i] = &together.drafts[i];
+        idss[i] = windows[i][0..w];
+    }
+    const n = widths.len;
+    try ops.verifyMany(seqs[0..n], ds[0..n], idss[0..n], shared);
+    try stream.synchronize();
+    const got = try gpa.alloc(u8, @max(row_bytes, hidden) * 16);
+    defer gpa.free(got);
+    row0 = 0;
+    for (widths, 0..) |w, i| {
+        try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(got.ptr, together.drafts[i].logits.ptr, w * row_bytes), "cuMemcpyDtoH");
+        try check.sameBytes("a shared window's scores equal its lone pass's", got[0 .. w * row_bytes], want_logits[row0 * row_bytes ..][0 .. w * row_bytes]);
+        try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(got.ptr, together.drafts[i].streams.ptr, w * hidden), "cuMemcpyDtoH");
+        try check.sameBytes("a shared window's hidden rows equal its lone pass's", got[0 .. w * hidden], want_hidden[row0 * hidden ..][0 .. w * hidden]);
+        try check.sameBytes("a shared window's last host row equals its lone pass's", together.seqs[i].logits, alone.seqs[i].logits);
+        row0 += w;
+    }
+    // every stream keeps part of its window (its rollback replays the kept rows), then one more token
+    for (widths, 0..) |w, i| {
+        const keep = 1 + (3 * i) % w;
+        try ops.rollback(&alone.seqs[i], &alone.drafts[i], keep);
+        try ops.rollback(&together.seqs[i], &together.drafts[i], keep);
+        try check.expect(alone.seqs[i].fed == together.seqs[i].fed, "kept positions match", .{});
+        try check.expect(std.mem.eql(i64, &alone.seqs[i].history, &together.seqs[i].history), "kept n-gram histories match", .{});
+        try e.forward(flash.Tri, ops.tri, &alone.seqs[i], &.{13}, .{});
+        try e.forward(flash.Tri, ops.tri, &together.seqs[i], &.{13}, .{});
+        try check.sameBytes("the row after a shared window's rollback equals the lone one's", together.seqs[i].logits, alone.seqs[i].logits);
+    }
+    check.pass("shared verify, {d} windows of {any} rows ({d} in all) at {d} prompt tokens: scores, hidden rows, rollback and next row byte-exact", .{ n, widths, total, ids.len });
+}
diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig
index 837f1bb070..61f1c3b788 100644
--- a/zig/tests/cuda/main.zig
+++ b/zig/tests/cuda/main.zig
@@ -70,6 +70,7 @@ const route_rows_tests = @import("route_rows_tests.zig");
 const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig");
 const flashnext_fidelity_tests = @import("flashnext_fidelity_tests.zig");
 const flashnext_window_tests = @import("flashnext_window_tests.zig");
+const flashnext_windows_tests = @import("flashnext_windows_tests.zig");
 const flashnext_rank_tests = @import("flashnext_rank_tests.zig");
 const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig");
 const flashnext_shared_tests = @import("flashnext_shared_tests.zig");
@@ -98,6 +99,7 @@ const usage =
     \\  flashnext-fidelity-single MODEL FIXTURE OUTPUT  single-rank reference isolation
     \\  flashnext-fidelity-rank MODEL RANK PORT FIXTURE OUTPUT  fixed-history full score rows
     \\  flashnext-window-rank MODEL RANK PORT IDS_FILE  two-rank serial/window and rollback parity
+    \\  flashnext-windows MODEL IDS_FILE  several streams' windows in one shared verify against each alone
     \\  flashnext-model-rank MODEL RANK PORT IDS_FILE  explicit full-model replay experiment (no drafts)
     \\  flashnext-rank RANK PORT [sum|eof|timeout|missing]  two-process collective and failure probes
     \\  flashnext-draft-affine    proposal-head affine projection against FP64 and row-count invariance
@@ -243,6 +245,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void {
     if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), false);
     if (std.mem.eql(u8, cmd, "flashnext-fidelity-single-serial")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), true);
     if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4));
+    if (std.mem.eql(u8, cmd, "flashnext-windows")) return flashnext_windows_tests.run(gpu, try arg(rest, 0), try arg(rest, 1));
     if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3));
     if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3));
     if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum");

From 5cc536ce6fda58fe30959d9325dbc48b187b8b84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:41:02 +0200
Subject: [PATCH 3/4] Flash Next CUDA: a shared round's greedy streams draft
 their MTP heads together (on #1069)

Depends on #1069 (lanes with row spans). The head drafted one stream at a time: at four streams a
round spent ~29 ms in head passes against ~5 ms at one. Now the core asks a shared round's drafts
in one call (.draft_streams) and draftBatched runs them together: each greedy stream's kept rows
absorbed in one head pass (a lane each, on its own head cache), then one pass a chain step for the
streams still drafting, their bests read in one download. Every row runs the ops a lone head pass
runs on it, so each stream gets the drafts it gets alone (same confidence stop, same count).
Sampled streams, trees and media keep the one-at-a-time path. TF_FLASHNEXT_BATCH_HEADS=0 restores
it; TF_FLASHNEXT_DEPTH_LOG=1 prints each stream's chain length by round width.

FP4-Local, GB10, main + #1069 + local stack, affine4 head, four streams: code 130.2 -> 150.7,
chat 88.9 -> 99.8 tok/s aggregate; one stream unchanged (70.7 / 44.2). Every request's rounds and
accepted/drafted equal the one-at-a-time run's; concurrent == alone 4/4, drafted == serial 1/1.
---
 zig/src/families/flashnext/cuda_lanes.zig   |  90 +++++++++++-
 zig/src/families/flashnext/cuda_mtp_ops.zig | 153 ++++++++++++++++++++
 2 files changed, 242 insertions(+), 1 deletion(-)

diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig
index 7d036518d2..3329498eac 100644
--- a/zig/src/families/flashnext/cuda_lanes.zig
+++ b/zig/src/families/flashnext/cuda_lanes.zig
@@ -6,6 +6,18 @@ const lanes = @import("lanes");
 const prompt = @import("cuda_prompt.zig");
 const weights = @import("cuda_prompt_wt.zig");
 const mtp = @import("cuda_mtp.zig");
+const mtp_ops = @import("cuda_mtp_ops.zig");
+
+/// TF_FLASHNEXT_BATCH_HEADS=0 drafts a shared round's streams one at a time (the batched head gives the same drafts).
+fn batchHeads() bool {
+    const v = std.c.getenv("TF_FLASHNEXT_BATCH_HEADS") orelse return true;
+    return !std.mem.eql(u8, std.mem.span(v), "0");
+}
+
+/// TF_FLASHNEXT_DEPTH_LOG=1: every 200 draft calls, the drafts a stream got by round width (streams drafting together).
+var depth_hist: [5][max_hist]u64 = @splat(@splat(0));
+var depth_calls: u64 = 0;
+const max_hist = 17;
 const snap = @import("cuda_snap.zig");
 const ranks = @import("cuda_rank_engine.zig");
 const config = @import("cuda_config.zig");
@@ -231,6 +243,8 @@ pub fn Namespace(comptime Tri: type) type {
                     .hidden_rows = true,
                     .batch_rows = if (drafting) 4 * mtp.max_rows else round_streams,
                     .max_streams = round_streams,
+                    // a shared round's greedy streams draft in one head pass a step (draftBatched)
+                    .draft_streams = drafting and batchHeads(),
                 };
             }
 
@@ -483,7 +497,12 @@ pub fn Namespace(comptime Tri: type) type {
         fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void {
             const self = of(ptr);
             if (self.head == null) return error.NoDraftHead;
-            for (requests) |r| {
+            defer depthLog(self, requests);
+            // greedy chains of a shared round: their heads in one pass a step (the drafts each would get alone)
+            var batched: [mtp_ops.max_lanes]bool = @splat(false);
+            if (batchHeads() and requests.len > 1 and requests.len <= mtp_ops.max_lanes) try draftBatched(self, requests, &batched);
+            for (requests, 0..) |r, ri| {
+                if (ri < batched.len and batched[ri]) continue;
                 if (r.lanes != null) return error.TreesNotBuilt;
                 const l = try self.lane(r.stream);
                 const d = if (l.draft) |*x| x else return error.NoDraftHead;
@@ -506,6 +525,75 @@ pub fn Namespace(comptime Tri: type) type {
             }
         }
 
+        /// The greedy requests of `requests` (no trees, no media, fewer rows than a head piece) drafted together: rollbacks
+        /// as draftFn makes them, then one head pass a step over the streams still chaining; `done[i]` marks those drafted.
+        fn draftBatched(self: *Cuda, requests: []const be.DraftRequest, done: []bool) !void {
+            var ds: [mtp_ops.max_lanes]*mtp.Draft = undefined;
+            var nexts: [mtp_ops.max_lanes][]const u32 = undefined;
+            var firsts: [mtp_ops.max_lanes][1]u32 = undefined;
+            var rows: [mtp_ops.max_lanes]u64 = undefined;
+            var counts: [mtp_ops.max_lanes]usize = undefined;
+            var outs: [mtp_ops.max_lanes]*mtp.Chain = undefined;
+            var which: [mtp_ops.max_lanes]usize = undefined;
+            var n: usize = 0;
+            for (requests, 0..) |r, ri| {
+                if (r.lanes != null) continue;
+                if (r.stream.sampling) |sm| if (sm.temperature > 0) continue;
+                const l = try self.lane(r.stream);
+                const d = if (l.draft) |*x| x else continue;
+                if (d.media != null) continue;
+                if (r.rows) |path| {
+                    for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt;
+                    if (r.follow.len < path.len) return error.PositionMismatch;
+                }
+                if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else {
+                    if (d.open) try self.ops().rollback(&l.seq, d, d.rows);
+                }
+                if (r.position != l.seq.fed + 1) return error.PositionMismatch;
+                const room = l.seq.cap - l.seq.fed - 1;
+                counts[n] = @min(@min(r.depth, self.depth), room);
+                if (r.rows) |path| {
+                    nexts[n] = r.follow[0..path.len];
+                    rows[n] = d.streams.ptr;
+                } else {
+                    firsts[n] = .{self.value(r.first orelse return error.NoFirstToken)};
+                    nexts[n] = &firsts[n];
+                    rows[n] = d.prompt_last;
+                }
+                ds[n] = d;
+                outs[n] = &l.chain;
+                which[n] = ri;
+                n += 1;
+            }
+            if (n < 2) {
+                // one greedy stream: it drafts alone below (its rollback is idempotent)
+                return;
+            }
+            try self.ops().draftMany(ds[0..n], nexts[0..n], rows[0..n], counts[0..n], self.confidence, outs[0..n]);
+            for (which[0..n]) |ri| done[ri] = true;
+        }
+
+        /// TF_FLASHNEXT_DEPTH_LOG=1: each stream's chain length by how many streams drafted together.
+        fn depthLog(self: *Cuda, requests: []const be.DraftRequest) void {
+            const v = std.c.getenv("TF_FLASHNEXT_DEPTH_LOG") orelse return;
+            if (!std.mem.eql(u8, std.mem.span(v), "1")) return;
+            const w = @min(requests.len, 5) - 1;
+            for (requests) |r| {
+                const l = self.lane(r.stream) catch continue;
+                depth_hist[w][@min(l.chain.n, max_hist - 1)] += 1;
+            }
+            depth_calls += 1;
+            if (depth_calls % 200 != 0) return;
+            for (depth_hist, 1..) |row, width| {
+                var any: u64 = 0;
+                for (row) |c| any += c;
+                if (any == 0) continue;
+                std.debug.print("DEPTH width={d}{s}", .{ width, if (width == 5) "+" else "" });
+                for (row, 0..) |c, d| if (c > 0) std.debug.print(" {d}:{d}", .{ d, c });
+                std.debug.print("\n", .{});
+            }
+        }
+
         /// The held chain as host tokens: the round loop verifies them as it would copied ones.
         fn treeFn(ptr: *anyopaque, s: *lanes.Stream, gpa: std.mem.Allocator) anyerror!?lanes.stream.Held {
             const self = of(ptr);
diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig
index 3adf0e7169..d507602fe2 100644
--- a/zig/src/families/flashnext/cuda_mtp_ops.zig
+++ b/zig/src/families/flashnext/cuda_mtp_ops.zig
@@ -30,6 +30,8 @@ const proj_bytes = mtp.proj_bytes;
 const fc_sk = mtp.fc_sk;
 const piece_rows = mtp.piece_rows;
 const pss_n = mtp.pss_n;
+/// Streams one batched head pass drafts for at most.
+pub const max_lanes: usize = 16;
 const Head = mtp.Head;
 const Draft = mtp.Draft;
 const Chain = mtp.Chain;
@@ -226,6 +228,157 @@ pub fn Ops(comptime Tri: type) type {
             }
         }
 
+        /// forwardHead over several drafts in one pass, a lane each (its own cache at d.len): `nexts[i]` rows of streams
+        /// `rows[i]`, each lane's last row in its d.out and (scores) its draft scores in its d.dlogits. Every row runs the
+        /// ops a lone forwardHead runs on it, each row's sums its own, so each draft gets what drafting alone gives it.
+        pub fn forwardHeadMany(o: Self, ds: []const *Draft, nexts: []const []const u32, rows: []const u64, scores: bool) !void {
+            const e = o.e;
+            const tri = o.tri;
+            const h = o.h;
+            const n_l = ds.len;
+            if (n_l == 0) return;
+            var total: usize = 0;
+            var cap: usize = 0;
+            for (ds, nexts) |d, nx| {
+                if (nx.len == 0 or nx.len > piece_rows or d.media != null) return error.UnexpectedTensor;
+                if (d.len + nx.len > d.cap) return error.PromptTooLong;
+                total += nx.len;
+                cap = @max(cap, d.cap);
+            }
+            const s = try e.reserve(@max(@as(usize, 32), 2 * total));
+            if (total > s.cap or n_l > max_lanes) return error.UnexpectedTensor;
+            const score_rows = try e.reserveScores(cap, piece_rows);
+            var lane_buf: [max_lanes]batch.Lane = undefined;
+            var ids_buf: [max_lanes * piece_rows]u32 = undefined;
+            var at: usize = 0;
+            for (ds, nexts, 0..) |d, nx, i| {
+                lane_buf[i] = .{ .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .pos = d.len, .rows = nx.len };
+                @memcpy(ids_buf[at..][0..nx.len], nx);
+                at += nx.len;
+            }
+            var p: batch.Pass(Tri) = .{
+                .gpa = e.gpa,
+                .io = e.io,
+                .stream = e.stream.*,
+                .tri = tri,
+                .k = &e.batch_k,
+                .s = s,
+                .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram },
+                .seq = lane_buf[0].seq,
+                .hooks = .{ .logits = .none },
+                .lanes = lane_buf[0..n_l],
+            };
+            p.rows = total;
+            p.pos0 = 0;
+            p.pending = null;
+            try p.stage(ids_buf[0..total]);
+            if (e.table.dense) {
+                try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, total);
+            } else {
+                try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, total);
+            }
+            try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, total, dims, dims);
+            try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, total, fc_sk);
+            try p.sumSlices(s.part, s.branch, true, total * dims, fc_sk);
+            // each lane's rows of main-model streams, side by side (the head reads them as one block of rows)
+            at = 0;
+            for (rows, nexts) |r, nx| {
+                try p.copy(s.ple_nrow + at * wide * 2, r, nx.len * wide * 2);
+                at += nx.len;
+            }
+            try tri.rmsnorm(s.ple_nrow, h.norm_h.ptr, s.norm, s.gq, eps, wide, total, wide, wide);
+            try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * total, fc_sk);
+            try p.sumSlices(s.part, s.ple_gated, true, streams * total * dims, fc_sk);
+            try tri.addStreams(s.branch, s.ple_gated, s.h, total, dims, streams);
+            try p.writeback(0, total);
+            try p.readout(h.layer.attn_hc, 0, total);
+            try p.attnBlock(&h.layer, 0);
+            try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", total * dims, total, dims, streams, 4, 1, 1, out_sk);
+            try p.readout(h.layer.mlp_hc, 0, total);
+            try p.moeBlock(&h.layer);
+            p.pending = 0;
+            try p.writeback(0, total);
+            // each lane's last row: its d.out, and gathered to row i for the draft head (row i <= its last row)
+            at = 0;
+            for (ds, nexts, 0..) |d, nx, i| {
+                const last = at + nx.len - 1;
+                try p.copy(d.out.ptr, s.h + last * wide * 2, wide * 2);
+                if (scores and last != i) {
+                    try p.copy(s.h + i * wide * 2, s.h + last * wide * 2, wide * 2);
+                    try p.copy(s.pss + i * pss_n * 4, s.pss + last * pss_n * 4, pss_n * 4);
+                }
+                d.len += nx.len;
+                at += nx.len;
+            }
+            if (!scores) return;
+            try p.readout(h.mixer, 0, n_l);
+            const n_ids = h.ids.len;
+            try p.lane(s.mixed, s.xsm, h.draft, s.head_local, s.dummy, n_l, 1);
+            for (ds, 0..) |d, i| try p.copy(d.dlogits.ptr, s.head_local + i * n_ids * 2, n_ids * 2);
+        }
+
+        /// draft over several greedy drafts at once: absorb each one's kept rows in one pass, then chain them a step a
+        /// pass, each stream picking and stopping as it would alone (same confidence rule, same count).
+        pub fn draftMany(o: Self, ds: []const *Draft, nexts: []const []const u32, rows: []const u64, counts: []const usize, confidence: f64, outs: []const *Chain) !void {
+            const h = o.h;
+            const n_l = ds.len;
+            for (ds, outs) |d, out| {
+                out.n = 0;
+                if (d.drafted > 0) {
+                    d.len -= d.drafted;
+                    d.drafted = 0;
+                }
+            }
+            try o.forwardHeadMany(ds, nexts, rows, true);
+            var act: [max_lanes]usize = undefined;
+            var n_act: usize = 0;
+            for (counts, 0..) |c, i| if (c > 0) {
+                act[n_act] = i;
+                n_act += 1;
+            };
+            var j: usize = 0;
+            var best_host: [max_lanes * 4]f32 = undefined;
+            var step_ds: [max_lanes]*Draft = undefined;
+            var step_next: [max_lanes][1]u32 = undefined;
+            var step_nexts: [max_lanes][]const u32 = undefined;
+            var step_rows: [max_lanes]u64 = undefined;
+            while (n_act > 0 and j < max_chain) : (j += 1) {
+                // every active lane's best draft (fn_row_best on its own scores), then one download of all of them
+                for (act[0..n_act], 0..) |i, r| try h.best(o.stream(), ds[i].dlogits.ptr, h.ids.len, h.ids.len, 1, ds[0].best.ptr + r * 16);
+                try download(o.stream(), std.mem.sliceAsBytes(best_host[0 .. n_act * 4]), ds[0].best.ptr);
+                var next_act: usize = 0;
+                var n_step: usize = 0;
+                for (act[0..n_act], 0..) |i, r| {
+                    const top = best_host[r * 4];
+                    const lse = best_host[r * 4 + 1];
+                    const col: usize = @intCast(@as(i32, @bitCast(best_host[r * 4 + 2])));
+                    const got: Pick = .{ .token = h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) };
+                    const out = outs[i];
+                    const low = confidence > 0 and got.prob < confidence;
+                    if (low and j > 0) continue;
+                    out.tokens[out.n] = got.token;
+                    out.probs[out.n] = got.prob;
+                    out.n += 1;
+                    if (low) continue;
+                    if (j + 1 < @min(counts[i], max_chain)) {
+                        step_ds[n_step] = ds[i];
+                        step_next[n_step] = .{got.token};
+                        step_nexts[n_step] = &step_next[n_step];
+                        step_rows[n_step] = ds[i].out.ptr;
+                        n_step += 1;
+                        act[next_act] = i;
+                        next_act += 1;
+                    }
+                }
+                if (n_step > 0) {
+                    try o.forwardHeadMany(step_ds[0..n_step], step_nexts[0..n_step], step_rows[0..n_step], true);
+                    for (step_ds[0..n_step]) |d| d.drafted += 1;
+                }
+                n_act = next_act;
+            }
+            _ = n_l;
+        }
+
         /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams.
         pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void {
             if (o.main) |m| return m.verify(m.ptr, seq, d, ids);

From 282feaea00ee34c604aa77b1d046f36ea14a0a94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Wed, 21 Oct 2026 15:32:17 +0200
Subject: [PATCH 4/4] Flash Next CUDA: one-line comments, the batched draft in
 cuda_lanes_draft.zig

The lean gate's new findings on this branch: multi-line comments folded to one line within 120 columns,
and draftBatched, depthLog and batchHeads moved unchanged to cuda_lanes_draft.zig (as cuda_lanes_verify.zig
holds the shared verify), which brings cuda_lanes.zig back to 587 lines; value() is pub for it, as lane()
and ops() are. lean_new.py: 0 new.
---
 zig/src/families/flashnext/cuda_lanes.zig     | 91 ++-----------------
 .../families/flashnext/cuda_lanes_draft.zig   | 85 +++++++++++++++++
 zig/src/families/flashnext/cuda_mtp_ops.zig   |  7 +-
 3 files changed, 93 insertions(+), 90 deletions(-)
 create mode 100644 zig/src/families/flashnext/cuda_lanes_draft.zig

diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig
index 3329498eac..228d795db0 100644
--- a/zig/src/families/flashnext/cuda_lanes.zig
+++ b/zig/src/families/flashnext/cuda_lanes.zig
@@ -7,17 +7,7 @@ const prompt = @import("cuda_prompt.zig");
 const weights = @import("cuda_prompt_wt.zig");
 const mtp = @import("cuda_mtp.zig");
 const mtp_ops = @import("cuda_mtp_ops.zig");
-
-/// TF_FLASHNEXT_BATCH_HEADS=0 drafts a shared round's streams one at a time (the batched head gives the same drafts).
-fn batchHeads() bool {
-    const v = std.c.getenv("TF_FLASHNEXT_BATCH_HEADS") orelse return true;
-    return !std.mem.eql(u8, std.mem.span(v), "0");
-}
-
-/// TF_FLASHNEXT_DEPTH_LOG=1: every 200 draft calls, the drafts a stream got by round width (streams drafting together).
-var depth_hist: [5][max_hist]u64 = @splat(@splat(0));
-var depth_calls: u64 = 0;
-const max_hist = 17;
+const lanes_draft = @import("cuda_lanes_draft.zig");
 const snap = @import("cuda_snap.zig");
 const ranks = @import("cuda_rank_engine.zig");
 const config = @import("cuda_config.zig");
@@ -243,8 +233,8 @@ pub fn Namespace(comptime Tri: type) type {
                     .hidden_rows = true,
                     .batch_rows = if (drafting) 4 * mtp.max_rows else round_streams,
                     .max_streams = round_streams,
-                    // a shared round's greedy streams draft in one head pass a step (draftBatched)
-                    .draft_streams = drafting and batchHeads(),
+                    // a shared round's greedy streams draft in one head pass a step (lanes_draft.batched)
+                    .draft_streams = drafting and lanes_draft.batchHeads(),
                 };
             }
 
@@ -255,7 +245,7 @@ pub fn Namespace(comptime Tri: type) type {
                 return h;
             }
 
-            fn value(self: *const Cuda, feed: be.Feed) u32 {
+            pub fn value(self: *const Cuda, feed: be.Feed) u32 {
                 return switch (feed) {
                     .handle => |h| self.drawn[h % ring],
                     .value => |v| v,
@@ -497,10 +487,10 @@ pub fn Namespace(comptime Tri: type) type {
         fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void {
             const self = of(ptr);
             if (self.head == null) return error.NoDraftHead;
-            defer depthLog(self, requests);
+            defer lanes_draft.depthLog(self, requests);
             // greedy chains of a shared round: their heads in one pass a step (the drafts each would get alone)
             var batched: [mtp_ops.max_lanes]bool = @splat(false);
-            if (batchHeads() and requests.len > 1 and requests.len <= mtp_ops.max_lanes) try draftBatched(self, requests, &batched);
+            if (lanes_draft.batchHeads() and requests.len > 1 and requests.len <= mtp_ops.max_lanes) try lanes_draft.batched(self, requests, &batched);
             for (requests, 0..) |r, ri| {
                 if (ri < batched.len and batched[ri]) continue;
                 if (r.lanes != null) return error.TreesNotBuilt;
@@ -525,75 +515,6 @@ pub fn Namespace(comptime Tri: type) type {
             }
         }
 
-        /// The greedy requests of `requests` (no trees, no media, fewer rows than a head piece) drafted together: rollbacks
-        /// as draftFn makes them, then one head pass a step over the streams still chaining; `done[i]` marks those drafted.
-        fn draftBatched(self: *Cuda, requests: []const be.DraftRequest, done: []bool) !void {
-            var ds: [mtp_ops.max_lanes]*mtp.Draft = undefined;
-            var nexts: [mtp_ops.max_lanes][]const u32 = undefined;
-            var firsts: [mtp_ops.max_lanes][1]u32 = undefined;
-            var rows: [mtp_ops.max_lanes]u64 = undefined;
-            var counts: [mtp_ops.max_lanes]usize = undefined;
-            var outs: [mtp_ops.max_lanes]*mtp.Chain = undefined;
-            var which: [mtp_ops.max_lanes]usize = undefined;
-            var n: usize = 0;
-            for (requests, 0..) |r, ri| {
-                if (r.lanes != null) continue;
-                if (r.stream.sampling) |sm| if (sm.temperature > 0) continue;
-                const l = try self.lane(r.stream);
-                const d = if (l.draft) |*x| x else continue;
-                if (d.media != null) continue;
-                if (r.rows) |path| {
-                    for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt;
-                    if (r.follow.len < path.len) return error.PositionMismatch;
-                }
-                if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else {
-                    if (d.open) try self.ops().rollback(&l.seq, d, d.rows);
-                }
-                if (r.position != l.seq.fed + 1) return error.PositionMismatch;
-                const room = l.seq.cap - l.seq.fed - 1;
-                counts[n] = @min(@min(r.depth, self.depth), room);
-                if (r.rows) |path| {
-                    nexts[n] = r.follow[0..path.len];
-                    rows[n] = d.streams.ptr;
-                } else {
-                    firsts[n] = .{self.value(r.first orelse return error.NoFirstToken)};
-                    nexts[n] = &firsts[n];
-                    rows[n] = d.prompt_last;
-                }
-                ds[n] = d;
-                outs[n] = &l.chain;
-                which[n] = ri;
-                n += 1;
-            }
-            if (n < 2) {
-                // one greedy stream: it drafts alone below (its rollback is idempotent)
-                return;
-            }
-            try self.ops().draftMany(ds[0..n], nexts[0..n], rows[0..n], counts[0..n], self.confidence, outs[0..n]);
-            for (which[0..n]) |ri| done[ri] = true;
-        }
-
-        /// TF_FLASHNEXT_DEPTH_LOG=1: each stream's chain length by how many streams drafted together.
-        fn depthLog(self: *Cuda, requests: []const be.DraftRequest) void {
-            const v = std.c.getenv("TF_FLASHNEXT_DEPTH_LOG") orelse return;
-            if (!std.mem.eql(u8, std.mem.span(v), "1")) return;
-            const w = @min(requests.len, 5) - 1;
-            for (requests) |r| {
-                const l = self.lane(r.stream) catch continue;
-                depth_hist[w][@min(l.chain.n, max_hist - 1)] += 1;
-            }
-            depth_calls += 1;
-            if (depth_calls % 200 != 0) return;
-            for (depth_hist, 1..) |row, width| {
-                var any: u64 = 0;
-                for (row) |c| any += c;
-                if (any == 0) continue;
-                std.debug.print("DEPTH width={d}{s}", .{ width, if (width == 5) "+" else "" });
-                for (row, 0..) |c, d| if (c > 0) std.debug.print(" {d}:{d}", .{ d, c });
-                std.debug.print("\n", .{});
-            }
-        }
-
         /// The held chain as host tokens: the round loop verifies them as it would copied ones.
         fn treeFn(ptr: *anyopaque, s: *lanes.Stream, gpa: std.mem.Allocator) anyerror!?lanes.stream.Held {
             const self = of(ptr);
diff --git a/zig/src/families/flashnext/cuda_lanes_draft.zig b/zig/src/families/flashnext/cuda_lanes_draft.zig
new file mode 100644
index 0000000000..1d6bae4fe3
--- /dev/null
+++ b/zig/src/families/flashnext/cuda_lanes_draft.zig
@@ -0,0 +1,85 @@
+//! A shared round's greedy streams drafted in one head pass a step (cuda_lanes.zig's draftFn), and the depth log.
+const std = @import("std");
+const lanes = @import("lanes");
+const mtp = @import("cuda_mtp.zig");
+const mtp_ops = @import("cuda_mtp_ops.zig");
+const be = lanes.backend;
+
+/// TF_FLASHNEXT_BATCH_HEADS=0 drafts a shared round's streams one at a time (the batched head gives the same drafts).
+pub fn batchHeads() bool {
+    const v = std.c.getenv("TF_FLASHNEXT_BATCH_HEADS") orelse return true;
+    return !std.mem.eql(u8, std.mem.span(v), "0");
+}
+
+/// TF_FLASHNEXT_DEPTH_LOG=1: every 200 draft calls, the drafts a stream got by round width (streams drafting together).
+var depth_hist: [5][max_hist]u64 = @splat(@splat(0));
+var depth_calls: u64 = 0;
+const max_hist = 17;
+
+/// Greedy requests (no trees, no media) drafted together: draftFn's rollbacks, one head pass a step; `done` marks them.
+pub fn batched(self: anytype, requests: []const be.DraftRequest, done: []bool) !void {
+    var ds: [mtp_ops.max_lanes]*mtp.Draft = undefined;
+    var nexts: [mtp_ops.max_lanes][]const u32 = undefined;
+    var firsts: [mtp_ops.max_lanes][1]u32 = undefined;
+    var rows: [mtp_ops.max_lanes]u64 = undefined;
+    var counts: [mtp_ops.max_lanes]usize = undefined;
+    var outs: [mtp_ops.max_lanes]*mtp.Chain = undefined;
+    var which: [mtp_ops.max_lanes]usize = undefined;
+    var n: usize = 0;
+    for (requests, 0..) |r, ri| {
+        if (r.lanes != null) continue;
+        if (r.stream.sampling) |sm| if (sm.temperature > 0) continue;
+        const l = try self.lane(r.stream);
+        const d = if (l.draft) |*x| x else continue;
+        if (d.media != null) continue;
+        if (r.rows) |path| {
+            for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt;
+            if (r.follow.len < path.len) return error.PositionMismatch;
+        }
+        if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else {
+            if (d.open) try self.ops().rollback(&l.seq, d, d.rows);
+        }
+        if (r.position != l.seq.fed + 1) return error.PositionMismatch;
+        const room = l.seq.cap - l.seq.fed - 1;
+        counts[n] = @min(@min(r.depth, self.depth), room);
+        if (r.rows) |path| {
+            nexts[n] = r.follow[0..path.len];
+            rows[n] = d.streams.ptr;
+        } else {
+            firsts[n] = .{self.value(r.first orelse return error.NoFirstToken)};
+            nexts[n] = &firsts[n];
+            rows[n] = d.prompt_last;
+        }
+        ds[n] = d;
+        outs[n] = &l.chain;
+        which[n] = ri;
+        n += 1;
+    }
+    if (n < 2) {
+        // one greedy stream: it drafts alone below (its rollback is idempotent)
+        return;
+    }
+    try self.ops().draftMany(ds[0..n], nexts[0..n], rows[0..n], counts[0..n], self.confidence, outs[0..n]);
+    for (which[0..n]) |ri| done[ri] = true;
+}
+
+/// TF_FLASHNEXT_DEPTH_LOG=1: each stream's chain length by how many streams drafted together.
+pub fn depthLog(self: anytype, requests: []const be.DraftRequest) void {
+    const v = std.c.getenv("TF_FLASHNEXT_DEPTH_LOG") orelse return;
+    if (!std.mem.eql(u8, std.mem.span(v), "1")) return;
+    const w = @min(requests.len, 5) - 1;
+    for (requests) |r| {
+        const l = self.lane(r.stream) catch continue;
+        depth_hist[w][@min(l.chain.n, max_hist - 1)] += 1;
+    }
+    depth_calls += 1;
+    if (depth_calls % 200 != 0) return;
+    for (depth_hist, 1..) |row, width| {
+        var any: u64 = 0;
+        for (row) |c| any += c;
+        if (any == 0) continue;
+        std.debug.print("DEPTH width={d}{s}", .{ width, if (width == 5) "+" else "" });
+        for (row, 0..) |c, d| if (c > 0) std.debug.print(" {d}:{d}", .{ d, c });
+        std.debug.print("\n", .{});
+    }
+}
diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig
index d507602fe2..956345d081 100644
--- a/zig/src/families/flashnext/cuda_mtp_ops.zig
+++ b/zig/src/families/flashnext/cuda_mtp_ops.zig
@@ -228,9 +228,7 @@ pub fn Ops(comptime Tri: type) type {
             }
         }
 
-        /// forwardHead over several drafts in one pass, a lane each (its own cache at d.len): `nexts[i]` rows of streams
-        /// `rows[i]`, each lane's last row in its d.out and (scores) its draft scores in its d.dlogits. Every row runs the
-        /// ops a lone forwardHead runs on it, each row's sums its own, so each draft gets what drafting alone gives it.
+        /// forwardHead over several drafts in one pass, a lane each: every row's ops and sums as a lone forwardHead's.
         pub fn forwardHeadMany(o: Self, ds: []const *Draft, nexts: []const []const u32, rows: []const u64, scores: bool) !void {
             const e = o.e;
             const tri = o.tri;
@@ -317,8 +315,7 @@ pub fn Ops(comptime Tri: type) type {
             for (ds, 0..) |d, i| try p.copy(d.dlogits.ptr, s.head_local + i * n_ids * 2, n_ids * 2);
         }
 
-        /// draft over several greedy drafts at once: absorb each one's kept rows in one pass, then chain them a step a
-        /// pass, each stream picking and stopping as it would alone (same confidence rule, same count).
+        /// draft over several greedy drafts: kept rows in one pass, then a pass a step, each chaining as it would alone
         pub fn draftMany(o: Self, ds: []const *Draft, nexts: []const []const u32, rows: []const u64, counts: []const usize, confidence: f64, outs: []const *Chain) !void {
             const h = o.h;
             const n_l = ds.len;