From 1806f7fda05187df9c7111cc79f44d7b963c6e71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Wed, 4 Nov 2026 15:57:10 +0100
Subject: [PATCH 1/4] qwen3_5 cuda: a pipeline stage holds only its layer
 range's weights and state, and runs its part of the forward

RFC #2005 step 2. Geometry.stage slices a model to layers [first, end): linear, bits and kinds stay indexed by
global layer, deltas() and attentions() count the stage's own, so every state, KV, log and keep size follows.
The loader names tensors by global layer and loads the embedding only from layer 0 and the final norm and head
only where the stage holds them (MLX checkpoints; ModelOpt, EXL3, plain BF16 and pockets refused). A stage
without the head carves no logits scratch (cuda_scratch.zig). The 27B's recurrent scratch path keys on the
whole model, as one engine takes it.

Engine.initStage builds such an engine. Its forward is cuda_stage_forward.run: the boundary's h and r come in
(past rank 0), its layers run, h and r go out, and rank 0 then runs the final norm and head over the rows the
ring brings back. Shape.part keeps the layers and head pieces in their own captured graphs (the head piece at
key bit 31; a stage engine never runs the whole); the TCP hop is outside both. The first layer of a later stage
adds the incoming residual, as layer i of one engine does. A prompt chunk's tail (the last layer's kept row
alone) runs only on the stage that holds the model's last layer: an earlier stage sends every row on.
One engine's path is unchanged: step, window and prompt logits equal the parent's bit for bit.
---
 zig/build/cuda_family_host.zig                |  3 +
 zig/src/families/qwen3_5/cuda.zig             |  1 +
 zig/src/families/qwen3_5/cuda_engine.zig      | 59 +++++++++----------
 zig/src/families/qwen3_5/cuda_forward.zig     | 10 ++--
 zig/src/families/qwen3_5/cuda_geometry.zig    | 32 +++++++++-
 zig/src/families/qwen3_5/cuda_recurrent.zig   |  6 +-
 zig/src/families/qwen3_5/cuda_scratch.zig     |  7 ++-
 zig/src/families/qwen3_5/cuda_shapes.zig      | 26 +++++++-
 .../families/qwen3_5/cuda_stage_forward.zig   | 57 ++++++++++++++++++
 zig/src/families/qwen3_5/cuda_weights.zig     | 57 ++++++++++++++----
 10 files changed, 204 insertions(+), 54 deletions(-)
 create mode 100644 zig/src/families/qwen3_5/cuda_stage_forward.zig

diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig
index dc409a0355..f6e9cac3e2 100644
--- a/zig/build/cuda_family_host.zig
+++ b/zig/build/cuda_family_host.zig
@@ -83,6 +83,9 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B
     const prompt_moe = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"prompt MoE precision:"} }));
     b.step("test-qwen35-prompt-moe-host", "Prompt expert arithmetic stays separate from verify capacity without devices").dependOn(&prompt_moe.step);
     step.dependOn(&prompt_moe.step);
+    const stages = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen stage:"} }));
+    b.step("test-qwen35-stage-host", "Pipeline stage geometry: the stages' weights and state sum to one engine's").dependOn(&stages.step);
+    step.dependOn(&stages.step);
     const int8_contract = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_w4a8_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } } });
     int8_contract.addCSourceFile(.{ .file = b.path("zig/tests/cuda/w4a8_normalize_host.c"), .flags = &.{ "-std=c11", "-ffp-contract=off" } });
     int8_contract.addCSourceFile(.{ .file = b.path("zig/tests/cuda/w4a8_store_host.cpp"), .flags = &.{ "-std=c++17", "-fno-exceptions", "-fno-rtti" } });
diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig
index a4353843b7..387d838a53 100644
--- a/zig/src/families/qwen3_5/cuda.zig
+++ b/zig/src/families/qwen3_5/cuda.zig
@@ -60,6 +60,7 @@ test {
 
 test {
     _ = @import("cuda_wide.zig");
+    _ = @import("cuda_weights.zig");
 }
 
 pub const train_ops = @import("cuda_train_ops.zig");
diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig
index e3691c9c75..df672d15e3 100644
--- a/zig/src/families/qwen3_5/cuda_engine.zig
+++ b/zig/src/families/qwen3_5/cuda_engine.zig
@@ -14,6 +14,7 @@ const modelopt = @import("cuda_modelopt.zig");
 const x3w = @import("cuda_exl3_weights.zig");
 const hosted = @import("cuda_host_embed.zig");
 pub const recurrent = @import("cuda_recurrent.zig");
+pub const stage_fwd = @import("cuda_stage_forward.zig");
 
 pub const Geometry = geometry.Geometry;
 pub const Cache = caches.Cache;
@@ -170,6 +171,9 @@ pub const Engine = struct {
     n_taps: usize = 0,
     /// During an image prompt chunk's forward: its feature rows (0: none); the chunk's src and positions are staged.
     features: u64 = 0,
+    /// A pipeline stage's ring (null: one engine, or a stage warming up) and its boundary rows in host memory.
+    hop: ?stage_fwd.Hop = null,
+    boundary: ?cuda.HostBuffer = null,
 
     pub const ids_at = Round.ids_at;
     pub const segs_at = Round.segs_at;
@@ -185,9 +189,16 @@ pub const Engine = struct {
     const tables_len = Round.bytes;
 
     pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kv8: bool, fit: ?hosted.Fit) !*Engine {
+        return initStage(gpa, io, ctx, dir, kv8, fit, null);
+    }
+
+    /// One engine, or with `span` a pipeline stage holding only its layers' weights and state (RFC #2005).
+    pub fn initStage(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kv8: bool, fit: ?hosted.Fit, span: ?stage_fwd.Span) !*Engine {
         if (!cuda.kernels.available) return error.BuiltWithoutKernels;
         const started = profile.now(io);
-        const g = try Geometry.read(gpa, io, dir);
+        const whole = try Geometry.read(gpa, io, dir);
+        if (span) |sp| if (sp.heads != (sp.first == 0)) return error.StageHeadOffRankZero; // v1 is a ring
+        const g = if (span) |sp| try whole.stage(sp.first, sp.end, sp.heads) else whole;
         const d = ctx.d;
         const e = try @import("cuda_prompt_int8.zig").Allocation(Engine).create(gpa);
         errdefer gpa.destroy(e);
@@ -217,8 +228,12 @@ pub const Engine = struct {
         e.rec_scratch = recurrent.Pool.init(gpa, @intFromPtr(e));
         e.n_taps = 0;
         e.features = 0;
+        e.hop = null;
+        e.boundary = null;
         e.stream = try cuda.Stream.init(d, true);
         errdefer e.stream.deinit();
+        if (span != null) e.boundary = try cuda.HostBuffer.alloc(d, stage_fwd.bytes(g.hidden, max_rows));
+        errdefer if (e.boundary) |*b| b.free();
         var ck = try core.Checkpoint.openModel(gpa, io, dir);
         defer core.checkpoint.cache.close(&ck);
         e.arena = try hosted.arena(e, io, ctx, &ck, fit);
@@ -311,6 +326,7 @@ pub const Engine = struct {
         e.staged.free();
         e.tables.free();
         hosted.release(e);
+        if (e.boundary) |*b| b.free();
         e.gpa.free(e.layers);
         e.gpa.free(e.dir);
         e.affine.deinit();
@@ -335,6 +351,7 @@ pub const Engine = struct {
 
     /// The layers a drafter taps, told before the streams' caches exist; each forward then fills their rings.
     pub fn setTaps(e: *Engine, taps: []const u32) !void {
+        if (e.g.whole != 0 and taps.len > 0) return error.StageTaps; // v1 stages serve no drafter
         if (taps.len > e.taps.len) return error.TooManyTaps;
         for (taps) |t| if (t >= e.g.layers) return error.BadTap;
         @memcpy(e.taps[0..taps.len], taps);
@@ -422,8 +439,6 @@ pub const Engine = struct {
         const total = ids.len;
         const record = keep.verifies();
         const heads: usize = if (record) total else 1;
-        // the prompt GEMM only for prompt chunks: a window's rows keep the bits of one-row steps
-        const prompt = !record;
         // the tables
         const st = e.staged.bytes;
         @memcpy(std.mem.bytesAsSlice(u32, st[ids_at..][0 .. total * 4]), ids);
@@ -475,37 +490,16 @@ pub const Engine = struct {
         shape.hidden = keep == .hidden;
         shape.rows = keep == .rows;
         shape.paged = @import("cuda_attention.zig").hasPages(segs);
-        shape.mt2 = @import("cuda_shapes.zig").useMt2(e.mt2.enabled, e.kv8, prompt, total, parts.len, parts[0].cache.len);
+        shape.mt2 = @import("cuda_shapes.zig").useMt2(e.mt2.enabled, e.kv8, !record, total, parts.len, parts[0].cache.len);
         if (e.learning != null or e.pockets.any()) {
-            var exact = shape;
-            exact.total = @intCast(total);
-            exact.nseg = @intCast(parts.len);
-            e.features = if (images) |im| im.features else 0;
-            defer e.features = 0;
-            try e.body(exact);
-        } else if (prompt) {
-            // uncaptured: a prompt chunk's GEMMs dwarf its launches, and each length's graph would pay the stall
-            e.features = if (images) |im| im.features else 0;
-            defer e.features = 0;
-            try e.body(shape);
-        } else {
-            const gop = try e.graphs.getOrPut(e.gpa, shape.key());
-            if (!gop.found_existing) {
-                errdefer _ = e.graphs.remove(shape.key());
-                // captured, nothing runs: the replay below is this shape's first run
-                try cuda.graph.beginCapture(e.stream, .thread_local);
-                e.body(shape) catch |err| {
-                    var broken = cuda.graph.endCapture(e.stream) catch return err;
-                    broken.deinit();
-                    return err;
-                };
-                var captured = try cuda.graph.endCapture(e.stream);
-                defer captured.deinit();
-                gop.value_ptr.* = try captured.instantiate();
-            }
-            try gop.value_ptr.launchOn(e.stream);
+            shape.total = @intCast(total);
+            shape.nseg = @intCast(parts.len); // the exact path: uncaptured at the live counts
         }
+        e.features = if (images) |im| im.features else 0;
+        defer e.features = 0;
+        if (g.whole == 0) try e.launchShape(shape) else try stage_fwd.run(e, shape, total);
         for (parts) |p| p.cache.advanced(p.rows, record);
+        if (!g.heads) return recurrent.completed(e); // a stage without the head: no logits to keep
         if (keep == .draws or keep == .logits) try e.penalties.apply(e.d, &e.draws, e.stream, e.logits, g.vocab) else e.penalties.clear();
         switch (keep) {
             .hidden => try e.stream.synchronize(), // mapped layer inputs are ready before the learner reads them
@@ -539,6 +533,9 @@ pub const Engine = struct {
 
     pub const Shape = @import("cuda_shapes.zig").Shape;
 
+    /// One shape's launches: uncaptured for a prompt chunk or the exact path, else its captured graph replayed.
+    pub const launchShape = @import("cuda_shapes.zig").launch;
+
     /// The same scan selected by body, exposed for host launch-contract fixtures.
     pub fn deltaScan(e: *Engine, record: bool, slot_count: u8, segments: usize, layer: c_int) !void {
         const scan = if (!record) e.fns.gdn_chain else switch (slot_count) {
diff --git a/zig/src/families/qwen3_5/cuda_forward.zig b/zig/src/families/qwen3_5/cuda_forward.zig
index 0de40f97f0..5e34e6f8cc 100644
--- a/zig/src/families/qwen3_5/cuda_forward.zig
+++ b/zig/src/families/qwen3_5/cuda_forward.zig
@@ -38,9 +38,10 @@ pub fn body(e: *E.Engine, shape: anytype) !void {
     const t_params = e.tables.ptr + params_at;
     const geo = e.geo;
     const group = g.q_heads / g.kv_heads;
+    if (shape.part == .head) return finish(e, total, heads, shape.hidden, shape.rows); // a stage's head over the ring's rows
 
     if (prompt and e.int8.on) @import("cuda_prompt_int8.zig").chunkStart(e);
-    try e.embed(t_ids, e.h, total, t_params);
+    if (g.embeds) try e.embed(t_ids, e.h, total, t_params); // a later stage's h and r came over the ring
     // the producers of the out, o and down inputs write their group sums; a prompt GEMM past 16 rows needs none
     const sums: u64 = if (prompt and total > 16) 0 else e.xs;
     // an image chunk's rows take the tower's features, and rotate at their staged t, h, w positions
@@ -49,12 +50,13 @@ pub fn body(e: *E.Engine, shape: anytype) !void {
     var di: usize = 0;
     var ai: c_int = 0;
     // a prompt chunk keeps its last row: past the last layer's keys and values, the kept row's tile runs alone
-    const tail_at: ?usize = if (prompt) @import("cuda_prompt_tail.zig").first(e, total, nseg, shape.hidden or shape.rows) else null;
+    const ends = g.whole == 0 or g.first + g.layers == g.whole; // an earlier stage sends every row on: no tail
+    const tail_at: ?usize = if (prompt and ends) @import("cuda_prompt_tail.zig").first(e, total, nseg, shape.hidden or shape.rows) else null;
     for (e.layers, 0..) |l, i| {
         const tail = tail_at != null and i + 1 == e.layers.len;
         const t0: usize = if (tail) tail_at.? else 0;
         const rows: usize = total - t0;
-        try e.norm(i != 0, l.input_norm, total);
+        try e.norm(g.first + i != 0, l.input_norm, total);
         if (e.learning) |k| try k.input(i, e.h, total);
         switch (l.mixer) {
             .delta => |m| {
@@ -140,7 +142,7 @@ pub fn body(e: *E.Engine, shape: anytype) !void {
         };
     }
     if (prompt and e.int8.on) @import("cuda_prompt_int8.zig").chunkEnd(e, total);
-    try finish(e, total, heads, shape.hidden, shape.rows);
+    if (shape.part == .whole) try finish(e, total, heads, shape.hidden, shape.rows);
 }
 
 /// A launch whose arguments are checked against the kernel's .cu signature: a wrong count or kind does not build.
diff --git a/zig/src/families/qwen3_5/cuda_geometry.zig b/zig/src/families/qwen3_5/cuda_geometry.zig
index f2b1088407..bf19081f65 100644
--- a/zig/src/families/qwen3_5/cuda_geometry.zig
+++ b/zig/src/families/qwen3_5/cuda_geometry.zig
@@ -92,6 +92,12 @@ pub const Geometry = struct {
     exl3: bool = false,
     /// A plain BF16 checkpoint (no quantization at all, HF tensor names): every projection bf16, decisions only.
     bf16: bool = false,
+    /// A pipeline stage's layers: `layers` of them from global layer `first`, of a model of `whole` (0: this one).
+    first: usize = 0,
+    whole: usize = 0,
+    /// Whether the stage holds the embedding, and the final norm and the head (a whole model holds both).
+    embeds: bool = true,
+    heads: bool = true,
 
     pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Geometry {
         const path = try std.fs.path.join(gpa, &.{ dir, "config.json" });
@@ -126,12 +132,36 @@ pub const Geometry = struct {
         return g.q_heads * head_dim;
     }
 
+    /// The stage's linear layers (`linear` is indexed by global layer).
     pub fn deltas(g: Geometry) usize {
         var n: usize = 0;
-        for (g.linear[0..g.layers]) |l| n += @intFromBool(l);
+        for (g.linear[g.first..][0..g.layers]) |l| n += @intFromBool(l);
         return n;
     }
 
+    /// The whole model's geometry, of which this may be a stage.
+    pub fn model(g: Geometry) Geometry {
+        var m = g;
+        m.layers = if (g.whole == 0) g.layers else g.whole;
+        m.first = 0;
+        m.whole = 0;
+        m.embeds = true;
+        m.heads = true;
+        return m;
+    }
+
+    /// Layers [first, end) of this whole model as a stage; it holds the embedding when it starts at layer 0.
+    pub fn stage(g: Geometry, first: usize, end: usize, heads: bool) error{BadStageRange}!Geometry {
+        if (g.whole != 0 or first >= end or end > g.layers) return error.BadStageRange;
+        var s = g;
+        s.first = first;
+        s.layers = end - first;
+        s.whole = g.layers;
+        s.embeds = first == 0;
+        s.heads = heads;
+        return s;
+    }
+
     pub fn attentions(g: Geometry) usize {
         return g.layers - g.deltas();
     }
diff --git a/zig/src/families/qwen3_5/cuda_recurrent.zig b/zig/src/families/qwen3_5/cuda_recurrent.zig
index 7361f29807..a7513684a0 100644
--- a/zig/src/families/qwen3_5/cuda_recurrent.zig
+++ b/zig/src/families/qwen3_5/cuda_recurrent.zig
@@ -6,9 +6,11 @@ const Cache = @import("cuda_cache.zig").Cache;
 pub const Pool = common.Workspace(cuda.DeviceBuffer);
 pub const Lease = Pool.Lease;
 
+/// The whole model's shape decides, so a pipeline stage keeps the scratch path one engine takes.
 pub fn enabled(e: *const Engine) bool {
-    return e.g.hidden == 5120 and e.g.layers == 64 and e.g.experts == 0 and
-        e.g.v_heads == 48 and e.g.k_heads == 16 and e.g.deltas() == 48;
+    const g = e.g.model();
+    return g.hidden == 5120 and g.layers == 64 and g.experts == 0 and
+        g.v_heads == 48 and g.k_heads == 16 and g.deltas() == 48;
 }
 
 pub const Provider = struct {
diff --git a/zig/src/families/qwen3_5/cuda_scratch.zig b/zig/src/families/qwen3_5/cuda_scratch.zig
index 0972958635..92820b7460 100644
--- a/zig/src/families/qwen3_5/cuda_scratch.zig
+++ b/zig/src/families/qwen3_5/cuda_scratch.zig
@@ -68,9 +68,10 @@ pub fn scratch(e: *E.Engine, mode: Mode) !usize {
     e.pm = try t.alloc(partials * 4);
     e.pl = try t.alloc(partials * 4);
     e.po = try t.alloc(partials * geometry.head_dim * 4);
-    e.logits = try t.alloc(E.most_verify * g.vocab * 2);
+    // a stage without the head keeps no logits
+    e.logits = if (g.heads) try t.alloc(E.most_verify * g.vocab * 2) else 0;
     if (g.moe()) e.routed = try routed.Scratch.carve(&t, g, R);
-    e.amax = try t.alloc(E.most_verify * E.argmaxSlices(g.vocab) * 8);
-    e.picks = try t.alloc(E.most_verify * cuda.sample.Draws.row_scratch);
+    e.amax = if (g.heads) try t.alloc(E.most_verify * E.argmaxSlices(g.vocab) * 8) else 0;
+    e.picks = if (g.heads) try t.alloc(E.most_verify * cuda.sample.Draws.row_scratch) else 0;
     return t.total;
 }
diff --git a/zig/src/families/qwen3_5/cuda_shapes.zig b/zig/src/families/qwen3_5/cuda_shapes.zig
index a9ca7e76da..445493a2c0 100644
--- a/zig/src/families/qwen3_5/cuda_shapes.zig
+++ b/zig/src/families/qwen3_5/cuda_shapes.zig
@@ -1,5 +1,6 @@
 //! Captured launch geometry, separate from the live counts read through the staged device tables.
 const std = @import("std");
+const cuda = @import("cuda");
 const Round = @import("cuda_round.zig");
 
 pub const Shape = struct {
@@ -13,6 +14,8 @@ pub const Shape = struct {
     rows: bool = false, // a readout piece: every row through the last layer and the final norm, no head
     paged: bool = false,
     mt2: bool = false,
+    /// A pipeline stage's piece of the forward (cuda_stage_forward.zig); one engine runs the whole.
+    part: @import("cuda_stage_forward.zig").Part = .whole,
 
     pub fn init(rows: usize, segments: usize, record: bool, taps: bool, slots: u8) Shape {
         std.debug.assert(rows > 0 and rows <= @as(usize, if (record) Round.most_verify else Round.max_rows) and segments > 0 and segments <= rows);
@@ -23,7 +26,7 @@ pub const Shape = struct {
 
     pub fn key(s: Shape) u32 {
         // captured shapes only: a window bucket's rows fit the low byte
-        return @as(u32, @as(u8, @intCast(s.total))) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18 | @as(u32, s.slots) << 19 | @as(u32, @intFromBool(s.hidden)) << 27 | @as(u32, @intFromBool(s.paged)) << 28 | @as(u32, @intFromBool(s.mt2)) << 29 | @as(u32, @intFromBool(s.rows)) << 30;
+        return @as(u32, @as(u8, @intCast(s.total))) | @as(u32, s.nseg) << 8 | @as(u32, @intFromBool(s.record)) << 16 | @as(u32, @intFromBool(s.prompt)) << 17 | @as(u32, @intFromBool(s.taps)) << 18 | @as(u32, s.slots) << 19 | @as(u32, @intFromBool(s.hidden)) << 27 | @as(u32, @intFromBool(s.paged)) << 28 | @as(u32, @intFromBool(s.mt2)) << 29 | @as(u32, @intFromBool(s.rows)) << 30 | @as(u32, @intFromBool(s.part == .head)) << 31;
     }
 };
 
@@ -71,3 +74,24 @@ test "every captured verify shape keys its own graph: a shape after another neve
     try std.testing.expectEqual(distinct.count(), keys.count());
     try std.testing.expect(keys.count() > 1000);
 }
+
+/// One shape's launches (the engine's `launchShape`): uncaptured for a prompt chunk or the exact path, else replayed.
+pub fn launch(e: anytype, shape: Shape) !void {
+    // a prompt chunk's GEMMs (a window's rows keep one-row steps' bits) dwarf its launches: never captured
+    if (e.learning != null or e.pockets.any() or shape.prompt) return e.body(shape);
+    const gop = try e.graphs.getOrPut(e.gpa, shape.key());
+    if (!gop.found_existing) {
+        errdefer _ = e.graphs.remove(shape.key());
+        // captured, nothing runs: the replay below is this shape's first run
+        try cuda.graph.beginCapture(e.stream, .thread_local);
+        e.body(shape) catch |err| {
+            var broken = cuda.graph.endCapture(e.stream) catch return err;
+            broken.deinit();
+            return err;
+        };
+        var captured = try cuda.graph.endCapture(e.stream);
+        defer captured.deinit();
+        gop.value_ptr.* = try captured.instantiate();
+    }
+    try gop.value_ptr.launchOn(e.stream);
+}
diff --git a/zig/src/families/qwen3_5/cuda_stage_forward.zig b/zig/src/families/qwen3_5/cuda_stage_forward.zig
new file mode 100644
index 0000000000..74ac552271
--- /dev/null
+++ b/zig/src/families/qwen3_5/cuda_stage_forward.zig
@@ -0,0 +1,57 @@
+//! A pipeline stage's forward (RFC #2005): its own layers between the boundary rows h and r, and the head on rank 0.
+const std = @import("std");
+const cuda = @import("cuda");
+const E = @import("cuda_engine.zig");
+
+/// A stage's layers [first, end) of the model; `heads`: it also runs the final norm and the head (rank 0 of a ring).
+pub const Span = struct { first: usize, end: usize, heads: bool };
+
+/// A forward's launches: all (one engine), a stage's layers (rank 0: from the embedding), or the norm and head.
+pub const Part = enum(u2) { whole, layers, head };
+
+/// The boundary's two planes (h, then r) of `rows` rows in host memory, sent or filled in place by the ring.
+pub const Hop = struct {
+    ctx: *anyopaque,
+    give: *const fn (ctx: *anyopaque, rows: usize, bytes: []const u8) anyerror!void,
+    take: *const fn (ctx: *anyopaque, rows: usize, bytes: []u8) anyerror!void,
+};
+
+/// Host bytes of a boundary of `rows` rows: two bf16 planes of `hidden` values a row.
+pub fn bytes(hidden: usize, rows: usize) usize {
+    return 2 * rows * hidden * 2;
+}
+
+/// A stage's launches: rows in (past rank 0), its layers, rows out, rank 0 then its head; `rows` unbucketed.
+pub fn run(e: *E.Engine, shape: E.Engine.Shape, rows: usize) !void {
+    const g = e.g;
+    var s = shape;
+    s.part = .layers;
+    if (!g.embeds) try take(e, rows);
+    try e.launchShape(s);
+    try give(e, rows);
+    if (!g.heads) return;
+    try take(e, rows);
+    s.part = .head;
+    try e.launchShape(s);
+}
+
+/// The boundary rows from the stage before, onto h and r before the next launch (without a ring: as they are).
+fn take(e: *E.Engine, rows: usize) !void {
+    const hop = e.hop orelse return;
+    const buf = e.boundary.?.bytes[0..bytes(e.g.hidden, rows)];
+    try hop.take(hop.ctx, rows, buf);
+    const plane = buf.len / 2;
+    try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(e.h, buf.ptr, plane, e.stream.handle), "cuMemcpyHtoDAsync");
+    try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(e.r, buf[plane..].ptr, plane, e.stream.handle), "cuMemcpyHtoDAsync");
+}
+
+/// This stage's h and r to the next one, once the layers that wrote them are done.
+fn give(e: *E.Engine, rows: usize) !void {
+    const hop = e.hop orelse return;
+    const buf = e.boundary.?.bytes[0..bytes(e.g.hidden, rows)];
+    const plane = buf.len / 2;
+    try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(buf.ptr, e.h, plane, e.stream.handle), "cuMemcpyDtoHAsync");
+    try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(buf[plane..].ptr, e.r, plane, e.stream.handle), "cuMemcpyDtoHAsync");
+    try e.stream.synchronize();
+    try hop.give(hop.ctx, rows, buf);
+}
diff --git a/zig/src/families/qwen3_5/cuda_weights.zig b/zig/src/families/qwen3_5/cuda_weights.zig
index 33e6b73170..c6b29638d8 100644
--- a/zig/src/families/qwen3_5/cuda_weights.zig
+++ b/zig/src/families/qwen3_5/cuda_weights.zig
@@ -81,10 +81,11 @@ fn raw3Bytes(count: usize, n: usize, k: usize) usize {
 /// What `load` takes from the arena: every packed projection and vector, each allocation aligned.
 pub fn total(g: Geometry) usize {
     const mo = g.kinds.modelopt;
-    var t: usize = if (mo) aligned(g.vocab * g.hidden * 2) else embedBytes(g);
-    t += if (mo) modelopt.packedBytes(g.vocab, g.hidden, g.kinds.head) else packedBytes(g.vocab, g.hidden, g.bits.head); // the head
-    t += aligned(g.hidden * 2);
-    for (0..g.layers) |i| {
+    var t: usize = 0;
+    if (g.embeds) t += if (mo) aligned(g.vocab * g.hidden * 2) else embedBytes(g);
+    if (g.heads) t += if (mo) modelopt.packedBytes(g.vocab, g.hidden, g.kinds.head) else packedBytes(g.vocab, g.hidden, g.bits.head);
+    if (g.heads) t += aligned(g.hidden * 2); // the final norm
+    for (g.first..g.first + g.layers) |i| {
         t += 2 * aligned(g.hidden * 2);
         if (g.moe()) t += if (mo) modelopt_experts.routedBytes(g) else routedBytes(g.experts, g.intermediate, g.hidden, g.expert_bits);
         for (slots(g, i)) |slot| {
@@ -104,15 +105,15 @@ pub fn total(g: Geometry) usize {
 pub fn stageBytes(g: Geometry) usize {
     if (g.kinds.modelopt) {
         var raw = modelopt.rawBytes(g.vocab, g.hidden, g.kinds.head);
-        for (0..g.layers) |i| for (slots(g, i)) |slot| {
+        for (g.first..g.first + g.layers) |i| for (slots(g, i)) |slot| {
             const nk = shape(g, slot);
             raw = @max(raw, modelopt.rawBytes(nk[0], nk[1], g.kinds.of(i, slot)));
         };
         return if (g.moe()) @max(raw, modelopt_experts.stageBytes(g)) else raw;
     }
-    var most = rawBytes(g.vocab, g.hidden, g.bits.head);
-    if (rowBits(g.bits.embed) != g.bits.embed) most = @max(most, rawBytes(g.vocab, g.hidden, g.bits.embed));
-    for (0..g.layers) |i| for (slots(g, i)) |slot| {
+    var most: usize = if (g.heads) rawBytes(g.vocab, g.hidden, g.bits.head) else 0;
+    if (g.embeds and rowBits(g.bits.embed) != g.bits.embed) most = @max(most, rawBytes(g.vocab, g.hidden, g.bits.embed));
+    for (g.first..g.first + g.layers) |i| for (slots(g, i)) |slot| {
         const nk = shape(g, slot);
         most = @max(most, rawBytes(nk[0], nk[1], g.bits.of(i, slot)));
     };
@@ -231,6 +232,25 @@ pub fn linearSlice(e: *Engine, into: anytype, ck: *core.Checkpoint, stage: *Devi
     return q;
 }
 
+/// A pipeline stage's ends (an MLX checkpoint's): the embedding from layer 0, the head where the stage holds it.
+fn loadStage(e: *Engine, ck: *core.Checkpoint, stage: *DeviceBuffer, m: []const u8) !void {
+    const g = e.g;
+    if (g.kinds.modelopt or g.exl3 or g.bf16 or g.pocket > 0) return error.StageNeedsMlxCheckpoint;
+    e.embed_bits = rowBits(g.bits.embed);
+    e.embed_w = 0;
+    e.embed_s = 0;
+    e.embed_b = 0;
+    if (g.embeds) {
+        e.embed_w = try embedWords(e, ck, stage, m);
+        e.embed_s = try tensor(e, ck, "{s}embed_tokens.scales", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 });
+        e.embed_b = try tensor(e, ck, "{s}embed_tokens.biases", .{m}, .bf16, &.{ g.vocab, g.hidden / 64 });
+    }
+    e.final_norm = 0;
+    e.head = .{ .a4 = std.mem.zeroes(Affine4) }; // never read: a stage without the head runs no finish
+    if (!g.heads) return;
+    e.head = if (g.tied) try linear(e, e, ck, stage, m, "embed_tokens", g.vocab, g.hidden, g.bits.head) else try linear(e, e, ck, stage, "language_model.", "lm_head", g.vocab, g.hidden, g.bits.head);
+}
+
 fn bf16f(h: u16) f32 {
     return @bitCast(@as(u32, h) << 16);
 }
@@ -337,7 +357,7 @@ pub fn load(e: *Engine, ck: *core.Checkpoint, stage: *DeviceBuffer) !void {
     const g = e.g;
     const mo = g.kinds.modelopt;
     const m = if (mo) modelopt.prefix else "language_model.model.";
-    if (mo) {
+    if (g.whole != 0) try loadStage(e, ck, stage, m) else if (mo) {
         e.embed_bits = 16; // bf16 rows as stored
         e.embed_w = try tensor(e, ck, "{s}embed_tokens.weight", .{m}, .bf16, &.{ g.vocab, g.hidden });
         e.embed_s = 0;
@@ -352,17 +372,17 @@ pub fn load(e: *Engine, ck: *core.Checkpoint, stage: *DeviceBuffer) !void {
     }
     // a plain BF16 checkpoint keeps HF's raw w: its norm kernels (the _shift ones) scale by 1 + w in FP32
     const shift = mo and !g.bf16;
-    e.final_norm = try normVec(e, ck, shift, "{s}norm.weight", .{m}, .bf16, &.{g.hidden});
+    if (g.heads) e.final_norm = try normVec(e, ck, shift, "{s}norm.weight", .{m}, .bf16, &.{g.hidden});
     const conv = g.convWidth();
     const taps_shape: []const usize = if (mo) &.{ conv, 1, geometry.taps } else &.{ conv, geometry.taps, 1 }; // the same bytes
     const readers = try core.checkpoint.loadReaders();
     // each layer read past the page cache, the next one while it packs (TF_LOAD_DIRECT=0: #1928's read-ahead)
     var direct = core.checkpoint.direct.Layers.init(e.gpa, ck, readers);
     defer direct.deinit();
-    for (e.layers, 0..) |*l, i| {
+    for (e.layers, g.first..) |*l, i| {
         var pbuf: [2][96]u8 = undefined;
         const p = try std.fmt.bufPrint(&pbuf[0], "{s}layers.{d}.", .{ m, i });
-        const last = i + 1 == e.layers.len;
+        const last = i + 1 == g.first + e.layers.len;
         try direct.enter(p, if (last) null else try std.fmt.bufPrint(&pbuf[1], "{s}layers.{d}.", .{ m, i + 1 }));
         // a read-ahead layer is on the device: its pages leave the cache (none: main's one release at close)
         defer if (readers > 0 and !direct.on()) ck.release();
@@ -426,6 +446,19 @@ test "a uniform 4-bit checkpoint keeps its sizes; other modules take their own w
     try std.testing.expectEqual(rawBytes(g.vocab, g.hidden, 6), stageBytes(g));
 }
 
+test "Qwen stage: a pipeline's stages load one engine's weights between them, the embedding and the head once" {
+    var g: Geometry = .{ .hidden = 2048, .intermediate = 6144, .layers = 8, .vocab = 248320, .k_heads = 16, .v_heads = 16, .q_heads = 8, .kv_heads = 2, .eps = 1e-6, .tied = true, .context = 4096, .linear = @splat(false) };
+    for (0..8) |i| g.linear[i] = i % 4 != 3;
+    const a = try g.stage(0, 3, true);
+    const b = try g.stage(3, 8, false);
+    try std.testing.expectEqual(total(g), total(a) + total(b));
+    try std.testing.expectEqual(g.stateBytes(), a.stateBytes() + b.stateBytes());
+    try std.testing.expectEqual(g.kvBytes(64, false), a.kvBytes(64, false) + b.kvBytes(64, false));
+    try std.testing.expect(a.embeds and !b.embeds and !b.heads and b.model().layers == 8 and b.deltas() == 3);
+    try std.testing.expectError(error.BadStageRange, g.stage(3, 9, false));
+    try std.testing.expectError(error.BadStageRange, b.stage(0, 1, false));
+}
+
 test {
     _ = @import("cuda_expert_pack.zig");
 }

From 7b4080396550dd4733c90f053989efcdcdbbdc9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Wed, 4 Nov 2026 15:57:18 +0100
Subject: [PATCH 2/4] qwen3_5 cuda: a Qwen engine as a stage of the ring, boot
 to rounds, and tf-qwen35-cuda stage

cuda_stage.zig drives the stage link (#2012) with the protocol (#2060): at boot each stage's description and
ready report go round the ring and rank 0 checks every one and the layout; each round rank 0's plan goes round
(rank 0 checks the last stage's echo byte for byte), then the rows, h and r as bf16 planes. Slots is the same
code on every rank: open takes a free slot at generation 0, continue keeps the previous window's first
`accepted` rows (a keep advances the generation) and checks the plan's generation and position, close frees it;
resume is refused until the prompt cache spans stages. Rank 0's Lead predicts each entry's position and
generation, sends the plan, then applies it through Slots like every follower.

tf-qwen35-cuda <model> stage <out> runs one scenario (a two-chunk prompt, greedy steps, a second stream beside
it, drafted windows kept to 3, 5, 0 and 2 rows, a keep-alive) on one engine or a ring. Qwen3.8-27B MLX 4-bit on
one GB10: 2 stages (0:32, 32:64) and 3 (0:20, 20:44, 44:64) give the one engine's 48 logits rows byte for byte,
bf16 and fp8 KV; device bytes 16.98 GiB whole, 9.04 + 7.64 GiB in two, 6.61 + 6.03 + 5.22 GiB in three. A
follower killed mid-scenario ends rank 0 with Closed 0.4 s later.
---
 zig/build/cuda.zig                      |   2 +
 zig/build/cuda_census.zig               |   2 +-
 zig/src/families/qwen3_5/cuda.zig       |   1 +
 zig/src/families/qwen3_5/cuda_stage.zig | 325 ++++++++++++++++++++++++
 zig/tests/qwen35_cuda_main.zig          |  12 +
 zig/tests/qwen35_stage.zig              | 157 ++++++++++++
 6 files changed, 498 insertions(+), 1 deletion(-)
 create mode 100644 zig/src/families/qwen3_5/cuda_stage.zig
 create mode 100644 zig/tests/qwen35_stage.zig

diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig
index cf9ec42a88..b151343b52 100644
--- a/zig/build/cuda.zig
+++ b/zig/build/cuda.zig
@@ -38,6 +38,7 @@ pub fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bui
     qwen35.addImport("cuda", cuda);
     qwen35.addImport("core", core);
     qwen35.addImport("lanes", lanes);
+    qwen35.addImport("stage", @import("stage.zig").module(b, target, optimize));
     qwen35.addImport("engine_api", b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} })); // core's prompt-cache types, the server's engine_api
     const qwen_vision = qwenVisionModule(b, target, optimize, cuda, core, lanes);
     qwen35.addImport("qwen_vision", qwen_vision);
@@ -133,6 +134,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     q35.addImport("core", mods.core);
     q35.addImport("qwen35", mods.qwen35);
     q35.addImport("lanes", mods.lanes);
+    q35.addImport("stage", mods.qwen35.import_table.get("stage").?);
     const q35_install = b.addInstallArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 }), .{});
     b.getInstallStep().dependOn(&q35_install.step);
     b.step("qwen35-cuda", "Build the Qwen CUDA forward and raw state receipt tool").dependOn(&q35_install.step);
diff --git a/zig/build/cuda_census.zig b/zig/build/cuda_census.zig
index 8944b4b7ee..bcb733c618 100644
--- a/zig/build/cuda_census.zig
+++ b/zig/build/cuda_census.zig
@@ -2,7 +2,7 @@
 const std = @import("std");
 
 pub fn steps(b: *std.Build, cuda_target: std.Build.ResolvedTarget, check_cuda: *std.Build.Module, qwen35: *std.Build.Module, core: *std.Build.Module, flashnext: *std.Build.Module) void {
-    const rounds_tool = b.createModule(.{ .root_source_file = b.path("zig/tests/qwen35_cuda_main.zig"), .target = cuda_target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = check_cuda }, .{ .name = "qwen35", .module = qwen35 }, .{ .name = "core", .module = core } } });
+    const rounds_tool = b.createModule(.{ .root_source_file = b.path("zig/tests/qwen35_cuda_main.zig"), .target = cuda_target, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = check_cuda }, .{ .name = "qwen35", .module = qwen35 }, .{ .name = "core", .module = core }, .{ .name = "stage", .module = qwen35.import_table.get("stage").? } } });
     const rounds_compile = b.addObject(.{ .name = "qwen35-rounds-checker-compile", .root_module = b.createModule(.{ .root_source_file = b.path("zig/qwen35_rounds_compile.zig"), .target = cuda_target, .link_libc = true, .imports = &.{.{ .name = "tool", .module = rounds_tool }} }) });
     _ = rounds_compile.getEmittedBin();
     b.step("compile-qwen35-rounds-host", "Compile the real CUDA 64-row checker for Linux, no images or execution").dependOn(&rounds_compile.step);
diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig
index 387d838a53..33b6ff069e 100644
--- a/zig/src/families/qwen3_5/cuda.zig
+++ b/zig/src/families/qwen3_5/cuda.zig
@@ -15,6 +15,7 @@ pub const mtp_check = @import("cuda_mtp_check.zig");
 pub const draft = @import("cuda_draft.zig");
 pub const vision = @import("cuda_vision.zig");
 pub const reuse = @import("cuda_reuse.zig");
+pub const stage = @import("cuda_stage.zig");
 
 test {
     _ = @import("cuda_attention.zig");
diff --git a/zig/src/families/qwen3_5/cuda_stage.zig b/zig/src/families/qwen3_5/cuda_stage.zig
new file mode 100644
index 0000000000..0a32857e36
--- /dev/null
+++ b/zig/src/families/qwen3_5/cuda_stage.zig
@@ -0,0 +1,325 @@
+//! A Qwen CUDA engine as one stage of the ring (RFC #2005 v1): its boot, its streams' slots, rank 0's rounds.
+const std = @import("std");
+const stage = @import("stage");
+const E = @import("cuda_engine.zig");
+const sf = @import("cuda_stage_forward.zig");
+
+const plan = stage.plan;
+const boot = stage.boot;
+
+/// Every stage of one placement reports this; rank 0 refuses another.
+pub const version = "tensorfold qwen3_5-cuda stage-v1";
+
+pub const Error = error{ StreamOpen, NoFreeSlot, UnknownStream, BadAccepted, OtherPosition, ResumeNotInV1, CutsNotInV1, OtherEcho, PromptParts };
+
+/// The ring's boundary for this engine: h and r as bf16 planes, no drafter taps (v1).
+pub fn boundary(e: *const E.Engine) boot.Boundary {
+    return .{ .hidden = @intCast(e.g.hidden), .row_dtype = .bf16, .planes = 2, .draft_width = 0 };
+}
+
+/// The shapes every stage runs: a prompt chunk and a round at most `max_rows` rows, `parallel` slots of `context`.
+pub fn shapes(parallel: u32, context: u32) boot.Shapes {
+    return .{ .prompt_chunk = E.max_rows, .max_rows = E.max_rows, .parallel = parallel, .context = context };
+}
+
+/// The checkpoint's and layout's fingerprint: SHA-256 over config.json, the version and every rank's range.
+pub fn fingerprint(config: []const u8, ranges: []const boot.Range) [32]u8 {
+    var h = std.crypto.hash.sha2.Sha256.init(.{});
+    h.update(config);
+    h.update(version);
+    for (ranges) |r| {
+        var b: [8]u8 = undefined;
+        std.mem.writeInt(u32, b[0..4], r.first, .little);
+        std.mem.writeInt(u32, b[4..8], r.end, .little);
+        h.update(&b);
+    }
+    return h.finalResult();
+}
+
+/// One stream's state on this stage: its cache (its own layers only) and its state generation.
+pub const Slot = struct { cache: E.Cache, id: u32 = 0, used: bool = false, generation: u32 = 0 };
+
+/// A stage's streams: the same code on rank 0 and on every follower, so each applies a plan's keeps alike.
+pub const Slots = struct {
+    e: *E.Engine,
+    all: []Slot,
+
+    pub fn init(gpa: std.mem.Allocator, e: *E.Engine, n: usize, context: usize) !Slots {
+        const all = try gpa.alloc(Slot, n);
+        var made: usize = 0;
+        errdefer {
+            for (all[0..made]) |*s| s.cache.deinit(gpa);
+            gpa.free(all);
+        }
+        for (all) |*s| {
+            s.* = .{ .cache = try E.Cache.init(gpa, e, context) };
+            made += 1;
+        }
+        return .{ .e = e, .all = all };
+    }
+
+    pub fn deinit(s: *Slots, gpa: std.mem.Allocator) void {
+        for (s.all) |*x| x.cache.deinit(gpa);
+        gpa.free(s.all);
+    }
+
+    pub fn find(s: *Slots, id: u32) ?*Slot {
+        for (s.all) |*x| if (x.used and x.id == id) return x;
+        return null;
+    }
+
+    /// The stream's previous window kept to its first `accepted` rows (a keep advances the generation).
+    fn settle(s: *Slots, x: *Slot, accepted: u32) !void {
+        const last = x.cache.last orelse return if (accepted != 0) error.BadAccepted;
+        if (accepted > last.rows) return error.BadAccepted;
+        var path: [E.most_window]u32 = undefined;
+        for (path[0..accepted], 0..) |*r, i| r.* = @intCast(i);
+        try x.cache.keep(s.e, path[0..accepted]);
+        x.generation +%= 1;
+    }
+
+    /// One stream entry before the round's forward: opened, kept and checked, or closed; its slot when it has rows.
+    pub fn apply(s: *Slots, st: plan.Stream) !?*Slot {
+        switch (st.op) {
+            .open => {
+                if (s.find(st.id) != null) return error.StreamOpen;
+                const x = for (s.all) |*x| {
+                    if (!x.used) break x;
+                } else return error.NoFreeSlot;
+                try x.cache.reset(s.e);
+                x.* = .{ .cache = x.cache, .id = st.id, .used = true };
+                try plan.sameGeneration(x.generation, st);
+                return x;
+            },
+            .@"continue" => {
+                const x = s.find(st.id) orelse return error.UnknownStream;
+                try s.settle(x, st.accepted);
+                try plan.sameGeneration(x.generation, st);
+                if (x.cache.len != st.first) return error.OtherPosition;
+                return if (st.rows > 0) x else null;
+            },
+            .close, .cancel => {
+                const x = s.find(st.id) orelse return error.UnknownStream;
+                try plan.sameGeneration(x.generation, st);
+                x.used = false;
+                return null;
+            },
+            .@"resume" => return error.ResumeNotInV1, // a prompt cache across stages is step 7
+        }
+    }
+};
+
+/// What rank 0's caller asks of a stream in one round.
+pub const Entry = struct { id: u32, op: plan.Op, rows: u32 = 0, accepted: u32 = 0 };
+
+/// A plan's forward parts, in stream order.
+const Round = struct {
+    parts: [E.max_segs]E.Part = undefined,
+    n: usize = 0,
+};
+
+/// The ring's hop for one round: rows to the next rank, rows from the one before (on rank 0, after the plan's echo).
+const Link = struct {
+    ring: *stage.Ring,
+    round: u64,
+    bound: boot.Boundary,
+    plan: plan.Plan,
+    /// Rank 0: the plan it sent, which the last stage's echo must equal.
+    sent: ?[]const u8 = null,
+    echo: []u8 = &.{},
+
+    fn hop(l: *Link) sf.Hop {
+        return .{ .ctx = l, .give = give, .take = take };
+    }
+
+    fn give(ctx: *anyopaque, rows: usize, bytes: []const u8) anyerror!void {
+        const l: *Link = @ptrCast(@alignCast(ctx));
+        try l.ring.send(.rows, l.round, @intCast(rows), &.{bytes});
+    }
+
+    fn take(ctx: *anyopaque, rows: usize, bytes: []u8) anyerror!void {
+        const l: *Link = @ptrCast(@alignCast(ctx));
+        if (l.sent) |p| {
+            try echoed(l.ring, l.round, p, l.echo);
+            l.sent = null;
+        }
+        const h = try l.ring.recv(.rows, l.round, bytes);
+        const r = try plan.checkRows(h, l.plan, l.bound);
+        if (r.rows != rows) return error.OtherRowCount;
+    }
+};
+
+/// The last stage's copy of rank 0's plan, which must be byte for byte the one sent.
+fn echoed(ring: *stage.Ring, round: u64, sent: []const u8, buf: []u8) !void {
+    const h = try ring.recv(.plan, round, buf);
+    if (!std.mem.eql(u8, buf[0..@intCast(h.bytes)], sent)) return error.OtherEcho;
+}
+
+/// Rank 0: the plan of each round to the ring, its own streams applied alike, then the forward with the hop.
+pub const Lead = struct {
+    e: *E.Engine,
+    /// null: one engine, no ring (the whole-model reference runs the same rounds).
+    ring: ?*stage.Ring,
+    slots: Slots,
+    limits: plan.Limits,
+    round: u64,
+    mem: []u8,
+
+    pub fn init(gpa: std.mem.Allocator, e: *E.Engine, ring: ?*stage.Ring, parallel: u32, context: u32, first_round: u64) !Lead {
+        if (parallel > E.max_segs) return error.TooManyStreams;
+        var slots = try Slots.init(gpa, e, parallel, context);
+        errdefer slots.deinit(gpa);
+        const mem = try gpa.alloc(u8, 2 * plan.planSize(parallel, 0));
+        return .{ .e = e, .ring = ring, .slots = slots, .limits = plan.Limits.of(shapes(parallel, context)), .round = first_round, .mem = mem };
+    }
+
+    pub fn deinit(l: *Lead, gpa: std.mem.Allocator) void {
+        l.slots.deinit(gpa);
+        gpa.free(l.mem);
+    }
+
+    pub fn cache(l: *Lead, id: u32) ?*E.Cache {
+        return if (l.slots.find(id)) |x| &x.cache else null;
+    }
+
+    /// One round over `entries` (ascending ids): every stage applies them, then one forward of `ids` and `keep`.
+    pub fn step(l: *Lead, entries: []const Entry, ids: []const u32, keep: E.Keep) !void {
+        var streams: [E.max_segs]plan.Stream = undefined;
+        if (entries.len > streams.len) return error.TooManyStreams;
+        for (entries, streams[0..entries.len]) |en, *st| st.* = l.expect(en);
+        const mode: plan.Mode = if (keep == .last) .prompt else .window;
+        const half = l.mem.len / 2;
+        const p = try plan.encodeMode(mode, streams[0..entries.len], &.{}, l.limits, l.mem[0..half]);
+        defer l.round += 1;
+        var r: Round = .{};
+        for (streams[0..entries.len]) |st| {
+            const x = try l.slots.apply(st) orelse continue; // rank 0 checks what it planned, as every stage does
+            r.parts[r.n] = .{ .cache = &x.cache, .rows = st.rows };
+            r.n += 1;
+        }
+        const ring = l.ring orelse {
+            if (r.n > 0) try l.e.forward(r.parts[0..r.n], ids, keep);
+            return;
+        };
+        try ring.send(.plan, l.round, 0, &.{p});
+        if (r.n == 0) return echoed(ring, l.round, p, l.mem[half..]);
+        var link: Link = .{ .ring = ring, .round = l.round, .bound = boundary(l.e), .plan = try plan.decode(p, l.limits), .sent = p, .echo = l.mem[half..] };
+        l.e.hop = link.hop();
+        defer l.e.hop = null;
+        try l.e.forward(r.parts[0..r.n], ids, keep);
+    }
+
+    /// The plan's entry for `en`: where the stream stands once its previous window keeps `accepted` rows.
+    fn expect(l: *Lead, en: Entry) plan.Stream {
+        var st: plan.Stream = .{ .id = en.id, .op = en.op, .first = 0, .rows = en.rows, .accepted = en.accepted };
+        const x = l.slots.find(en.id) orelse return st;
+        st.generation = x.generation;
+        if (en.op != .@"continue") return st;
+        st.first = @intCast(x.cache.len);
+        if (x.cache.last) |w| {
+            st.first = @intCast(w.start + en.accepted);
+            st.generation +%= 1;
+        }
+        return st;
+    }
+
+    /// A keep-alive: a plan with no streams, which every stage forwards and none acts on.
+    pub fn beat(l: *Lead) !void {
+        try l.step(&.{}, &.{}, .logits);
+    }
+};
+
+/// A follower's rounds until rank 0 stops the ring: each plan passed on, its streams applied, its rows run.
+pub fn serve(gpa: std.mem.Allocator, e: *E.Engine, ring: *stage.Ring, parallel: u32, context: u32, first_round: u64) !void {
+    var slots = try Slots.init(gpa, e, parallel, context);
+    defer slots.deinit(gpa);
+    const limits = plan.Limits.of(shapes(parallel, context));
+    const buf = try gpa.alloc(u8, plan.planSize(parallel, 0));
+    defer gpa.free(buf);
+    const zeros = try gpa.alloc(u32, E.max_rows); // later stages read no token ids
+    defer gpa.free(zeros);
+    @memset(zeros, 0);
+    var round = first_round;
+    while (true) : (round += 1) {
+        const h = ring.recv(.plan, round, buf) catch |err| return if (err == error.Stopped) {} else err;
+        const bytes = buf[0..@intCast(h.bytes)];
+        const p = try plan.decode(bytes, limits);
+        try ring.send(.plan, round, 0, &.{bytes}); // on round, so the last stage echoes it to rank 0
+        if (p.cuts != 0) return error.CutsNotInV1;
+        var r: Round = .{};
+        for (0..p.streams) |i| {
+            const st = p.stream(i);
+            const x = try slots.apply(st) orelse continue;
+            r.parts[r.n] = .{ .cache = &x.cache, .rows = st.rows };
+            r.n += 1;
+        }
+        if (r.n == 0) continue;
+        if (p.mode == .prompt and r.n != 1) return error.PromptParts;
+        var link: Link = .{ .ring = ring, .round = round, .bound = boundary(e), .plan = p };
+        e.hop = link.hop();
+        defer e.hop = null;
+        try e.forward(r.parts[0..r.n], zeros[0..p.rows], if (p.mode == .prompt) .last else .logits);
+    }
+}
+
+/// This stage's boot description.
+pub fn describe(e: *const E.Engine, rank: u32, world: u32, range: boot.Range, parallel: u32, context: u32) !boot.Description {
+    return .{
+        .rank = rank,
+        .world = world,
+        .version = try boot.Text.init(version),
+        .backend = .cuda,
+        .kv_dtype = if (e.kv8) .f8e4m3 else .bf16,
+        .range = range,
+        .boundary = boundary(e),
+        .shapes = shapes(parallel, context),
+        .taps = .{},
+        .kv_rows = context,
+        .states = parallel,
+    };
+}
+
+/// Boot on the ring, rounds 1 and 2: descriptions, then ready reports, each stage's appended; rank 0 checks all.
+pub fn bootRing(gpa: std.mem.Allocator, ring: *stage.Ring, own: *const boot.Description, model: *const boot.Model, ranges: []const boot.Range) !u64 {
+    try boot.checkLayout(model.layers, ranges);
+    const world = ring.world;
+    const most = world * (boot.describe_fixed + 4 * boot.most_taps + boot.most_version);
+    const buf = try gpa.alloc(u8, most);
+    defer gpa.free(buf);
+    var mine: [boot.describe_fixed + 4 * boot.most_taps + boot.most_version]u8 = undefined;
+    const d = try boot.encodeDescription(own, &mine);
+    var ready: [boot.ready_size]u8 = undefined;
+    const rd = try boot.encodeReady(.{ .rank = own.rank, .loaded = true, .warm = true }, &ready);
+    for ([_][]const u8{ d, rd }, 1..) |msg, round| {
+        if (ring.rank == 0) {
+            try ring.send(.plan, round, 0, &.{msg});
+            const h = try ring.recv(.plan, round, buf);
+            try checkAll(buf[0..@intCast(h.bytes)], world, model, ranges, round == 1);
+        } else {
+            const h = try ring.recv(.plan, round, buf);
+            try ring.send(.plan, round, 0, &.{ buf[0..@intCast(h.bytes)], msg });
+        }
+    }
+    return 3;
+}
+
+/// Every stage's description (or ready report) in rank order, each against rank 0's plan for it.
+fn checkAll(all: []const u8, world: u32, model: *const boot.Model, ranges: []const boot.Range, describes: bool) !void {
+    var at: usize = 0;
+    for (0..world) |rank| {
+        if (describes) {
+            if (all.len < at + boot.describe_fixed) return error.BadLength;
+            const n = boot.describe_fixed + 4 * @as(usize, all[at + 20]) + all[at + 21];
+            if (all.len < at + n) return error.BadLength;
+            const desc = try boot.decodeDescription(all[at..][0..n]);
+            try boot.check(.{ .model = model, .rank = @intCast(rank), .world = world, .backend = .cuda, .range = ranges[rank] }, &desc);
+            at += n;
+        } else {
+            if (all.len < at + boot.ready_size) return error.BadLength;
+            const r = try boot.decodeReady(all[at..][0..boot.ready_size]);
+            if (r.rank != rank or !r.ready()) return error.NotReady;
+            at += boot.ready_size;
+        }
+    }
+    if (at != all.len) return error.BadLength;
+}
diff --git a/zig/tests/qwen35_cuda_main.zig b/zig/tests/qwen35_cuda_main.zig
index f5e1b30e3f..6ec0a4cf29 100644
--- a/zig/tests/qwen35_cuda_main.zig
+++ b/zig/tests/qwen35_cuda_main.zig
@@ -20,6 +20,10 @@ const usage =
     \\       tf-qwen35-cuda <model dir> rows <ids file> <out> [piece] | rows-last <ids file> | rows-cuts <ids file>
     \\rows: final-norm rows as a readout keeps them (raw bf16); rows-last checks them against .last; rows-cuts reports.
     \\       tf-qwen35-cuda <model dir> readout <ids file> [streams]: a test readout's rows alone, together, after a cancel
+    \\       tf-qwen35-cuda <model dir> stage <out> [fp8]: a scenario's logits, one engine or a ring of stages (RFC #2005)
+    \\stage: without TF_STAGE_RANK one engine; with TF_STAGE_RANK, _WORLD, _PEERS, _SECRET_FILE and TF_STAGE_LAYOUT
+    \\(first:end a rank, e.g. 0:32,32:64) each process loads its layers only and rank 0 writes <out>. The scenario: a
+    \\two-chunk prompt, greedy steps, a second stream beside it, drafted windows kept short or rolled back.
     \\       tf-qwen35-cuda <model dir> kl <ids file> <out.json> <skip> <rows> [control|int8]
     \\kl: bf16 caches against fp8 ones, teacher-forced: the first <skip> ids as prompt chunks, then <rows> ids in
     \\windows, each row's KL(bf16 || fp8) of the next-token distributions, top-1 agreement and both losses on the next
@@ -51,6 +55,14 @@ pub fn main(init: std.process.Init) !u8 {
         if (code == 2) std.debug.print("{s}", .{usage});
         return code;
     }
+    if (args.len >= 4 and std.mem.eql(u8, args[2], "stage")) {
+        const fp8 = args.len == 5 and std.mem.eql(u8, args[4], "fp8");
+        if (args.len > 5 or (args.len == 5 and !fp8)) {
+            std.debug.print("{s}", .{usage});
+            return 2;
+        }
+        return @import("qwen35_stage.zig").run(init, a, args[1], args[3], fp8);
+    }
     if (args.len >= 3 and std.mem.eql(u8, args[2], "kl")) {
         if (args.len != 7 and args.len != 8) {
             std.debug.print("{s}", .{usage});
diff --git a/zig/tests/qwen35_stage.zig b/zig/tests/qwen35_stage.zig
new file mode 100644
index 0000000000..f44824f582
--- /dev/null
+++ b/zig/tests/qwen35_stage.zig
@@ -0,0 +1,157 @@
+//! `tf-qwen35-cuda <model> stage <out> [fp8]`: one scenario's logits from one engine, or rank 0 of a stage ring.
+const std = @import("std");
+const cuda = @import("cuda");
+const stage = @import("stage");
+const q = @import("qwen35");
+const E = q.engine;
+const S = q.stage;
+
+const parallel = 2;
+const context = 8192;
+
+pub fn run(init: std.process.Init, a: std.mem.Allocator, dir: []const u8, out_path: []const u8, kv8: bool) !u8 {
+    const config = try std.Io.Dir.cwd().readFileAlloc(init.io, try std.fs.path.join(a, &.{ dir, "config.json" }), a, .limited(1 << 22));
+    var driver = try cuda.Driver.open();
+    defer driver.close();
+    var ctx = try cuda.Context.init(&driver, 0);
+    defer ctx.deinit();
+    if (std.c.getenv("TF_STAGE_RANK") == null) {
+        const e = try E.Engine.init(init.gpa, init.io, &ctx, dir, kv8, null);
+        defer e.deinit();
+        var lead = try S.Lead.init(init.gpa, e, null, parallel, context, 1);
+        defer lead.deinit(init.gpa);
+        return finish(init, a, e, &lead, out_path, "whole");
+    }
+    var set: stage.config.Settings = undefined;
+    try stage.config.parse(stage.config.Env{}, &set);
+    var ranges: [stage.config.most_ranks]stage.boot.Range = undefined;
+    const layout = try parseLayout(std.mem.span(std.c.getenv("TF_STAGE_LAYOUT") orelse return error.MissingStageLayout), &ranges);
+    if (layout.len != set.world) return error.BadStageLayout;
+    const mine = layout[set.rank];
+    const e = try E.Engine.initStage(init.gpa, init.io, &ctx, dir, kv8, null, .{ .first = mine.first, .end = mine.end, .heads = set.rank == 0 });
+    defer e.deinit();
+    const loaded = cuda.usage(false);
+    std.debug.print("stage {d}/{d} layers [{d},{d}): {d:.2} GiB device, {d:.2} GiB pinned after load\n", .{ set.rank, set.world, mine.first, mine.end, gib(loaded.device), gib(loaded.host) });
+    var ring = try stage.Ring.open(set.options(S.fingerprint(config, layout)));
+    defer ring.close();
+    defer ring.stop();
+    const own = try S.describe(e, set.rank, set.world, mine, parallel, context);
+    const model: stage.boot.Model = .{ .version = S.version, .kv_dtype = own.kv_dtype, .layers = @intCast(e.g.model().layers), .boundary = own.boundary, .shapes = own.shapes, .taps = &.{} };
+    const next = try S.bootRing(init.gpa, &ring, &own, &model, layout);
+    if (set.rank != 0) {
+        try S.serve(init.gpa, e, &ring, parallel, context, next);
+        report(set.rank, "served");
+        return 0;
+    }
+    var lead = try S.Lead.init(init.gpa, e, &ring, parallel, context, next);
+    defer lead.deinit(init.gpa);
+    return finish(init, a, e, &lead, out_path, "rank 0");
+}
+
+fn gib(n: u64) f64 {
+    return @as(f64, @floatFromInt(n)) / (1 << 30);
+}
+
+fn report(rank: u32, what: []const u8) void {
+    const u = cuda.usage(false);
+    std.debug.print("stage {d} {s}: {d:.2} GiB device now, {d:.2} GiB peak\n", .{ rank, what, gib(u.device), gib(u.peak) });
+}
+
+/// `first:end` ranges, comma-separated, in rank order.
+fn parseLayout(text: []const u8, out: []stage.boot.Range) ![]const stage.boot.Range {
+    var n: usize = 0;
+    var it = std.mem.splitScalar(u8, text, ',');
+    while (it.next()) |r| {
+        if (n == out.len) return error.BadStageLayout;
+        const colon = std.mem.indexOfScalar(u8, r, ':') orelse return error.BadStageLayout;
+        out[n] = .{ .first = try std.fmt.parseInt(u32, r[0..colon], 10), .end = try std.fmt.parseInt(u32, r[colon + 1 ..], 10) };
+        n += 1;
+    }
+    return out[0..n];
+}
+
+fn finish(init: std.process.Init, a: std.mem.Allocator, e: *E.Engine, lead: *S.Lead, out_path: []const u8, role: []const u8) !u8 {
+    var out: std.ArrayList(u8) = .empty;
+    const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds();
+    const toks = try scenario(a, e, lead, &out);
+    const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0)) / 1e6;
+    std.debug.print("stage {s}: scenario {d:.0} ms, {d} logits rows, tokens", .{ role, ms, out.items.len / (e.g.vocab * 2) });
+    for (toks) |t| std.debug.print(" {d}", .{t});
+    std.debug.print("\n", .{});
+    report(0, role);
+    try std.Io.Dir.cwd().writeFile(init.io, .{ .sub_path = out_path, .data = out.items });
+    return 0;
+}
+
+/// The rounds both runs take; every kept logits row appended to `out`; the greedy tokens of stream 1 returned.
+fn scenario(a: std.mem.Allocator, e: *E.Engine, lead: *S.Lead, out: *std.ArrayList(u8)) ![]const u32 {
+    var prng = std.Random.DefaultPrng.init(7);
+    const rnd = prng.random();
+    const prompt = try a.alloc(u32, 4096 + 904);
+    for (prompt) |*t| t.* = rnd.intRangeLessThan(u32, 1000, 60000);
+    const other = try a.alloc(u32, 300);
+    for (other) |*t| t.* = rnd.intRangeLessThan(u32, 1000, 60000);
+    var toks: std.ArrayList(u32) = .empty;
+    // stream 1's prompt in two chunks, then stream 2's beside it
+    try lead.step(&.{.{ .id = 1, .op = .open, .rows = 4096 }}, prompt[0..4096], .last);
+    try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(1).?.logits));
+    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 904 }}, prompt[4096..], .last);
+    try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(1).?.logits));
+    var t1 = argmax(lead.cache(1).?.logits);
+    try toks.append(a, t1);
+    var kept1: u32 = 0; // stream 1's previous window: no window after a prompt chunk
+    for (0..8) |_| {
+        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .logits);
+        try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(0)));
+        t1 = argmax(e.rowLogits(0));
+        try toks.append(a, t1);
+        kept1 = 1;
+    }
+    try lead.step(&.{.{ .id = 2, .op = .open, .rows = 300 }}, other, .last);
+    try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(2).?.logits));
+    var t2 = argmax(lead.cache(2).?.logits);
+    var kept2: u32 = 0;
+    for (0..8) |_| {
+        const both = [_]S.Entry{ .{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }, .{ .id = 2, .op = .@"continue", .rows = 1, .accepted = kept2 } };
+        try lead.step(&both, &.{ t1, t2 }, .logits);
+        for (0..2) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r)));
+        t1 = argmax(e.rowLogits(0));
+        t2 = argmax(e.rowLogits(1));
+        try toks.append(a, t1);
+        kept2 = 1;
+    }
+    try lead.beat();
+    // drafted windows on stream 1: its next token, repeats of it as drafts up to a cut, then junk
+    const cut = [_]usize{ 3, 5, 3, 2 }; // a greedy check keeps a prefix of 1 to 5 rows; window 2 rolls back whole
+    for (0..4) |w| {
+        var win: [5]u32 = @splat(t1);
+        for (win[cut[w]..]) |*d| d.* = rnd.intRangeLessThan(u32, 1000, 60000);
+        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 5, .accepted = kept1 }}, &win, .logits);
+        var take: u32 = 1;
+        while (take < 5 and argmax(e.rowLogits(take - 1)) == win[take]) take += 1;
+        for (0..5) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r)));
+        t1 = argmax(e.rowLogits(take - 1));
+        try toks.append(a, t1);
+        kept1 = if (w == 2) 0 else take; // window 2 rolls back whole: its first row runs again
+        std.debug.print("window {d}: {d} of 5 rows kept\n", .{ w, kept1 });
+        if (w == 2) t1 = win[0];
+    }
+    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .logits);
+    try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(0)));
+    try toks.append(a, argmax(e.rowLogits(0)));
+    try lead.step(&.{ .{ .id = 1, .op = .close }, .{ .id = 2, .op = .close } }, &.{}, .logits);
+    return toks.items;
+}
+
+fn argmax(row: []const u16) u32 {
+    var best: usize = 0;
+    var bv: f32 = -std.math.inf(f32);
+    for (row, 0..) |h, i| {
+        const v: f32 = @bitCast(@as(u32, h) << 16);
+        if (v > bv) {
+            bv = v;
+            best = i;
+        }
+    }
+    return @intCast(best);
+}

