Flash Next CUDA: a shared round's greedy streams draft their MTP heads together

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

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;

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

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

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

zig/src/families/flashnext/cuda_lanes.zig

@@ -6,6 +6,8 @@ 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");
+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");
@@ -18,6 +20,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 +178,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,
@@ -229,6 +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 (lanes_draft.batched)
+                    .draft_streams = drafting and lanes_draft.batchHeads(),
                 };
             }
 
@@ -239,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,
@@ -255,7 +261,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 +435,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 +462,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);
@@ -475,7 +487,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 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 (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;
                 const l = try self.lane(r.stream);
                 const d = if (l.draft) |*x| x else return error.NoDraftHead;