qwen3_5 cuda: the 27B as a pipeline stage, logits byte-equal to the whole model

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

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

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

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

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");
@@ -60,6 +61,7 @@ test {
 
 test {
     _ = @import("cuda_wide.zig");
+    _ = @import("cuda_weights.zig");
 }
 
 pub const train_ops = @import("cuda_train_ops.zig");

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;

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

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.

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

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 {

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;

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

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