From b4b6969ff20f16f119f43f6676c027d7ab5d7172 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Wed, 4 Nov 2026 16:19:52 +0100
Subject: [PATCH 3/4] qwen3_5 cuda: the load-scope check reads initStage, and
 test geometries built field by field are one engine

The engine's checkpoint scope moved into initStage with the stage span. Test geometries that start undefined set
first and whole to 0, or deltas() reads past the layer table.
---
 zig/cuda_checkpoint_cache_host.zig                 | 6 +++---
 zig/src/families/qwen3_5/cuda_disk_test.zig        | 2 ++
 zig/src/families/qwen3_5/cuda_reuse.zig            | 2 ++
 zig/src/families/qwen3_5/cuda_reuse_pages_test.zig | 2 ++
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/zig/cuda_checkpoint_cache_host.zig b/zig/cuda_checkpoint_cache_host.zig
index ae6fa8ff69..816caa056c 100644
--- a/zig/cuda_checkpoint_cache_host.zig
+++ b/zig/cuda_checkpoint_cache_host.zig
@@ -50,17 +50,17 @@ const mm = "families/minimax_m3/";
 const glm = "families/glm/";
 
 test "CUDA load cache: Qwen dense backbone closes after all load reads" {
-    try scope(qwen ++ "cuda_engine.zig", "init", "ck");
+    try scope(qwen ++ "cuda_engine.zig", "initStage", "ck");
     try scope(qwen ++ "cuda_train_load.zig", "weights", "ck");
 }
 test "CUDA load cache: Qwen MoE backbone and its separate heads" {
-    try scope(qwen ++ "cuda_engine.zig", "init", "ck");
+    try scope(qwen ++ "cuda_engine.zig", "initStage", "ck");
     try scope(qwen ++ "cuda_mtp_weights.zig", "load", "ck");
     try scope(qwen ++ "cuda_mtp_bf16.zig", "load", "ck");
     try scope(qwen ++ "cuda_draft_weights.zig", "load", "ck");
 }
 test "CUDA load cache: Qwen EXL3 shares the backbone checkpoint owner" {
-    try scope(qwen ++ "cuda_engine.zig", "init", "ck");
+    try scope(qwen ++ "cuda_engine.zig", "initStage", "ck");
 }
 test "CUDA load cache: Gemma backbone, vision and both draft loaders" {
     try scope(gemma ++ "cuda_weights.zig", "loadConfig", "ck"); // load reads the config, then calls it
diff --git a/zig/src/families/qwen3_5/cuda_disk_test.zig b/zig/src/families/qwen3_5/cuda_disk_test.zig
index 91e5b796d3..9e03c7ab12 100644
--- a/zig/src/families/qwen3_5/cuda_disk_test.zig
+++ b/zig/src/families/qwen3_5/cuda_disk_test.zig
@@ -36,6 +36,8 @@ fn testGeometry() geometry.Geometry {
     g.q_heads = 4;
     g.kv_heads = 2;
     g.layers = 8;
+    g.first = 0; // one engine, not a pipeline stage
+    g.whole = 0;
     @memset(&g.linear, false);
     for (0..g.layers) |i| g.linear[i] = i % 4 != 3;
     return g;
diff --git a/zig/src/families/qwen3_5/cuda_reuse.zig b/zig/src/families/qwen3_5/cuda_reuse.zig
index 2c30be44d7..82f197c8c6 100644
--- a/zig/src/families/qwen3_5/cuda_reuse.zig
+++ b/zig/src/families/qwen3_5/cuda_reuse.zig
@@ -340,6 +340,8 @@ fn testGeometry(layers: usize, attention_every: usize, kv_heads: usize) Geometry
     g.q_heads = 32;
     g.kv_heads = kv_heads;
     g.layers = layers;
+    g.first = 0; // one engine, not a pipeline stage
+    g.whole = 0;
     @memset(&g.linear, false);
     for (0..layers) |i| g.linear[i] = i % attention_every != attention_every - 1;
     return g;
diff --git a/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig b/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig
index 6bc3a14607..5bf6958c95 100644
--- a/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig
+++ b/zig/src/families/qwen3_5/cuda_reuse_pages_test.zig
@@ -86,6 +86,8 @@ test "Qwen paged reuse: a bf16 save and restore copy a page each into the copy's
     e.g.q_heads = 32;
     e.g.kv_heads = 4;
     e.g.layers = 8;
+    e.g.first = 0; // one engine, not a pipeline stage
+    e.g.whole = 0;
     @memset(&e.g.linear, false);
     for (0..8) |i| e.g.linear[i] = i % 4 != 3;
     e.stream = .{ .d = &d, .handle = null };

From da54eef4e9c3baf3e195626104a1fe86a9ba822f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Thu, 5 Nov 2026 08:45:54 +0100
Subject: [PATCH 4/4] qwen3_5 cuda: the stage is a binding of fabric/stage's
 driver, the boot and rounds now shared with other backends

cuda_stage.zig keeps only what is Qwen's: the version, the boundary (h and r, bf16), the shapes, the KV dtype,
a slot's cache, its position, window and prefix-path rollback, and the forward of the stage's layers with the
round's hop, a prompt chunk keeping its last row's logits and a window every row's. The boot, the slots, rank 0's
rounds and the follower's serve loop are the driver's. Lead.step takes the plan's mode instead of the engine's keep.
---
 zig/src/families/qwen3_5/cuda_stage.zig | 341 ++++--------------------
 zig/tests/qwen35_stage.zig              |  16 +-
 2 files changed, 58 insertions(+), 299 deletions(-)

diff --git a/zig/src/families/qwen3_5/cuda_stage.zig b/zig/src/families/qwen3_5/cuda_stage.zig
index 0a32857e36..de241101ee 100644
--- a/zig/src/families/qwen3_5/cuda_stage.zig
+++ b/zig/src/families/qwen3_5/cuda_stage.zig
@@ -1,16 +1,20 @@
-//! A Qwen CUDA engine as one stage of the ring (RFC #2005 v1): its boot, its streams' slots, rank 0's rounds.
+//! A Qwen CUDA engine as one stage of the ring (RFC #2005 v1): its binding of fabric/stage's driver.
 const std = @import("std");
 const stage = @import("stage");
 const E = @import("cuda_engine.zig");
-const sf = @import("cuda_stage_forward.zig");
 
-const plan = stage.plan;
 const boot = stage.boot;
+const plan = stage.plan;
+const driver = stage.driver;
 
-/// Every stage of one placement reports this; rank 0 refuses another.
+/// Every stage of one placement reports this; rank 0 refuses another. This file is the driver's binding `B`.
 pub const version = "tensorfold qwen3_5-cuda stage-v1";
 
-pub const Error = error{ StreamOpen, NoFreeSlot, UnknownStream, BadAccepted, OtherPosition, ResumeNotInV1, CutsNotInV1, OtherEcho, PromptParts };
+pub const Engine = E.Engine;
+pub const Cache = E.Cache;
+pub const Part = E.Part;
+pub const backend: stage.payload.Backend = .cuda;
+pub const most_streams = E.max_segs;
 
 /// The ring's boundary for this engine: h and r as bf16 planes, no drafter taps (v1).
 pub fn boundary(e: *const E.Engine) boot.Boundary {
@@ -22,304 +26,59 @@ pub fn shapes(parallel: u32, context: u32) boot.Shapes {
     return .{ .prompt_chunk = E.max_rows, .max_rows = E.max_rows, .parallel = parallel, .context = context };
 }
 
-/// The checkpoint's and layout's fingerprint: SHA-256 over config.json, the version and every rank's range.
-pub fn fingerprint(config: []const u8, ranges: []const boot.Range) [32]u8 {
-    var h = std.crypto.hash.sha2.Sha256.init(.{});
-    h.update(config);
-    h.update(version);
-    for (ranges) |r| {
-        var b: [8]u8 = undefined;
-        std.mem.writeInt(u32, b[0..4], r.first, .little);
-        std.mem.writeInt(u32, b[4..8], r.end, .little);
-        h.update(&b);
-    }
-    return h.finalResult();
+pub fn kvDtype(e: *const E.Engine) stage.payload.Dtype {
+    return if (e.kv8) .f8e4m3 else .bf16;
 }
 
-/// One stream's state on this stage: its cache (its own layers only) and its state generation.
-pub const Slot = struct { cache: E.Cache, id: u32 = 0, used: bool = false, generation: u32 = 0 };
-
-/// A stage's streams: the same code on rank 0 and on every follower, so each applies a plan's keeps alike.
-pub const Slots = struct {
-    e: *E.Engine,
-    all: []Slot,
-
-    pub fn init(gpa: std.mem.Allocator, e: *E.Engine, n: usize, context: usize) !Slots {
-        const all = try gpa.alloc(Slot, n);
-        var made: usize = 0;
-        errdefer {
-            for (all[0..made]) |*s| s.cache.deinit(gpa);
-            gpa.free(all);
-        }
-        for (all) |*s| {
-            s.* = .{ .cache = try E.Cache.init(gpa, e, context) };
-            made += 1;
-        }
-        return .{ .e = e, .all = all };
-    }
-
-    pub fn deinit(s: *Slots, gpa: std.mem.Allocator) void {
-        for (s.all) |*x| x.cache.deinit(gpa);
-        gpa.free(s.all);
-    }
-
-    pub fn find(s: *Slots, id: u32) ?*Slot {
-        for (s.all) |*x| if (x.used and x.id == id) return x;
-        return null;
-    }
-
-    /// The stream's previous window kept to its first `accepted` rows (a keep advances the generation).
-    fn settle(s: *Slots, x: *Slot, accepted: u32) !void {
-        const last = x.cache.last orelse return if (accepted != 0) error.BadAccepted;
-        if (accepted > last.rows) return error.BadAccepted;
-        var path: [E.most_window]u32 = undefined;
-        for (path[0..accepted], 0..) |*r, i| r.* = @intCast(i);
-        try x.cache.keep(s.e, path[0..accepted]);
-        x.generation +%= 1;
-    }
-
-    /// One stream entry before the round's forward: opened, kept and checked, or closed; its slot when it has rows.
-    pub fn apply(s: *Slots, st: plan.Stream) !?*Slot {
-        switch (st.op) {
-            .open => {
-                if (s.find(st.id) != null) return error.StreamOpen;
-                const x = for (s.all) |*x| {
-                    if (!x.used) break x;
-                } else return error.NoFreeSlot;
-                try x.cache.reset(s.e);
-                x.* = .{ .cache = x.cache, .id = st.id, .used = true };
-                try plan.sameGeneration(x.generation, st);
-                return x;
-            },
-            .@"continue" => {
-                const x = s.find(st.id) orelse return error.UnknownStream;
-                try s.settle(x, st.accepted);
-                try plan.sameGeneration(x.generation, st);
-                if (x.cache.len != st.first) return error.OtherPosition;
-                return if (st.rows > 0) x else null;
-            },
-            .close, .cancel => {
-                const x = s.find(st.id) orelse return error.UnknownStream;
-                try plan.sameGeneration(x.generation, st);
-                x.used = false;
-                return null;
-            },
-            .@"resume" => return error.ResumeNotInV1, // a prompt cache across stages is step 7
-        }
-    }
-};
-
-/// What rank 0's caller asks of a stream in one round.
-pub const Entry = struct { id: u32, op: plan.Op, rows: u32 = 0, accepted: u32 = 0 };
-
-/// A plan's forward parts, in stream order.
-const Round = struct {
-    parts: [E.max_segs]E.Part = undefined,
-    n: usize = 0,
-};
-
-/// The ring's hop for one round: rows to the next rank, rows from the one before (on rank 0, after the plan's echo).
-const Link = struct {
-    ring: *stage.Ring,
-    round: u64,
-    bound: boot.Boundary,
-    plan: plan.Plan,
-    /// Rank 0: the plan it sent, which the last stage's echo must equal.
-    sent: ?[]const u8 = null,
-    echo: []u8 = &.{},
-
-    fn hop(l: *Link) sf.Hop {
-        return .{ .ctx = l, .give = give, .take = take };
-    }
-
-    fn give(ctx: *anyopaque, rows: usize, bytes: []const u8) anyerror!void {
-        const l: *Link = @ptrCast(@alignCast(ctx));
-        try l.ring.send(.rows, l.round, @intCast(rows), &.{bytes});
-    }
-
-    fn take(ctx: *anyopaque, rows: usize, bytes: []u8) anyerror!void {
-        const l: *Link = @ptrCast(@alignCast(ctx));
-        if (l.sent) |p| {
-            try echoed(l.ring, l.round, p, l.echo);
-            l.sent = null;
-        }
-        const h = try l.ring.recv(.rows, l.round, bytes);
-        const r = try plan.checkRows(h, l.plan, l.bound);
-        if (r.rows != rows) return error.OtherRowCount;
-    }
-};
-
-/// The last stage's copy of rank 0's plan, which must be byte for byte the one sent.
-fn echoed(ring: *stage.Ring, round: u64, sent: []const u8, buf: []u8) !void {
-    const h = try ring.recv(.plan, round, buf);
-    if (!std.mem.eql(u8, buf[0..@intCast(h.bytes)], sent)) return error.OtherEcho;
+/// A slot's cache, for this stage's own layers only.
+pub fn open(gpa: std.mem.Allocator, e: *E.Engine, context: u32) !E.Cache {
+    return E.Cache.init(gpa, e, context);
 }
 
-/// Rank 0: the plan of each round to the ring, its own streams applied alike, then the forward with the hop.
-pub const Lead = struct {
-    e: *E.Engine,
-    /// null: one engine, no ring (the whole-model reference runs the same rounds).
-    ring: ?*stage.Ring,
-    slots: Slots,
-    limits: plan.Limits,
-    round: u64,
-    mem: []u8,
-
-    pub fn init(gpa: std.mem.Allocator, e: *E.Engine, ring: ?*stage.Ring, parallel: u32, context: u32, first_round: u64) !Lead {
-        if (parallel > E.max_segs) return error.TooManyStreams;
-        var slots = try Slots.init(gpa, e, parallel, context);
-        errdefer slots.deinit(gpa);
-        const mem = try gpa.alloc(u8, 2 * plan.planSize(parallel, 0));
-        return .{ .e = e, .ring = ring, .slots = slots, .limits = plan.Limits.of(shapes(parallel, context)), .round = first_round, .mem = mem };
-    }
-
-    pub fn deinit(l: *Lead, gpa: std.mem.Allocator) void {
-        l.slots.deinit(gpa);
-        gpa.free(l.mem);
-    }
-
-    pub fn cache(l: *Lead, id: u32) ?*E.Cache {
-        return if (l.slots.find(id)) |x| &x.cache else null;
-    }
+pub fn close(gpa: std.mem.Allocator, c: *E.Cache) void {
+    c.deinit(gpa);
+}
 
-    /// One round over `entries` (ascending ids): every stage applies them, then one forward of `ids` and `keep`.
-    pub fn step(l: *Lead, entries: []const Entry, ids: []const u32, keep: E.Keep) !void {
-        var streams: [E.max_segs]plan.Stream = undefined;
-        if (entries.len > streams.len) return error.TooManyStreams;
-        for (entries, streams[0..entries.len]) |en, *st| st.* = l.expect(en);
-        const mode: plan.Mode = if (keep == .last) .prompt else .window;
-        const half = l.mem.len / 2;
-        const p = try plan.encodeMode(mode, streams[0..entries.len], &.{}, l.limits, l.mem[0..half]);
-        defer l.round += 1;
-        var r: Round = .{};
-        for (streams[0..entries.len]) |st| {
-            const x = try l.slots.apply(st) orelse continue; // rank 0 checks what it planned, as every stage does
-            r.parts[r.n] = .{ .cache = &x.cache, .rows = st.rows };
-            r.n += 1;
-        }
-        const ring = l.ring orelse {
-            if (r.n > 0) try l.e.forward(r.parts[0..r.n], ids, keep);
-            return;
-        };
-        try ring.send(.plan, l.round, 0, &.{p});
-        if (r.n == 0) return echoed(ring, l.round, p, l.mem[half..]);
-        var link: Link = .{ .ring = ring, .round = l.round, .bound = boundary(l.e), .plan = try plan.decode(p, l.limits), .sent = p, .echo = l.mem[half..] };
-        l.e.hop = link.hop();
-        defer l.e.hop = null;
-        try l.e.forward(r.parts[0..r.n], ids, keep);
-    }
+pub fn reset(e: *E.Engine, c: *E.Cache) !void {
+    try c.reset(e);
+}
 
-    /// The plan's entry for `en`: where the stream stands once its previous window keeps `accepted` rows.
-    fn expect(l: *Lead, en: Entry) plan.Stream {
-        var st: plan.Stream = .{ .id = en.id, .op = en.op, .first = 0, .rows = en.rows, .accepted = en.accepted };
-        const x = l.slots.find(en.id) orelse return st;
-        st.generation = x.generation;
-        if (en.op != .@"continue") return st;
-        st.first = @intCast(x.cache.len);
-        if (x.cache.last) |w| {
-            st.first = @intCast(w.start + en.accepted);
-            st.generation +%= 1;
-        }
-        return st;
-    }
+pub fn position(c: *const E.Cache) u32 {
+    return @intCast(c.len);
+}
 
-    /// A keep-alive: a plan with no streams, which every stage forwards and none acts on.
-    pub fn beat(l: *Lead) !void {
-        try l.step(&.{}, &.{}, .logits);
-    }
-};
+pub fn window(c: *const E.Cache) ?driver.Window {
+    const w = c.last orelse return null;
+    return .{ .start = @intCast(w.start), .rows = @intCast(w.rows) };
+}
 
-/// A follower's rounds until rank 0 stops the ring: each plan passed on, its streams applied, its rows run.
-pub fn serve(gpa: std.mem.Allocator, e: *E.Engine, ring: *stage.Ring, parallel: u32, context: u32, first_round: u64) !void {
-    var slots = try Slots.init(gpa, e, parallel, context);
-    defer slots.deinit(gpa);
-    const limits = plan.Limits.of(shapes(parallel, context));
-    const buf = try gpa.alloc(u8, plan.planSize(parallel, 0));
-    defer gpa.free(buf);
-    const zeros = try gpa.alloc(u32, E.max_rows); // later stages read no token ids
-    defer gpa.free(zeros);
-    @memset(zeros, 0);
-    var round = first_round;
-    while (true) : (round += 1) {
-        const h = ring.recv(.plan, round, buf) catch |err| return if (err == error.Stopped) {} else err;
-        const bytes = buf[0..@intCast(h.bytes)];
-        const p = try plan.decode(bytes, limits);
-        try ring.send(.plan, round, 0, &.{bytes}); // on round, so the last stage echoes it to rank 0
-        if (p.cuts != 0) return error.CutsNotInV1;
-        var r: Round = .{};
-        for (0..p.streams) |i| {
-            const st = p.stream(i);
-            const x = try slots.apply(st) orelse continue;
-            r.parts[r.n] = .{ .cache = &x.cache, .rows = st.rows };
-            r.n += 1;
-        }
-        if (r.n == 0) continue;
-        if (p.mode == .prompt and r.n != 1) return error.PromptParts;
-        var link: Link = .{ .ring = ring, .round = round, .bound = boundary(e), .plan = p };
-        e.hop = link.hop();
-        defer e.hop = null;
-        try e.forward(r.parts[0..r.n], zeros[0..p.rows], if (p.mode == .prompt) .last else .logits);
-    }
+/// The last window kept to its first `accepted` rows, a prefix path.
+pub fn rollback(e: *E.Engine, c: *E.Cache, accepted: u32) !void {
+    var path: [E.most_window]u32 = undefined;
+    for (path[0..accepted], 0..) |*r, i| r.* = @intCast(i);
+    try c.keep(e, path[0..accepted]);
 }
 
-/// This stage's boot description.
-pub fn describe(e: *const E.Engine, rank: u32, world: u32, range: boot.Range, parallel: u32, context: u32) !boot.Description {
-    return .{
-        .rank = rank,
-        .world = world,
-        .version = try boot.Text.init(version),
-        .backend = .cuda,
-        .kv_dtype = if (e.kv8) .f8e4m3 else .bf16,
-        .range = range,
-        .boundary = boundary(e),
-        .shapes = shapes(parallel, context),
-        .taps = .{},
-        .kv_rows = context,
-        .states = parallel,
-    };
+/// This stage's layers over the round's rows: a prompt chunk keeps its last row's logits, a window every row's.
+pub fn forward(e: *E.Engine, parts: []const E.Part, ids: []const u32, mode: plan.Mode, hop: ?driver.Hop) !void {
+    e.hop = if (hop) |h| .{ .ctx = h.ctx, .give = h.give, .take = h.take } else null; // the engine builds without stage
+    defer e.hop = null;
+    try e.forward(parts, ids, if (mode == .prompt) .last else .logits);
 }
 
-/// Boot on the ring, rounds 1 and 2: descriptions, then ready reports, each stage's appended; rank 0 checks all.
-pub fn bootRing(gpa: std.mem.Allocator, ring: *stage.Ring, own: *const boot.Description, model: *const boot.Model, ranges: []const boot.Range) !u64 {
-    try boot.checkLayout(model.layers, ranges);
-    const world = ring.world;
-    const most = world * (boot.describe_fixed + 4 * boot.most_taps + boot.most_version);
-    const buf = try gpa.alloc(u8, most);
-    defer gpa.free(buf);
-    var mine: [boot.describe_fixed + 4 * boot.most_taps + boot.most_version]u8 = undefined;
-    const d = try boot.encodeDescription(own, &mine);
-    var ready: [boot.ready_size]u8 = undefined;
-    const rd = try boot.encodeReady(.{ .rank = own.rank, .loaded = true, .warm = true }, &ready);
-    for ([_][]const u8{ d, rd }, 1..) |msg, round| {
-        if (ring.rank == 0) {
-            try ring.send(.plan, round, 0, &.{msg});
-            const h = try ring.recv(.plan, round, buf);
-            try checkAll(buf[0..@intCast(h.bytes)], world, model, ranges, round == 1);
-        } else {
-            const h = try ring.recv(.plan, round, buf);
-            try ring.send(.plan, round, 0, &.{ buf[0..@intCast(h.bytes)], msg });
-        }
-    }
-    return 3;
+pub const Driver = driver.Driver(@This());
+pub const Entry = driver.Entry;
+pub const Slots = Driver.Slots;
+pub const Lead = Driver.Lead;
+pub const serve = Driver.serve;
+pub const describe = Driver.describe;
+
+/// The checkpoint's and layout's fingerprint: SHA-256 over config.json, the version and every rank's range.
+pub fn fingerprint(config: []const u8, ranges: []const boot.Range) [32]u8 {
+    return driver.fingerprint(version, config, ranges);
 }
 
-/// Every stage's description (or ready report) in rank order, each against rank 0's plan for it.
-fn checkAll(all: []const u8, world: u32, model: *const boot.Model, ranges: []const boot.Range, describes: bool) !void {
-    var at: usize = 0;
-    for (0..world) |rank| {
-        if (describes) {
-            if (all.len < at + boot.describe_fixed) return error.BadLength;
-            const n = boot.describe_fixed + 4 * @as(usize, all[at + 20]) + all[at + 21];
-            if (all.len < at + n) return error.BadLength;
-            const desc = try boot.decodeDescription(all[at..][0..n]);
-            try boot.check(.{ .model = model, .rank = @intCast(rank), .world = world, .backend = .cuda, .range = ranges[rank] }, &desc);
-            at += n;
-        } else {
-            if (all.len < at + boot.ready_size) return error.BadLength;
-            const r = try boot.decodeReady(all[at..][0..boot.ready_size]);
-            if (r.rank != rank or !r.ready()) return error.NotReady;
-            at += boot.ready_size;
-        }
-    }
-    if (at != all.len) return error.BadLength;
+/// Boot on the ring, every rank a CUDA stage; the next round returned.
+pub fn bootRing(gpa: std.mem.Allocator, ring: *stage.Ring, own: *const boot.Description, model: *const boot.Model, ranges: []const boot.Range) !u64 {
+    return driver.bootRing(gpa, ring, own, model, ranges, .cuda);
 }
diff --git a/zig/tests/qwen35_stage.zig b/zig/tests/qwen35_stage.zig
index f44824f582..92d559f9b4 100644
--- a/zig/tests/qwen35_stage.zig
+++ b/zig/tests/qwen35_stage.zig
@@ -93,27 +93,27 @@ fn scenario(a: std.mem.Allocator, e: *E.Engine, lead: *S.Lead, out: *std.ArrayLi
     for (other) |*t| t.* = rnd.intRangeLessThan(u32, 1000, 60000);
     var toks: std.ArrayList(u32) = .empty;
     // stream 1's prompt in two chunks, then stream 2's beside it
-    try lead.step(&.{.{ .id = 1, .op = .open, .rows = 4096 }}, prompt[0..4096], .last);
+    try lead.step(&.{.{ .id = 1, .op = .open, .rows = 4096 }}, prompt[0..4096], .prompt);
     try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(1).?.logits));
-    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 904 }}, prompt[4096..], .last);
+    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 904 }}, prompt[4096..], .prompt);
     try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(1).?.logits));
     var t1 = argmax(lead.cache(1).?.logits);
     try toks.append(a, t1);
     var kept1: u32 = 0; // stream 1's previous window: no window after a prompt chunk
     for (0..8) |_| {
-        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .logits);
+        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .window);
         try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(0)));
         t1 = argmax(e.rowLogits(0));
         try toks.append(a, t1);
         kept1 = 1;
     }
