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/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/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.zig b/zig/src/families/qwen3_5/cuda.zig index a4353843b7..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"); @@ -60,6 +61,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_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_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_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 }; 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.zig b/zig/src/families/qwen3_5/cuda_stage.zig new file mode 100644 index 0000000000..de241101ee --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_stage.zig @@ -0,0 +1,84 @@ +//! 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 boot = stage.boot; +const plan = stage.plan; +const driver = stage.driver; + +/// 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 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 { + 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 }; +} + +pub fn kvDtype(e: *const E.Engine) stage.payload.Dtype { + return if (e.kv8) .f8e4m3 else .bf16; +} + +/// 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); +} + +pub fn close(gpa: std.mem.Allocator, c: *E.Cache) void { + c.deinit(gpa); +} + +pub fn reset(e: *E.Engine, c: *E.Cache) !void { + try c.reset(e); +} + +pub fn position(c: *const E.Cache) u32 { + return @intCast(c.len); +} + +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) }; +} + +/// 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 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); +} + +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); +} + +/// 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/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"); } 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 rows [piece] | rows-last | rows-cuts \\rows: final-norm rows as a readout keeps them (raw bf16); rows-last checks them against .last; rows-cuts reports. \\ tf-qwen35-cuda readout [streams]: a test readout's rows alone, together, after a cancel + \\ tf-qwen35-cuda stage [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 . The scenario: a + \\two-chunk prompt, greedy steps, a second stream beside it, drafted windows kept short or rolled back. \\ tf-qwen35-cuda kl [control|int8] \\kl: bf16 caches against fp8 ones, teacher-forced: the first ids as prompt chunks, then 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..92d559f9b4 --- /dev/null +++ b/zig/tests/qwen35_stage.zig @@ -0,0 +1,157 @@ +//! `tf-qwen35-cuda stage [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], .prompt); + try out.appendSlice(a, std.mem.sliceAsBytes(lead.cache(1).?.logits)); + 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}, .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, .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 }, .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)); + 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, .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))); + 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}, .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 } }, &.{}, .window); + 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); +}