gemma4 cuda: TF_GEMMA4_CHUNK_SOLO, on a solo-chunk rule shared in lanes (#1554 option C)

CHANGELOG.md

@@ -7,6 +7,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 - Internal: HIP record counts are pinned in their owning arrays, so additions in different record files no longer edit a shared numeric total; omissions still fail host checks (#2461).
 - Qwen3.8 Flash Next on CUDA: BF16 prompt projections take a 128 × 128 tile on sm_12x wherever it fills the GPU, with the same output bytes, so prompts are 4 to 6% faster on a GB10 (NVFP4 checkpoint) and replies are unchanged. `TF_FLASHNEXT_PROMPT_BF16=0` turns the tile off, and `TF_FLASHNEXT_PROMPT_ROWS` is gone (#1846, refs #1754). Thanks to @laudney.
 - Internal: `compile-cuda-tests` builds again: `contract_controls.cu` declares the prompt-row fragment's page settings (`P256`, `IDS`, `PN`), which its control kernel predated (#2463).
+- Gemma 4 on CUDA: `TF_GEMMA4_CHUNK_SOLO=2048` lets a prompt pass take 2,048 rows while no other stream decodes, and the usual 512 once one does. On a GB10 with the 31B 4-bit release, a cold prompt's first token comes 4-5% sooner at 4K, 16K and 29K tokens with the same replies, and decode beside a filling prompt keeps its pace; each stream's sliding rings take 1.2 GiB more. Unset, nothing changes. The rule now lives in `lanes.solo_chunk`, which Flash Next's `TF_FLASHNEXT_CHUNK_SOLO` and GLM's `TF_GLM_CHUNK_SOLO` also call.
+- Gemma 4 on CUDA: `TF_GEMMA4_CHUNK_SOLO=2048` lets a prompt pass take 2,048 rows while no other stream decodes, and the usual 512 once one does. On a GB10 with the 31B 4-bit release, a cold prompt's first token comes 4-5% sooner at 4K, 16K and 29K tokens with the same replies, and decode beside a filling prompt keeps its pace; each stream's sliding rings take 1.2 GiB more. Unset, nothing changes. The rule now lives in `lanes.solo_chunk`, which Flash Next's `TF_FLASHNEXT_CHUNK_SOLO` and GLM's `TF_GLM_CHUNK_SOLO` also call (#2391).
 
 - Qwen3.5 and Qwen3.6 on CUDA: a stream's rollback copy shares its verify log's bytes, which are never in use together, so each pooled stream takes 147 MiB less on Qwen3.8-27B and 61 MiB less on the 35B-A3B; tokens and speed unchanged (#2447, part of #1943). Thanks to @jschmied.
 - LightOnOCR-3 on HIP: the native server reads `lightonai/LightOnOCR-3-0.8B`, `-1B` and `-4B` (BF16) and serves their

docs/optimization-ledger.md

@@ -64,7 +64,7 @@ last two hold the families with only a Metal, HIP or scoring backend.
 | 17. DSpark | yes | yes | yes | yes | yes |
 | 18. Copy drafts | done [#1477] | done [#468] | done [#2146] | done [#2022] | done [#1974] |
 | 19. Shared drafts | done [#1562] | no (null [#1762]) | done [#1779] | done [#1779] | done [#768] |
-| 20. Solo chunks | yes | done C[#1733] M[#2167] | done [#1823] | yes | yes |
+| 20. Solo chunks | yes | done C[#1733] M[#2167] | done [#1823] | done C[#2391] | yes |
 | 21. Read-ahead | done [#1928] | no (row 22) | done [#1616] | done [#1928] | done [#2033] |
 | 22. Direct load | done [#2282] | done [#1901] | yes | done [#2282] | done [#458] |
 | 23. Page drop | done [#1884] | done [#1880] | done [#600] | done [#1884] | done [#1884] |
@@ -374,6 +374,7 @@ below; these are estimates for planning a port, not results.
     | flashnext (M) | 1,024-token reply beside a 64K prompt 48 → 8 s, prompt's first token 90 → 54 s | [#2167] |
     | flashnext (C), one GPU | prompts 15–23 % faster at `--parallel 1`, up to 8,192 rows a pass | [#2246] |
     | glm (C), TP=2 | prompts about 2 % sooner at 8K–128K, two chunks a step, replies unchanged | [#1423] |
+    | gemma4 (C) | first token 31B −4..−5 %, E4B −9..−10 %, opt-in with `TF_GEMMA4_CHUNK_SOLO` | [#2391] |
 
     `TF_GLM_CHUNK_SOLO=4096` is slower with GLM's NVFP4 checkpoint on two GB10s ([#2288]).
 
@@ -551,3 +552,4 @@ below; these are estimates for planning a port, not results.
 [#2283]: https://github.com/ashhart/TensorFold/pull/2283
 [#2288]: https://github.com/ashhart/TensorFold/pull/2288
 [#2306]: https://github.com/ashhart/TensorFold/pull/2306
+[#2391]: https://github.com/ashhart/TensorFold/pull/2391

zig/src/core/lanes/lanes.zig

@@ -12,6 +12,7 @@ pub const recall_store = @import("recall_store.zig");
 pub const sampling = @import("sampling.zig");
 pub const fill = @import("fill.zig");
 pub const prefill_widths = @import("prefill_widths.zig");
+pub const solo_chunk = @import("solo_chunk.zig");
 pub const shape = @import("shape.zig");
 pub const plan_lanes = @import("plan_lanes.zig");
 pub const cost_cache = @import("cost_cache.zig");

zig/src/core/lanes/solo_chunk.zig

@@ -0,0 +1,55 @@
+//! Solo prompt passes (#1554 C), backend-free: wider chunks while no other stream decodes, the busy chunk else.
+const std = @import("std");
+
+/// A family's limits: its busy chunk, its widest solo pass, the grid a solo pass keeps (1: any row count).
+pub const Limits = struct { busy: usize, most: usize, grid: usize = 1 };
+
+/// A solo value's rows: unset, `l.busy` (off); `l.busy` to `l.most` on `l.grid`, the rest error.BadChunkSolo.
+pub fn parse(v: ?[]const u8, l: Limits) !usize {
+    const s = v orelse return l.busy;
+    const n = std.fmt.parseInt(usize, s, 10) catch return error.BadChunkSolo;
+    if (n < l.busy or n > l.most or n % l.grid != 0) return error.BadChunkSolo;
+    return n;
+}
+
+/// The named variable's solo rows (parse).
+pub fn rows(name: [:0]const u8, l: Limits) !usize {
+    return parse(if (std.c.getenv(name)) |v| std.mem.span(v) else null, l);
+}
+
+/// Whether a stream besides `own` is past its prompt pass: its lane's `fill` field (where its pass stands) is null.
+pub fn othersDecode(lanes: anytype, own: anytype, comptime fill: []const u8) bool {
+    var it = lanes.iterator();
+    while (it.next()) |kv| if (kv.key_ptr.* != own and @field(kv.value_ptr.*, fill) == null) return true;
+    return false;
+}
+
+/// The rows the pass of stream `own` takes as it starts: `solo` unless another stream decodes, then `busy`.
+pub fn passRows(lanes: anytype, own: anytype, comptime fill: []const u8, solo: usize, busy: usize) usize {
+    if (solo == busy) return busy;
+    return if (othersDecode(lanes, own, fill)) busy else solo;
+}
+
+test "a solo value is off unset, takes busy to most on the grid and refuses the rest" {
+    const l: Limits = .{ .busy = 512, .most = 4096, .grid = 512 };
+    try std.testing.expectEqual(@as(usize, 512), try parse(null, l));
+    try std.testing.expectEqual(@as(usize, 2048), try parse("2048", l));
+    try std.testing.expectEqual(@as(usize, 4096), try parse("4096", l));
+    const bad = [_][]const u8{ "0", "256", "1000", "4608", "-1", "2k", "" };
+    for (bad) |v| try std.testing.expectError(error.BadChunkSolo, parse(v, l));
+    try std.testing.expectEqual(@as(usize, 1537), try parse("1537", .{ .busy = 1536, .most = 16384 }));
+}
+
+test "a pass takes the solo rows unless another stream decodes; prompts still filling do not count" {
+    const L = struct { fill: ?usize };
+    var m: std.AutoHashMapUnmanaged(u32, L) = .empty;
+    defer m.deinit(std.testing.allocator);
+    try m.put(std.testing.allocator, 0, .{ .fill = 0 });
+    try std.testing.expectEqual(@as(usize, 2048), passRows(&m, 0, "fill", 2048, 512));
+    try m.put(std.testing.allocator, 1, .{ .fill = 1024 });
+    try std.testing.expectEqual(@as(usize, 2048), passRows(&m, 0, "fill", 2048, 512));
+    m.getPtr(1).?.fill = null;
+    try std.testing.expectEqual(@as(usize, 512), passRows(&m, 0, "fill", 2048, 512));
+    try std.testing.expectEqual(@as(usize, 2048), passRows(&m, 1, "fill", 2048, 512)); // its own lane never counts
+    try std.testing.expectEqual(@as(usize, 512), passRows(&m, 0, "fill", 512, 512));
+}

zig/src/families/flashnext/cuda_chunk.zig

@@ -1,6 +1,7 @@
 //! Prompt rows a pass: TF_FLASHNEXT_CHUNK while another stream decodes, TF_FLASHNEXT_CHUNK_SOLO when none does.
 const std = @import("std");
 const batch = @import("cuda_batch.zig");
+const solo_chunk = @import("lanes").solo_chunk;
 
 /// TF_FLASHNEXT_CHUNK's rows a pass: unset, the default; otherwise 1 to batch.max_chunk, anything else refused.
 pub fn parseChunk(v: ?[]const u8) !usize {
@@ -19,22 +20,17 @@ pub const lone_rows: usize = 8192;
 
 /// TF_FLASHNEXT_CHUNK_SOLO's rows: unset, `lone_rows` when `lone`, else `busy` (off); below `busy`, refused.
 pub fn parseSolo(v: ?[]const u8, busy: usize, lone: bool) !usize {
-    const s = v orelse return if (lone) @max(busy, lone_rows) else busy;
-    const n = parseChunk(s) catch return error.BadChunk;
-    if (n < busy) return error.BadChunk;
-    return n;
+    if (v == null and lone) return @max(busy, lone_rows);
+    return solo_chunk.parse(v, .{ .busy = busy, .most = batch.max_chunk }) catch error.BadChunk;
 }
 
 pub fn soloRows(busy: usize, lone: bool) !usize {
     return parseSolo(if (std.c.getenv("TF_FLASHNEXT_CHUNK_SOLO")) |v| std.mem.span(v) else null, busy, lone);
 }
 
-/// The rows a pass of `own` takes when it starts: `solo` unless another lane's prompt is in (it decodes).
+/// The rows the pass of stream `own` takes when it starts: lanes.solo_chunk's rule on the lanes' `filled`.
 pub fn passRows(lanes: anytype, own: anytype, solo: usize, busy: usize) usize {
-    if (solo == busy) return busy;
-    var it = lanes.valueIterator();
-    while (it.next()) |l| if (l != own and l.filled == null) return busy;
-    return solo;
+    return solo_chunk.passRows(lanes, own, "filled", solo, busy);
 }
 
 var readout_rows: ?usize = null;
@@ -73,12 +69,10 @@ test "a pass takes the solo rows unless another lane decodes; prompts still fill
     var m: std.AutoHashMapUnmanaged(u32, L) = .empty;
     defer m.deinit(std.testing.allocator);
     try m.put(std.testing.allocator, 0, .{ .filled = 0 });
-    const own = m.getPtr(0).?;
-    try std.testing.expectEqual(@as(usize, 8192), passRows(&m, own, 8192, 1536));
+    try std.testing.expectEqual(@as(usize, 8192), passRows(&m, 0, 8192, 1536));
     try m.put(std.testing.allocator, 1, .{ .filled = 3072 });
-    const mine = m.getPtr(0).?;
-    try std.testing.expectEqual(@as(usize, 8192), passRows(&m, mine, 8192, 1536));
+    try std.testing.expectEqual(@as(usize, 8192), passRows(&m, 0, 8192, 1536));
     m.getPtr(1).?.filled = null;
-    try std.testing.expectEqual(@as(usize, 1536), passRows(&m, m.getPtr(0).?, 8192, 1536));
-    try std.testing.expectEqual(@as(usize, 1536), passRows(&m, m.getPtr(0).?, 1536, 1536));
+    try std.testing.expectEqual(@as(usize, 1536), passRows(&m, 0, 8192, 1536));
+    try std.testing.expectEqual(@as(usize, 1536), passRows(&m, 0, 1536, 1536));
 }

zig/src/families/flashnext/cuda_lanes.zig

@@ -338,7 +338,7 @@ pub fn Namespace(comptime Tri: type) type {
                     s.keepInterruptedPrefill(@intCast(at));
                     return error.Cancelled;
                 }
-                const rows = @import("cuda_chunk.zig").passRows(&self.lanes, l, self.engine.chunk, self.engine.busy_chunk);
+                const rows = @import("cuda_chunk.zig").passRows(&self.lanes, s, self.engine.chunk, self.engine.busy_chunk);
                 var marks: [prompt.batch.max_marks]prompt.batch.Mark = .{ .{}, .{} };
                 defer snaps.dropPrepared(self, l);
                 const cut = snaps.cut(self, s, l, ids, at, @min(ids.len, at + rows), &marks);

zig/src/families/gemma4/cuda.zig

@@ -48,6 +48,7 @@ test {
     _ = @import("cuda_forward.zig");
     _ = @import("cuda_snap.zig");
     _ = @import("cuda_backend.zig");
+    _ = @import("cuda_solo.zig");
     _ = @import("cuda_shared.zig");
     _ = @import("cuda_assistant.zig");
     _ = @import("cuda_assistant_weights.zig");

zig/src/families/gemma4/cuda_assistant.zig

@@ -79,7 +79,7 @@ pub const Drafter = struct {
         const hid = c.hidden;
         const qw = c.heads * cfg.global_dim;
         const wide: Window = .{ .h = c.heads, .hk = geo.global_kv_heads, .d = cfg.global_dim, .win = 0, .ring = m.context, .scale = 1 };
-        const narrow: Window = .{ .h = c.heads, .hk = geo.sliding_kv_heads, .d = cfg.sliding_dim, .win = cfg.window, .ring = fwd.ring, .scale = 1 };
+        const narrow: Window = .{ .h = c.heads, .hk = geo.sliding_kv_heads, .d = cfg.sliding_dim, .win = cfg.window, .ring = m.ringRows(), .scale = 1 };
         const part = @max(cuda.window_attention.WindowAttention.scratchBytes(wide, 16), cuda.window_attention.WindowAttention.scratchBytes(narrow, 16));
         x.x = .{
             .ids = try x.alloc(n * 4),

zig/src/families/gemma4/cuda_backend.zig

@@ -22,6 +22,7 @@ pub const max_window = 16;
 pub const pagesFor = shared.pagesFor;
 /// Drawn tokens a handle names (a ring; read before the round loop could wrap it).
 const ring = 1024;
+const max_injects = 8; // images one prompt pass can take; a solo pass stays under it
 
 const Lane = struct {
     slot: usize,
@@ -337,20 +338,23 @@ pub const Gemma = struct {
         return k.at;
     }
 
-    /// One chunk from `l.fill` to the next grid line, image or mark; true once the prompt's last row is in. `_: u32`.
+    /// One chunk from `l.fill` to a grid line, image or mark; true once the prompt's last row is in. `_: u32`.
     fn chunk(x: *Gemma, s: *Stream, l: *Lane, streamed: bool, _: u32) !bool {
         const p = s.prompt();
         const at = l.fill.?;
         if (s.isCancelled()) return error.Cancelled;
-        var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows);
+        // solo passes (TF_GEMMA4_CHUNK_SOLO) only without a drafter's taps, at most `injects` images a pass
+        const solo = lanes.solo_chunk.passRows(&x.lanes, s, "fill", x.m.prompt_most, prompt_rows);
+        const rows = if (x.tapped or s.images.len > max_injects) prompt_rows else solo;
+        var end = @min(p.len, at / prompt_rows * prompt_rows + rows);
         for (if (l.shared) l.marks else s.reuse.marks) |mk| if (mk > at and mk < end and !insideImage(s, mk)) {
             end = mk;
         };
         // never through an image: its rows see each other, so a chunk holds all of them (one is at most 280)
         for (s.images) |im| if (im.at < end and end < im.at + im.tokens) {
             end = if (im.at > at) im.at else im.at + im.tokens;
         };
-        var injects: [8]fwd.Inject = undefined;
+        var injects: [max_injects]fwd.Inject = undefined;
         var n: usize = 0;
         var off: usize = 0;
         for (s.images) |im| {