-    try lead.step(&.{.{ .id = 2, .op = .open, .rows = 300 }}, other, .last);
+    try lead.step(&.{.{ .id = 2, .op = .open, .rows = 300 }}, other, .prompt);
     try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(2).?.logits));
     var t2 = argmax(lead.cache(2).?.logits);
     var kept2: u32 = 0;
     for (0..8) |_| {
         const both = [_]S.Entry{ .{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }, .{ .id = 2, .op = .@"continue", .rows = 1, .accepted = kept2 } };
-        try lead.step(&both, &.{ t1, t2 }, .logits);
+        try lead.step(&both, &.{ t1, t2 }, .window);
         for (0..2) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r)));
         t1 = argmax(e.rowLogits(0));
         t2 = argmax(e.rowLogits(1));
@@ -126,7 +126,7 @@ fn scenario(a: std.mem.Allocator, e: *E.Engine, lead: *S.Lead, out: *std.ArrayLi
     for (0..4) |w| {
         var win: [5]u32 = @splat(t1);
         for (win[cut[w]..]) |*d| d.* = rnd.intRangeLessThan(u32, 1000, 60000);
-        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 5, .accepted = kept1 }}, &win, .logits);
+        try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 5, .accepted = kept1 }}, &win, .window);
         var take: u32 = 1;
         while (take < 5 and argmax(e.rowLogits(take - 1)) == win[take]) take += 1;
         for (0..5) |r| try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(r)));
@@ -136,10 +136,10 @@ fn scenario(a: std.mem.Allocator, e: *E.Engine, lead: *S.Lead, out: *std.ArrayLi
         std.debug.print("window {d}: {d} of 5 rows kept\n", .{ w, kept1 });
         if (w == 2) t1 = win[0];
     }
-    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .logits);
+    try lead.step(&.{.{ .id = 1, .op = .@"continue", .rows = 1, .accepted = kept1 }}, &.{t1}, .window);
     try out.appendSlice(a, std.mem.sliceAsBytes(e.rowLogits(0)));
     try toks.append(a, argmax(e.rowLogits(0)));
-    try lead.step(&.{ .{ .id = 1, .op = .close }, .{ .id = 2, .op = .close } }, &.{}, .logits);
+    try lead.step(&.{ .{ .id = 1, .op = .close }, .{ .id = 2, .op = .close } }, &.{}, .window);
     return toks.items;
 }