From 58bd8ac44e732e2e7c4329d3dc9f75a86e3038b5 Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:10:48 -0800 Subject: [PATCH 001/186] flashnext cuda: the rank collective's in-place FP32 reduce-scatter, bounded and aborted as its sum sumBlock leaves block `rank` of the two ranks' values summed in place (ncclReduceScatter), with sum's peer check, deadline, asynchronous-error wait and abort. flashnext-rank's eof-scatter and timeout-scatter modes make it the first collective a dead or silent peer fails; a mode it does not know is an error instead of a sum. --- .../families/flashnext/cuda_collective.zig | 11 ++++++++++ zig/tests/cuda/flashnext_rank_tests.zig | 22 +++++++++++++------ zig/tests/cuda/main.zig | 2 +- zig/tests/fixtures/gradient-serving.tsv | 2 +- 4 files changed, 28 insertions(+), 9 deletions(-) diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index c1ef75bf9c2..b24d45d31d9 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -8,6 +8,7 @@ const Advanced = struct { ncclCommInitRankConfig: *const fn (*nccl.Comm, c_int, nccl.UniqueId, c_int, *nccl.Config) callconv(.c) nccl.Result, ncclCommAbort: *const fn (nccl.Comm) callconv(.c) nccl.Result, ncclCommGetAsyncError: *const fn (nccl.Comm, *nccl.Result) callconv(.c) nccl.Result, + ncclReduceScatter: @FieldType(nccl.Api, "ncclAllReduce"), }; pub const Collective = struct { @@ -143,4 +144,14 @@ pub const Collective = struct { try c.submitted(c.lib.api.ncclAllReduce(values, values, count, .f32, .sum, c.comm, stream.handle)); try c.wait(null, wire.deadline(c.timeout_ms)); } + + /// sum's FP32 sum of block `rank` of `values` (blocks of `count`) alone, left in place there: a reduce-scatter. + pub fn sumBlock(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.api.ncclReduceScatter(values, values + c.rank * count * 4, count, .f32, .sum, c.comm, stream.handle)); + try c.wait(null, wire.deadline(c.timeout_ms)); + } }; diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig index d0e4d36a23b..cea115b95d2 100644 --- a/zig/tests/cuda/flashnext_rank_tests.zig +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -6,7 +6,14 @@ const flash = @import("flashnext"); const check = @import("check.zig"); const wire = flash.rank_wire; -pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { +pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode_arg: []const u8) !void { + const modes = [_][]const u8{ "sum", "eof", "timeout", "missing", "eof-scatter", "timeout-scatter" }; + for (modes) |m| { + if (std.mem.eql(u8, m, mode_arg)) break; + } else return error.UnknownRankMode; + // eof-scatter and timeout-scatter: the peer fails during a reduce-scatter, the first collective + const scatter = std.mem.endsWith(u8, mode_arg, "-scatter"); + const mode = if (scatter) mode_arg[0 .. mode_arg.len - "-scatter".len] else mode_arg; const missing = std.mem.eql(u8, mode, "missing"); const started = wire.now(); if (missing) { @@ -52,18 +59,19 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { } else { sleepSixTenths(); } - check.pass("rank 1: injected {s}", .{mode}); + check.pass("rank 1: injected {s}", .{mode_arg}); return; } - var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&@as(f32, 1))); + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&[2]f32{ 1, 1 })); defer values.free(); const start = wire.now(); - completeSum(&collective, stream, values.ptr, 1) catch |err| { + completeSum(&collective, stream, values.ptr, 1, scatter) catch |err| { try check.expect(err == error.RankDisconnected or err == error.RankTimeout or err == error.NcclFailed, "failure is surfaced ({s})", .{@errorName(err)}); try check.expect(collective.failed and collective.comm == null, "failed communicator is aborted and cannot be reused", .{}); try check.expect(wire.now() - start < 3 * std.time.ns_per_s, "failed collective terminates within bound", .{}); try std.testing.expectError(error.RankDisconnected, collective.sum(stream, values.ptr, 1)); - check.pass("rank 0: {s} surfaced as {s}, aborted, and refused reuse within bound", .{ mode, @errorName(err) }); + try std.testing.expectError(error.RankDisconnected, collective.sumBlock(stream, values.ptr, 1)); + check.pass("rank 0: {s} surfaced as {s}, aborted, and refused reuse within bound", .{ mode_arg, @errorName(err) }); return; }; return error.ExpectedRankFailure; @@ -124,8 +132,8 @@ fn meet(link: wire.Link, rank: usize, sequence: u64) !void { } } -fn completeSum(c: *flash.collective.Collective, stream: cuda.Stream, values: u64, count: usize) !void { - try c.sum(stream, values, count); +fn completeSum(c: *flash.collective.Collective, stream: cuda.Stream, values: u64, count: usize, scatter: bool) !void { + if (scatter) try c.sumBlock(stream, values, count) else try c.sum(stream, values, count); try c.finish(stream); } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 0f4d6c266d6..c525f198598 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -139,7 +139,7 @@ const usage = \\ flashnext-mtp-parity MODEL IDS_FILE several streams' MTP heads in one shared pass against each alone \\ flashnext-score MODEL IDS_FILE decisions' label scores against a plain pass's last row \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) - \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes + \\ flashnext-rank RANK PORT [sum|eof|timeout|missing|eof-scatter|timeout-scatter] two-process collectives and failure probes \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) diff --git a/zig/tests/fixtures/gradient-serving.tsv b/zig/tests/fixtures/gradient-serving.tsv index f88f1da6237..284c03a8cf2 100644 --- a/zig/tests/fixtures/gradient-serving.tsv +++ b/zig/tests/fixtures/gradient-serving.tsv @@ -117,7 +117,7 @@ zig/src/cuda/vit.zig 72c4301251f50d0f25907f0d8ad03a943281b234c506a455e413627a3ff zig/src/families/flashnext/cuda_attn.zig 99dfe79f919a8599490f1cdba0bec20d6e73ccbeafe26d5848419a465fdfa0d0 raw zig/src/families/flashnext/cuda_batch.zig 7fa858c564356e9cc80ee50a86602a5d1c1c3d1ef9ef6d0d69b2ac7a8a4b4e5e raw main 42ff6c510f (disk states keyed by the resolved prompt passes; serving unchanged) zig/src/families/flashnext/cuda_batch_state.zig 5da910105c3b0fa3a776c582198677bd54d977c821f7314a6494d7e9f7a1cbea raw main 4f6017ced (prompt cache marks kept inside the pass), dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) -zig/src/families/flashnext/cuda_collective.zig 898e15d936062b5a26d0aacc0df8b40b5ae4983111d6aec35db1128aaf5f3f76 raw +zig/src/families/flashnext/cuda_collective.zig bf475ffb46b3440fcad280f46e34de25cd22b79ed96a0b587d0869e7451d3519 raw adds sumBlock, the in-place FP32 reduce-scatter with sum's peer check, deadline and abort; no served change on its own zig/src/families/flashnext/cuda_compare_main.zig ade014dfec88f0b386076cdc093ea20710d7bf36c33ff566ca2caf472064695b raw main dfd9ac5ef4 (#2585: the TF_FLASHNEXT_* startup check), c274b45cc1 (the check through the flashnext module) zig/src/families/flashnext/cuda_config.zig 19366ee591270c719fa6f710af0efd4d5a55cfc3e30c923f5ed08771b5c14001 raw zig/src/families/flashnext/cuda_draft_affine.zig 2ec22bec1a38a6d2cc7e16a65d41deeb452d1a679661351a473652b5f39f2da8 raw From 77b02600cd73f39c3f47098b71c441ca57aeb84c Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:10:48 -0800 Subject: [PATCH 002/186] flashnext cuda: a prompt pass at --tp 2 reduce-scatters its experts' slots by rows and gathers their sums A pass wider than any decode round or verify window (more than 64 rows) reduce-scatters y [n, 11, 2560] by row halves instead of all-reducing it, sums each rank's rows' slots with fn_moe_sum into moe_f32, then all-gathers the FP32 sums; the next writeback takes the pending_f32 form (mode 4, one slice). Every summed element is one IEEE add of the two ranks' values, as before, and fn_moe_sum is the writeback's slot chain, so the written rows keep their bytes while the exchange moves 61,440 B a row each way instead of 112,640 B. Rounds and windows keep the all-reduce. y and moe_f32 take one spare row for an odd pass's padded block. flashnext-rank's exchange mode compares the written rows, square sums and normalized rows with the all-reduce's at 65 to 8,192 rows, odd and even, with a changed-sum control, and checks that the exchange writes no row of y or moe_f32 past the spare one. flashnext-model-rank takes TF_FLASHNEXT_CHUNK's rows a pass, as its engine does, where it fixed 32 and so stopped with RankIdentity unless the variable was 32. --- zig/src/families/flashnext/cuda_batch_moe.zig | 21 +- .../families/flashnext/cuda_batch_state.zig | 10 +- zig/src/families/flashnext/cuda_launch.zig | 3 + .../families/flashnext/cuda_moe_exchange.zig | 39 +++ .../flashnext/cuda_moe_exchange_test.zig | 193 ++++++++++++++ .../families/flashnext/cuda_moe_native.zig | 5 +- zig/tests/cuda/flashnext_exchange_tests.zig | 239 ++++++++++++++++++ zig/tests/cuda/flashnext_model_rank_tests.zig | 4 +- zig/tests/cuda/flashnext_rank_tests.zig | 6 +- zig/tests/cuda/main.zig | 2 +- zig/tests/fixtures/gradient-serving.tsv | 4 +- 11 files changed, 508 insertions(+), 18 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_moe_exchange.zig create mode 100644 zig/src/families/flashnext/cuda_moe_exchange_test.zig create mode 100644 zig/tests/cuda/flashnext_exchange_tests.zig diff --git a/zig/src/families/flashnext/cuda_batch_moe.zig b/zig/src/families/flashnext/cuda_batch_moe.zig index a595380b0ad..992607e5f87 100644 --- a/zig/src/families/flashnext/cuda_batch_moe.zig +++ b/zig/src/families/flashnext/cuda_batch_moe.zig @@ -12,14 +12,21 @@ pub fn moeAdds(self: anytype, layer: usize, keep: LayerKeep) !void { if (add == null and keep.moe_branch == 0) return; const s = self.s; const n = self.rows; - var a: cuda.Args = .{}; - a.add(s.y); - a.add(s.wts); - a.add(s.moe_f32); - a.add(batch.cint(dims)); - a.add(batch.cint(slots)); - try cuda.launch.launch(self.k.moe_sum, .{ .grid = .{ .x = @intCast(n), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); + // a prompt pass's rank exchange (cuda_moe_exchange.zig) already left the sum in moe_f32 + if (!self.pending_f32) try sumSlots(self, 0, n); if (add) |m| try m.add(m.ctx, layer, s.mixed, s.picks, s.wts, s.moe_f32, n); if (keep.moe_branch != 0) try self.sumSlices(s.moe_f32, keep.moe_branch, true, n * dims, 1); self.pending_f32 = true; } + +/// moe_f32's rows [r0, r0 + rows): each row's slots of y times its weights, summed in slot order (fn_moe_sum). +pub fn sumSlots(self: anytype, r0: usize, rows: usize) !void { + const s = self.s; + var a: cuda.Args = .{}; + a.add(s.y + r0 * slots * dims * 4); + a.add(s.wts + r0 * slots * 4); + a.add(s.moe_f32 + r0 * dims * 4); + a.add(batch.cint(dims)); + a.add(batch.cint(slots)); + try cuda.launch.launch(self.k.moe_sum, .{ .grid = .{ .x = @intCast(rows), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index beaf2b8f7b3..959ced3b65b 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -302,7 +302,7 @@ pub const Scratch = struct { shared_out: u64 = 0, y: u64 = 0, branch: u64 = 0, - /// A layer's MoE output in fp32 [cap, dims] when something adds to it, and the pocket's activations [cap, 128]. + /// A layer's MoE output in fp32 [cap + 1, dims] when something adds to it, and the pocket's activations [cap, 128]. moe_f32: u64 = 0, pocket_act: u64 = 0, ple_emb: u64 = 0, @@ -393,10 +393,10 @@ pub const Scratch = struct { n += fused_rows * (idw + 2) * 4 + lanes_bytes; n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4 + pairs * 4; n += cap * moe_w * 6 + cap * dims * 4; - n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; + n += pairs * moe_w * 2 + (pairs + slots) * dims * 4 + cap * dims * 2; n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + head_rows * vocab + half * 4 + 4096; - n += cap * dims * 4 + cap * pocket_stride * 4; + n += (cap + 1) * dims * 4 + cap * pocket_stride * 4; n += cap * out_k * 4 + cap * dims * 6; return n; } @@ -469,9 +469,9 @@ pub const Scratch = struct { s.shared_act = try s.add(gpa, d, cap * moe_w * 2); s.shared_out = try s.add(gpa, d, cap * dims * 4); s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); - s.y = try s.add(gpa, d, pairs * dims * 4); + s.y = try s.add(gpa, d, (pairs + slots) * dims * 4); // a spare row: the exchange's odd-row pad s.branch = try s.add(gpa, d, cap * dims * 2); - s.moe_f32 = try s.add(gpa, d, cap * dims * 4); + s.moe_f32 = try s.add(gpa, d, (cap + 1) * dims * 4); s.pocket_act = try s.add(gpa, d, cap * pocket_stride * 4); s.ple_emb = try s.add(gpa, d, cap * dims * 2); s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 083d1c112d1..061d30ddfe9 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -15,6 +15,8 @@ pub const rank_wire = @import("cuda_rank_wire.zig"); pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); /// The MoE pass's native routed experts (NVFP4, block FP8), whose FP8 call a device test takes. pub const moe_native = @import("cuda_moe_native.zig"); +/// A prompt pass's expert exchange at --tp 2, which the two-rank device test runs against the all-reduce. +pub const moe_exchange = @import("cuda_moe_exchange.zig"); pub const hc = @import("cuda_hc.zig"); pub const draft_affine = @import("cuda_draft_affine.zig"); pub const qmm = @import("cuda_qmm.zig"); @@ -471,6 +473,7 @@ test { _ = @import("cuda_round_graph.zig"); _ = @import("cuda_rank_memory.zig"); _ = @import("cuda_rank_wire.zig"); + _ = @import("cuda_moe_exchange.zig"); _ = @import("cuda_modelopt_experts.zig"); _ = @import("cuda_modelopt_threads.zig"); _ = row_path; diff --git a/zig/src/families/flashnext/cuda_moe_exchange.zig b/zig/src/families/flashnext/cuda_moe_exchange.zig new file mode 100644 index 00000000000..f1ca6c5e756 --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe_exchange.zig @@ -0,0 +1,39 @@ +//! A prompt pass's routed-expert exchange at --tp 2: each rank sums its share of rows' slots, then gathers the sums. +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const moe = @import("cuda_batch_moe.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const share = cuda.rank_scatter.share; + +/// Rows past any decode round or verify window: `world` and `n` are what both ranks' passes agree on. +pub fn rowSplit(world: usize, n: usize) bool { + return world == 2 and n > mtp.round_rows; +} + +/// Rows of each rank's reduce-scatter block: rank 0's share, which rank 1's pads to (row n when n is odd). +pub fn block(n: usize) usize { + return share(n, 2, 0)[1]; +} + +/// Rows of y the clear spans: a split pass's two blocks, else the pass's own. +pub fn clearedRows(world: usize, n: usize) usize { + return if (rowSplit(world, n)) 2 * block(n) else n; +} + +/// y's slots, a rank's values and the other's +0, to their fp32 sum in moe_f32 on both ranks: the pending_f32 form. +pub fn run(pass: anytype, c: anytype, n: usize) !void { + const s = pass.s; + const h = block(n); + const own = share(n, 2, c.rank); + try c.sumBlock(pass.stream, s.y, h * slots * dims); + try moe.sumSlots(pass, own[0], own[1] - own[0]); + try c.gatherF32(pass.stream, s.moe_f32 + own[0] * dims * 4, s.moe_f32, h * dims); + pass.pending_f32 = true; +} + +test { + if (@import("builtin").os.tag != .windows) _ = @import("cuda_moe_exchange_test.zig"); // over a Unix socket pair +} diff --git a/zig/src/families/flashnext/cuda_moe_exchange_test.zig b/zig/src/families/flashnext/cuda_moe_exchange_test.zig new file mode 100644 index 00000000000..7af7371174d --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe_exchange_test.zig @@ -0,0 +1,193 @@ +//! The exchange's rows and calls over a faked NCCL and driver, and its reduce-scatter's failures against the sum's. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const chunk = @import("cuda_chunk.zig"); +const mtp = @import("cuda_mtp.zig"); +const exchange = @import("cuda_moe_exchange.zig"); +const Collective = @import("cuda_collective.zig").Collective; +const nccl = cuda.nccl; +const abi = cuda.abi; +const expect = std.testing.expect; +const expectEqual = std.testing.expectEqual; + +const dims = wt.dims; +const slots = wt.slots; +const row_y = slots * dims * 4; + +// Each NCCL call and launch in order; an asynchronous state the waits read. +const Fake = struct { + const Call = struct { kind: enum { scatter, gather, reduce, launch }, send: u64 = 0, recv: u64 = 0, count: usize = 0, f32: bool = true, grid: [6]u32 = @splat(0), args: [5]u64 = @splat(0) }; + var calls: [4]Call = undefined; + var len: usize = 0; + var state: nccl.Result = 0; + var aborts: usize = 0; + + fn reset() void { + len = 0; + state = 0; + aborts = 0; + } + fn record(c: Call) void { + calls[len] = c; + len += 1; + } + fn scatter(send: abi.DevicePtr, recv: abi.DevicePtr, count: usize, t: nccl.DataType, op: nccl.RedOp, _: nccl.Comm, _: abi.Stream) callconv(.c) nccl.Result { + record(.{ .kind = .scatter, .send = send, .recv = recv, .count = count, .f32 = t == .f32 and op == .sum }); + return 0; + } + fn reduce(send: abi.DevicePtr, recv: abi.DevicePtr, count: usize, t: nccl.DataType, op: nccl.RedOp, _: nccl.Comm, _: abi.Stream) callconv(.c) nccl.Result { + record(.{ .kind = .reduce, .send = send, .recv = recv, .count = count, .f32 = t == .f32 and op == .sum }); + return 0; + } + fn gather(send: abi.DevicePtr, recv: abi.DevicePtr, count: usize, t: nccl.DataType, _: nccl.Comm, _: abi.Stream) callconv(.c) nccl.Result { + record(.{ .kind = .gather, .send = send, .recv = recv, .count = count, .f32 = t == .f32 }); + return 0; + } + fn read(comptime T: type, args: ?[*]?*anyopaque, i: usize) T { + return @as(*const T, @ptrCast(@alignCast(args.?[i].?))).*; + } + fn launch(_: abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, bx: c_uint, by: c_uint, bz: c_uint, _: c_uint, _: abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result { + var c: Call = .{ .kind = .launch, .grid = .{ gx, gy, gz, bx, by, bz } }; + for (0..3) |i| c.args[i] = read(u64, args, i); + for (3..5) |i| c.args[i] = @intCast(read(c_int, args, i)); + record(c); + return 0; + } + fn asyncError(_: nccl.Comm, out: *nccl.Result) callconv(.c) nccl.Result { + out.* = state; + return 0; + } + fn abort(_: nccl.Comm) callconv(.c) nccl.Result { + aborts += 1; + return 0; + } +}; + +const Bufs = struct { y: u64 = 0x1000_0000, wts: u64 = 0x2000_0000, moe_f32: u64 = 0x3000_0000 }; + +const Pass = struct { + s: *const Bufs, + k: struct { moe_sum: cuda.Function }, + stream: cuda.Stream, + pending_f32: bool = false, +}; + +/// Rank `rank`'s collective over one end of `pair`, its NCCL faked. +fn collective(pair: *[2]c_int, rank: usize, timeout_ms: u64) !Collective { + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, pair) != 0) return error.SocketPair; + var c: Collective = undefined; + c.failed = false; + c.comm = @ptrFromInt(0x1000); + c.link = .{ .fd = pair[0] }; + c.timeout_ms = timeout_ms; + c.rank = rank; + c.lib.api.ncclAllReduce = Fake.reduce; + c.lib.api.ncclAllGather = Fake.gather; + c.api.ncclReduceScatter = Fake.scatter; + c.api.ncclCommGetAsyncError = Fake.asyncError; + c.api.ncclCommAbort = Fake.abort; + return c; +} + +fn close(pair: [2]c_int) void { + for (pair) |fd| if (fd >= 0) { + _ = std.c.close(fd); + }; +} + +test "rounds and windows keep the all-reduce; a wider prompt pass splits, at --tp 2 only" { + for (1..mtp.round_rows + 1) |n| { + try expect(!exchange.rowSplit(2, n)); + try expectEqual(n, exchange.clearedRows(2, n)); + } + for ([_]usize{ mtp.round_rows + 1, batch.default_chunk, chunk.lone_rows, batch.max_chunk }) |n| { + try expect(exchange.rowSplit(2, n)); + try expect(!exchange.rowSplit(1, n)); + try expectEqual(n, exchange.clearedRows(1, n)); + } +} + +test "the ranks' rows cover a split pass once, in blocks that pad at most the scratch's one spare row" { + var n: usize = mtp.round_rows + 1; + while (n <= batch.max_chunk) : (n += if (n < 600) 1 else 997) { + const h = exchange.block(n); + const r0 = cuda.rank_scatter.share(n, 2, 0); + const r1 = cuda.rank_scatter.share(n, 2, 1); + try expect(r0[0] == 0 and r0[1] == h and r1[0] == h and r1[1] == n); + try expect(exchange.clearedRows(2, n) == 2 * h and 2 * h - n <= 1); + } + try expectEqual(@as(usize, 2 * 8192), exchange.clearedRows(2, batch.max_chunk)); +} + +test "a split pass reduce-scatters its rank's block in place, sums that block's rows and gathers the sums in place" { + var d: cuda.Driver = undefined; + d.api.cuLaunchKernel = Fake.launch; + const bufs: Bufs = .{}; + for ([_]usize{ 0, 1 }) |rank| for ([_]usize{ 65, 66, 1537, batch.default_chunk, batch.max_chunk }) |n| { + var pair: [2]c_int = .{ -1, -1 }; + var c = try collective(&pair, rank, 1000); + defer close(pair); + var pass: Pass = .{ .s = &bufs, .k = .{ .moe_sum = .{ .d = &d, .handle = @ptrFromInt(1) } }, .stream = .{ .d = &d, .handle = null } }; + Fake.reset(); + try exchange.run(&pass, &c, n); + const h = (n + 1) / 2; + const lo = rank * h; + const rows = if (rank == 0) h else n - h; + try expectEqual(@as(usize, 3), Fake.len); + const rs = Fake.calls[0]; + try expect(rs.kind == .scatter and rs.f32 and rs.send == bufs.y and rs.recv == bufs.y + lo * row_y and rs.count == h * slots * dims); + const k = Fake.calls[1]; + try expect(k.kind == .launch); + try expectEqual([6]u32{ @intCast(rows), dims / 256, 1, 256, 1, 1 }, k.grid); + try expectEqual([5]u64{ bufs.y + lo * row_y, bufs.wts + lo * slots * 4, bufs.moe_f32 + lo * dims * 4, dims, slots }, k.args); + const ag = Fake.calls[2]; + try expect(ag.kind == .gather and ag.f32 and ag.send == bufs.moe_f32 + rank * h * dims * 4 and ag.recv == bufs.moe_f32 and ag.count == h * dims); + try expect(pass.pending_f32 and !c.failed and Fake.aborts == 0); + }; +} + +const Outcome = struct { err: ?anyerror, calls: usize, aborted: bool, refused: bool }; + +/// sum or sumBlock under one failure: what it returned, what reached NCCL, and whether it left the communicator dead. +fn outcome(comptime block: bool, failure: enum { none, empty, peer, async_error, stuck }) !Outcome { + var pair: [2]c_int = .{ -1, -1 }; + var c = try collective(&pair, 1, 50); + defer close(pair); + var d: cuda.Driver = undefined; + const stream: cuda.Stream = .{ .d = &d, .handle = null }; + Fake.reset(); + switch (failure) { + .peer => { + _ = std.c.close(pair[1]); + pair[1] = -1; + }, + .async_error => Fake.state = 3, + .stuck => Fake.state = nccl.in_progress, + else => {}, + } + const count: usize = if (failure == .empty) 0 else 4099; + const f = if (block) Collective.sumBlock else Collective.sum; + var err: ?anyerror = null; + f(&c, stream, 0x5000_0000, count) catch |e| { + err = e; + }; + const calls = Fake.len; + const refused = if (f(&c, stream, 0x5000_0000, 1)) |_| false else |e| e == error.RankDisconnected; + return .{ .err = err, .calls = calls, .aborted = c.failed and c.comm == null and Fake.aborts == 1, .refused = refused }; +} + +test "the reduce-scatter keeps the sum's contract: a dead peer, an asynchronous error or the deadline aborts it for good" { + const cases = .{ .none, .empty, .peer, .async_error, .stuck }; + inline for (cases) |failure| { + const sum = try outcome(false, failure); + const scatter = try outcome(true, failure); + try expectEqual(sum, scatter); + } + try expectEqual(Outcome{ .err = null, .calls = 1, .aborted = false, .refused = false }, try outcome(true, .none)); + try expectEqual(Outcome{ .err = null, .calls = 0, .aborted = false, .refused = false }, try outcome(true, .empty)); + try expectEqual(Outcome{ .err = error.RankDisconnected, .calls = 0, .aborted = true, .refused = true }, try outcome(true, .peer)); + try expectEqual(Outcome{ .err = error.NcclFailed, .calls = 1, .aborted = true, .refused = true }, try outcome(true, .async_error)); + try expectEqual(Outcome{ .err = error.RankTimeout, .calls = 1, .aborted = true, .refused = true }, try outcome(true, .stuck)); +} diff --git a/zig/src/families/flashnext/cuda_moe_native.zig b/zig/src/families/flashnext/cuda_moe_native.zig index d6b58e3d57c..5498da8e0ce 100644 --- a/zig/src/families/flashnext/cuda_moe_native.zig +++ b/zig/src/families/flashnext/cuda_moe_native.zig @@ -5,6 +5,7 @@ const mixed = @import("cuda_mixed_experts.zig"); const nvfp4_tile = @import("cuda_nvfp4_tile.zig"); const zs = @import("cuda_zero_slice.zig"); const shared = @import("cuda_batch_side.zig"); +const exchange = @import("cuda_moe_exchange.zig"); const Kernels = @import("cuda_batch_kernels.zig").Kernels; const dims = wt.dims; @@ -19,7 +20,8 @@ pub fn run(self: anytype, native: *const mixed.Owned, n: usize, pairs: usize, pl const e = owned.count + 1; const skip: i32 = @intCast(owned.count); const picks = if (native.world == 1) s.picks else try ownPicks(self, s.picks, s.picks_ep, pairs, owned.first, owned.count); - if (zs.clearsY(native.world, native.rank)) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); + const y_bytes = exchange.clearedRows(native.world, n) * slots * dims * 4; + if (zs.clearsY(native.world, native.rank)) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, y_bytes, self.stream.handle), "cuMemsetD8Async"); const side = try shared.fork(self, native, n); switch (owned.format) { .nvfp4 => { @@ -35,6 +37,7 @@ pub fn run(self: anytype, native: *const mixed.Owned, n: usize, pairs: usize, pl if (side != null) try shared.join(self, side, native, n) else try @import("cuda_shared_split.zig").run(self, native, n); if (native.world != 1) { const collective = self.m.collective orelse return error.TensorParallelUnavailable; + if (exchange.rowSplit(native.world, n)) return exchange.run(self, collective, n); try collective.sum(self.stream, s.y, pairs * dims); } } diff --git a/zig/tests/cuda/flashnext_exchange_tests.zig b/zig/tests/cuda/flashnext_exchange_tests.zig new file mode 100644 index 00000000000..79a381a908a --- /dev/null +++ b/zig/tests/cuda/flashnext_exchange_tests.zig @@ -0,0 +1,239 @@ +//! Two ranks' prompt-pass expert exchange against the all-reduce it replaces: the written rows, sums and norms' bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); + +const Collective = flash.collective.Collective; +const dims = 2560; +const slots = 11; +const streams = 4; +const wide = dims * streams; +const pss_n = dims / 256 * streams; +const eps: f32 = 1e-6; +const neg_zero: u32 = 0x8000_0000; + +/// -0, +0, the least subnormals, the largest subnormal and finite values, the infinities, a NaN and two plain values. +const specials = [_]u32{ neg_zero, 0, 1, 0x8000_0001, 0x007f_ffff, 0x7f7f_ffff, 0xff7f_ffff, 0x7f80_0000, 0xff80_0000, 0x7fc0_0000, 0x3f80_0000, 0xbe80_0000 }; + +pub fn run(gpu: check.Gpu, c: *Collective, stream: cuda.Stream, rank: usize) !void { + try blocks(gpu, c, stream, rank); + try planted(gpu, c, stream, rank); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + defer mod.unload(); + const k: Kernels = .{ .moe_sum = try mod.function("fn_moe_sum") }; + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + for ([_]usize{ 65, 66, 67, 255, 1536, 1537, 8192 }) |n| try pass(gpu, c, stream, tri, k, rank, n); + check.pass("rank {d}: the exchange's written rows, square sums and normalized rows equal the all-reduce's for 65 to 8,192 rows, odd and even; a changed sum fails", .{rank}); +} + +/// sumBlock leaves block `rank` of the two ranks' values summed exactly, from one value to a vocabulary half. +fn blocks(gpu: check.Gpu, c: *Collective, stream: cuda.Stream, rank: usize) !void { + for ([_]usize{ 1, 17, 4099, 124_160 }) |count| { + const host = try gpu.gpa.alloc(f32, 2 * count); + defer gpu.gpa.free(host); + for (host, 0..) |*v, i| v.* = @floatFromInt(i % 127 + rank); + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer values.free(); + try c.sumBlock(stream, values.ptr, count); + try c.finish(stream); + try values.download(0, std.mem.sliceAsBytes(host)); + for (host[rank * count ..][0..count], rank * count..) |v, at| try check.expect(v == @as(f32, @floatFromInt(2 * (at % 127) + 1)), "rank {d}: reduce-scatter block at {d} of {d}", .{ rank, at, count }); + } + check.pass("rank {d}: in-place FP32 reduce-scatter exact for a scalar, tail and vocabulary-half blocks", .{rank}); +} + +/// One rank's specials beside the other's +0, then both ranks' values: the block's bytes equal the all-reduce's. +fn planted(gpu: check.Gpu, c: *Collective, stream: cuda.Stream, rank: usize) !void { + const m = specials.len * 3; + var host: [2 * m]u32 = undefined; + for (&host, 0..) |*v, i| { + const j = i % specials.len; + const case = (i / specials.len) % 3; + v.* = if (case == 2) specials[(j + rank * 5) % specials.len] else if (case == rank) specials[j] else 0; + } + var all = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&host)); + defer all.free(); + var mine = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&host)); + defer mine.free(); + try c.sum(stream, all.ptr, 2 * m); + try c.sumBlock(stream, mine.ptr, m); + try c.finish(stream); + var got_all: [2 * m]u32 = undefined; + var got: [2 * m]u32 = undefined; + try all.download(0, std.mem.asBytes(&got_all)); + try mine.download(0, std.mem.asBytes(&got)); + try check.sameBytes("reduce-scatter block against the all-reduce, planted", std.mem.sliceAsBytes(got[rank * m ..][0..m]), std.mem.sliceAsBytes(got_all[rank * m ..][0..m])); + for (rank * m..rank * m + m) |i| { + const j = i % specials.len; + if ((i / specials.len) % 3 == 2) continue; + const v: f32 = @bitCast(specials[j]); + const want: u32 = if (specials[j] == neg_zero) 0 else specials[j]; + try check.expect(if (std.math.isNan(v)) std.math.isNan(@as(f32, @bitCast(got[i]))) else got[i] == want, "rank {d}: x + 0 keeps x (-0 gives +0) at {d}: {x:0>8} from {x:0>8}", .{ rank, i, got[i], specials[j] }); + } + check.pass("rank {d}: reduce-scatter equals the all-reduce on planted -0, subnormals, extremes, infinities, NaN and two-rank values", .{rank}); +} + +const Kernels = struct { moe_sum: cuda.Function }; +const Bufs = struct { y: u64, wts: u64, moe_f32: u64 }; +const Pass = struct { s: *const Bufs, k: Kernels, stream: cuda.Stream, pending_f32: bool = false }; + +/// A writeback's outputs over the pass's rows and the spare one, poisoned so a row no launch writes shows. +const Out = struct { + hout: cuda.DeviceBuffer, + pss: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + + fn init(d: *const cuda.Driver, stream: cuda.Stream, rows: usize) !Out { + var o: Out = .{ .hout = try cuda.DeviceBuffer.alloc(d, rows * wide * 2), .pss = undefined, .norm = undefined }; + errdefer o.hout.free(); + o.pss = try cuda.DeviceBuffer.alloc(d, rows * pss_n * 4); + errdefer o.pss.free(); + o.norm = try cuda.DeviceBuffer.alloc(d, rows * wide * 2); + for ([_]cuda.DeviceBuffer{ o.hout, o.pss, o.norm }) |b| try b.fill8(0x5a, stream.handle); + return o; + } + + fn deinit(o: *Out) void { + o.hout.free(); + o.pss.free(); + o.norm.free(); + } + + fn same(o: Out, gpu: check.Gpu, want: Out, what: []const u8) !void { + inline for (.{ "hout", "pss", "norm" }) |name| { + const got = try check.download(gpu, @field(o, name)); + defer gpu.gpa.free(got); + const ref = try check.download(gpu, @field(want, name)); + defer gpu.gpa.free(ref); + try check.sameBytes(what, got, ref); + } + } +}; + +fn mix(x: u64) u64 { + var z = x +% 0x9e37_79b9_7f4a_7c15; + z = (z ^ (z >> 30)) *% 0xbf58_476d_1ce4_e5b9; + z = (z ^ (z >> 27)) *% 0x94d0_49bb_1331_11eb; + return z ^ (z >> 31); +} + +/// A value in [-1, 1) times 2^-16 to 2^7, or now and then one of the specials. +fn value(h: u64) f32 { + if (h % 4093 < 4) return @bitCast(specials[@intCast((h >> 20) % specials.len)]); + const unit = @as(f32, @floatFromInt(h >> 40)) / (1 << 24); + return (unit * 2 - 1) * std.math.ldexp(@as(f32, 1), @as(i32, @intCast((h >> 8) % 24)) - 16); +} + +fn bf16(v: f32) u16 { + return @truncate(@as(u32, @bitCast(v)) >> 16); +} + +/// Routed pairs owned by a rank drawn per pair, the shared slot's columns in halves; a few elements on both ranks. +fn fill(y: []f32, wts: []f32, n: usize, rank: usize) void { + @memset(y, 0); + @memset(wts, 0); + for (0..n) |r| for (0..slots) |k| { + const pair = r * slots + k; + const w = mix(pair ^ 0x77); + wts[pair] = if (w % 211 == 5) @bitCast(neg_zero) else if (w % 211 == 6) @bitCast(@as(u32, 3)) else @as(f32, @floatFromInt(w >> 40)) / (1 << 26); + if (r == 3 and k == 4) wts[pair] = std.math.nan(f32); + if (r == 5 and k == 2) wts[pair] = std.math.inf(f32); + const owner = mix(pair ^ 0x3a) & 1; + for (0..dims) |d| { + const at = pair * dims + d; + const h = mix(at); + const own = if (k == slots - 1) (d >= dims / 2) == (rank == 1) else owner == rank; + if (own) y[at] = value(h) else if (h % 50_021 == 7) y[at] = value(mix(h)) * 0.75; + } + }; +} + +/// The all-reduce and both mode-2 writebacks against the exchange and both mode-4 writebacks, on the same inputs. +fn pass(gpu: check.Gpu, c: *Collective, stream: cuda.Stream, tri: flash.Tri, k: Kernels, rank: usize, n: usize) !void { + const a = gpu.gpa; + const rows = n + 1; + const y = try a.alloc(f32, rows * slots * dims); + defer a.free(y); + const wts = try a.alloc(f32, rows * slots); + defer a.free(wts); + fill(y, wts, n, rank); + const hs = try a.alloc(u16, n * wide + n * streams); + defer a.free(hs); + for (hs, 0..) |*v, i| v.* = bf16(@as(f32, @floatFromInt(mix(i ^ 0x51) >> 40)) / (1 << 23) - 1); + for (hs[n * wide ..], 0..) |*v, i| v.* = bf16(0.5 + @as(f32, @floatFromInt(mix(i ^ 0x1f) >> 40)) / (1 << 24)); + var scale: [wide]f32 = undefined; + for (&scale, 0..) |*v, i| v.* = 0.5 + @as(f32, @floatFromInt(mix(i ^ 0x2b) >> 40)) / (1 << 24); + var y_a = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(y)); + defer y_a.free(); + // the exchange's y and moe_f32 carry a poisoned guard row past the spare one + const guarded = try a.alloc(f32, (rows + 1) * slots * dims); + defer a.free(guarded); + @memcpy(guarded[0..y.len], y); + @memset(std.mem.sliceAsBytes(guarded[y.len..]), 0x5a); + var y_b = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(guarded)); + defer y_b.free(); + var wd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(wts)); + defer wd.free(); + var hd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(hs)); + defer hd.free(); + var sd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&scale)); + defer sd.free(); + var moe = try cuda.DeviceBuffer.alloc(gpu.d, (rows + 1) * dims * 4); + defer moe.free(); + try moe.fill8(0x5a, stream.handle); + var dummy = try cuda.DeviceBuffer.alloc(gpu.d, 256); + defer dummy.free(); + const h = hd.ptr; + const inj = hd.ptr + n * wide * 2; + const sc = sd.ptr; + var outs: [5]Out = undefined; + var made: usize = 0; + defer for (outs[0..made]) |*o| o.deinit(); + while (made < outs.len) : (made += 1) outs[made] = try Out.init(gpu.d, stream, rows); + const rs = n * dims; + + try c.sum(stream, y_a.ptr, n * slots * dims); + try tri.hcWriteback(h, outs[0].hout.ptr, outs[0].pss.ptr, dummy.ptr, "*bf16", inj, y_a.ptr, "*fp32", wd.ptr, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); + try tri.hcWritebackNorm(h, outs[1].hout.ptr, outs[1].pss.ptr, dummy.ptr, inj, y_a.ptr, wd.ptr, sc, outs[1].norm.ptr, eps, rs, n, dims, streams, 2, 10, 11, 1); + + const bufs: Bufs = .{ .y = y_b.ptr, .wts = wd.ptr, .moe_f32 = moe.ptr }; + var p: Pass = .{ .s = &bufs, .k = k, .stream = stream }; + try flash.moe_exchange.run(&p, c, n); + try check.expect(p.pending_f32, "rank {d}: the exchange leaves its sum pending in fp32", .{rank}); + try tri.hcWriteback(h, outs[2].hout.ptr, outs[2].pss.ptr, moe.ptr, "*fp32", inj, dummy.ptr, "*bf16", dummy.ptr, "*bf16", rs, n, dims, streams, 4, 1, 1, 1); + try tri.hcWritebackNorm(h, outs[3].hout.ptr, outs[3].pss.ptr, moe.ptr, inj, dummy.ptr, dummy.ptr, sc, outs[3].norm.ptr, eps, rs, n, dims, streams, 4, 1, 1, 1); + try c.finish(stream); + try outs[2].same(gpu, outs[0], "exchange then writeback against all-reduce then writeback"); + try outs[3].same(gpu, outs[1], "exchange then writeback-norm against all-reduce then writeback-norm"); + const tail = try check.download(gpu, outs[3].hout); + defer a.free(tail); + for (tail[n * wide * 2 ..]) |b| try check.expect(b == 0x5a, "rank {d}: the spare row past {d} rows is not written", .{ rank, n }); + const y_after = try check.download(gpu, y_b); + defer a.free(y_after); + for (y_after[y.len * 4 ..]) |b| try check.expect(b == 0x5a, "rank {d}: {d} rows: the exchange writes no y row past its spare one", .{ rank, n }); + const moe_after = try check.download(gpu, moe); + defer a.free(moe_after); + for (moe_after[rows * dims * 4 ..]) |b| try check.expect(b == 0x5a, "rank {d}: {d} rows: the exchange writes no moe_f32 row past its spare one", .{ rank, n }); + + // a control on the compared path: one changed sum must change the written rows + const sums = try a.alloc(f32, n * dims); + defer a.free(sums); + try moe.download(0, std.mem.sliceAsBytes(sums)); + var at: usize = 0; + for (sums, 0..) |v, i| { + if (std.math.isFinite(v) and (!std.math.isFinite(sums[at]) or @abs(v) > @abs(sums[at]))) at = i; + } + const flipped: u32 = @as(u32, @bitCast(sums[at])) ^ (1 << 22); + try moe.upload(at * 4, std.mem.asBytes(&flipped)); + try tri.hcWriteback(h, outs[4].hout.ptr, outs[4].pss.ptr, moe.ptr, "*fp32", inj, dummy.ptr, "*bf16", dummy.ptr, "*bf16", rs, n, dims, streams, 4, 1, 1, 1); + try c.finish(stream); + const changed = try check.download(gpu, outs[4].hout); + defer a.free(changed); + const ref = try check.download(gpu, outs[0].hout); + defer a.free(ref); + try check.expect(!std.mem.eql(u8, changed, ref), "rank {d}: a flipped sum bit at {d} changes the written rows", .{ rank, at }); + check.pass("rank {d}: {d} rows: exchange bytes equal the all-reduce's; nothing written past the spare row; the changed-sum control fails as planted", .{ rank, n }); +} diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig index 92418815f67..0a79c81c66f 100644 --- a/zig/tests/cuda/flashnext_model_rank_tests.zig +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -8,7 +8,8 @@ const ranks = flash.rank_engine; pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: []const u8) !void { const context = 256; - const chunk = 32; + // TF_FLASHNEXT_CHUNK's rows a pass: past 64 a pass takes the expert exchange, at most 64 the all-reduce + const chunk = try prompt.chunkRows(); const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); defer gpu.gpa.free(source); var ids: std.ArrayList(u32) = .empty; @@ -48,6 +49,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] defer session.drop(&first); try session.forward(flash.Tri, tri, &first, ids.items, .{}); check.pass("two-rank prompt completed ({d} tokens)", .{ids.items.len}); + std.debug.print("LOGITS {x:0>16}\n", .{std.hash.Wyhash.hash(0, first.logits)}); for (0..prompt.vocabulary) |i| { const value: f32 = @bitCast(@as(u32, std.mem.readInt(u16, first.logits[i * 2 ..][0..2], .little)) << 16); try check.expect(std.math.isFinite(value), "finite initial logits at {d}", .{i}); diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig index cea115b95d2..5a3827c9ad2 100644 --- a/zig/tests/cuda/flashnext_rank_tests.zig +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -7,7 +7,7 @@ const check = @import("check.zig"); const wire = flash.rank_wire; pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode_arg: []const u8) !void { - const modes = [_][]const u8{ "sum", "eof", "timeout", "missing", "eof-scatter", "timeout-scatter" }; + const modes = [_][]const u8{ "sum", "eof", "timeout", "missing", "eof-scatter", "timeout-scatter", "exchange" }; for (modes) |m| { if (std.mem.eql(u8, m, mode_arg)) break; } else return error.UnknownRankMode; @@ -49,6 +49,10 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode_arg: []const u8) !void { if (std.mem.eql(u8, mode, "timeout")) collective.timeout_ms = 200; var stream = try cuda.Stream.init(gpu.d, true); defer stream.deinit(); + if (std.mem.eql(u8, mode, "exchange")) { + try @import("flashnext_exchange_tests.zig").run(gpu, &collective, stream, rank); + return meet(link, rank, 99); + } if (std.mem.eql(u8, mode, "eof") or std.mem.eql(u8, mode, "timeout")) { // both communicators are up before the peer fails, so the failure meets the collective, not the init try meet(link, rank, 98); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index c525f198598..f9f39206584 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -139,7 +139,7 @@ const usage = \\ flashnext-mtp-parity MODEL IDS_FILE several streams' MTP heads in one shared pass against each alone \\ flashnext-score MODEL IDS_FILE decisions' label scores against a plain pass's last row \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) - \\ flashnext-rank RANK PORT [sum|eof|timeout|missing|eof-scatter|timeout-scatter] two-process collectives and failure probes + \\ flashnext-rank RANK PORT [sum|eof|timeout|missing|eof-scatter|timeout-scatter|exchange] two-process collectives, failures, the expert exchange \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) diff --git a/zig/tests/fixtures/gradient-serving.tsv b/zig/tests/fixtures/gradient-serving.tsv index 284c03a8cf2..bb4da22d6c5 100644 --- a/zig/tests/fixtures/gradient-serving.tsv +++ b/zig/tests/fixtures/gradient-serving.tsv @@ -116,7 +116,7 @@ zig/src/cuda/triton.zig 256323865c6a4676b8b75aae9e32d2d9732cf2623efedb9462cf6f95 zig/src/cuda/vit.zig 72c4301251f50d0f25907f0d8ad03a943281b234c506a455e413627a3ffe09c0 raw zig/src/families/flashnext/cuda_attn.zig 99dfe79f919a8599490f1cdba0bec20d6e73ccbeafe26d5848419a465fdfa0d0 raw zig/src/families/flashnext/cuda_batch.zig 7fa858c564356e9cc80ee50a86602a5d1c1c3d1ef9ef6d0d69b2ac7a8a4b4e5e raw main 42ff6c510f (disk states keyed by the resolved prompt passes; serving unchanged) -zig/src/families/flashnext/cuda_batch_state.zig 5da910105c3b0fa3a776c582198677bd54d977c821f7314a6494d7e9f7a1cbea raw main 4f6017ced (prompt cache marks kept inside the pass), dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) +zig/src/families/flashnext/cuda_batch_state.zig ae9d16a76aa56cd3448b3414140cbabd2f918bfb69a094b02be82544f371603c raw y and moe_f32 take one spare row for a --tp 2 prompt pass's padded rank block; replies unchanged zig/src/families/flashnext/cuda_collective.zig bf475ffb46b3440fcad280f46e34de25cd22b79ed96a0b587d0869e7451d3519 raw adds sumBlock, the in-place FP32 reduce-scatter with sum's peer check, deadline and abort; no served change on its own zig/src/families/flashnext/cuda_compare_main.zig ade014dfec88f0b386076cdc093ea20710d7bf36c33ff566ca2caf472064695b raw main dfd9ac5ef4 (#2585: the TF_FLASHNEXT_* startup check), c274b45cc1 (the check through the flashnext module) zig/src/families/flashnext/cuda_config.zig 19366ee591270c719fa6f710af0efd4d5a55cfc3e30c923f5ed08771b5c14001 raw @@ -133,7 +133,7 @@ zig/src/families/flashnext/cuda_lanes.zig df7c17cb079b4a00ff16964c54bd3895bf997a zig/src/families/flashnext/cuda_lanes_admission_test.zig 4b8052db793ba81f0e73e1b9bd77b1780bbc4fd550d90976d017da47e76faaa2 raw zig/src/families/flashnext/cuda_lanes_load.zig 241d091d4b5199325665d25ea5d4def81fd15660351a199773c596ca2025fa18 raw main dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) zig/src/families/flashnext/cuda_lanes_restore_test.zig f9947521a6d9c7d7380ca1b784a93a43b364c6fa5861cad15efd3538d2b8ceff raw -zig/src/families/flashnext/cuda_launch.zig 0fec43d2c2ba7436beb3884287fedaa8099e0dc51ee2464d3722e7d458167b94 raw main ab2d08ac3 (#2802: pins that Flash Next declares no chunk_invariant; no served change), c274b45cc1 (env_validate exported for the tool mains; no served change) +zig/src/families/flashnext/cuda_launch.zig 849584af8bc3cd31f990911d449ec0f828318b05d480fe98fc0c287f5b6fb14d raw exports cuda_moe_exchange.zig for its device test; no served change zig/src/families/flashnext/cuda_media.zig 394601e8a807f5516f501efcaefa6fb3308f392ee322db7296b11c90814f8eea raw zig/src/families/flashnext/cuda_media_test.zig 4292f8f7c8f8bf4efd0dfaaa5a6c63e11568f85534127a0ce9feb0ab0be05120 raw zig/src/families/flashnext/cuda_mix.zig 248a2a86b3f377531f5b713c408ba99f73604510f4afe34ffae52615d5863c2c raw From 119070d6f1bd6d48c29e06577fd26230f0fd8f4c Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 17:09:56 -0800 Subject: [PATCH 003/186] changelog, recipe and ledger: Flash Next's prompt-pass expert exchange at --tp 2 --- CHANGELOG.md | 1 + docs/optimization-ledger.md | 5 +++++ docs/recipes/flash-next-cuda.md | 3 +++ 3 files changed, 9 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c663dc897f..9ddc59ee0c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -59,6 +59,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. +- Flash Next at `--tp 2` on CUDA: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows, then gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every summed value is the same IEEE add and the slots are summed in the writeback's order, so the written rows keep their bytes and replies their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (8,192 tokens: 7.12 to 5.62 s, and 3.61 to 2.70 s); prompts of 128 to 512 tokens are 6.0% to 18.3% sooner on the GB10s (#2977). - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). diff --git a/docs/optimization-ledger.md b/docs/optimization-ledger.md index 2b772874bc9..88e18546774 100644 --- a/docs/optimization-ledger.md +++ b/docs/optimization-ledger.md @@ -447,6 +447,10 @@ below; these are estimates for planning a port, not results. 25. **Tensor parallel ranks.** A model splits over two or more GPUs or hosts, so a model that doesn't fit one box runs, and decode reads each rank's share of the weights. Mixed Metal and CUDA ranks are tracked in [#1236]. + | Family | Gain (measured) | Where | + | --- | --- | --- | + | flashnext (C), TP=2 | prompt passes reduce-scatter the experts' slots by rows: rank 0 sends 5.81 → 3.43 MiB a prompt token, cold prefill +26.5..+26.8 % on two GB10s, +30.9..+33.9 % on two RTX PRO 6000s, same replies | [#2977] | + 26. **Pipeline stages.** Layer ranges on different nodes, one hidden row a token over Ethernet, so a GB10 and a Mac can serve one model ([#2005]). The stage link and its protocol are [#2012] and [#2060]; the 27B as a stage is [#2115], logits byte-equal to the whole model at two and three stages. @@ -615,3 +619,4 @@ below; these are estimates for planning a port, not results. [#2634]: https://github.com/ashhart/TensorFold/pull/2634 [#2784]: https://github.com/ashhart/TensorFold/pull/2784 [#2764]: https://github.com/ashhart/TensorFold/issues/2764 +[#2977]: https://github.com/ashhart/TensorFold/pull/2977 diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index dd413c25a96..9f9a7f520e8 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -39,6 +39,9 @@ tensorfold-native serve nvidia/Qwen3.8-Flash-Next-NVFP4 --tp 2 --rank 0 --master checkpoint's 124 GiB on a GB10), and both wait up to 30 minutes for the other. `TENSORFOLD_CHECKPOINT_SHA256`, when set the same on both ranks, replaces the digest and keeps the wait at a minute. - `--prompt-cache-gib 0` is required with drafts at `--tp 2`, and `--idle-unload` is refused with `--tp`. +- A prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's + rows and gathers the sums: rank 0 sends 3.43 MiB a prompt token over two GB10s' RoCE link instead of 5.81, with the + same bytes in every written row. Decode rounds and verify windows keep the all-reduce. - At `--tp 2` each stream's window verifies alone, so concurrent streams do not yet share a verify pass. ## Experimental K32 prompt experts From f803fa3a854259797dcaf4d348496718be5c2fdb Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:15:45 +0000 Subject: [PATCH 004/186] changelog: #2977's line from lines/2977.txt in place of its own --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ddc59ee0c1..5aaa3437a91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Flash Next on CUDA at `--tp 2`: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows and gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every reduced value is the same single add and the slots are summed in the writeback's order, so replies keep their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (an 8,192-token prompt: 7.12 to 5.62 s and 3.61 to 2.70 s) (#2977). Thanks to @swbiggart. + - Tests: `tf-cuda-test flashnext-rank RANK PORT timeout` passes again on two GPUs. NCCL start-up keeps its 30 s bound and only the stalled collective gets 200 ms, and both ranks meet once their communicators are up, so `timeout` and `eof` fail the collective under test, not the start-up (#2976, fixes #2975). Thanks to @swbiggart. - Tokenizer: a split pattern is held to a work allowance per piece (2,048 instructions a byte, counted, not timed) and refused at load when too deep or too large, so a crafted tokenizer or prompt can't hang the pre-tokenizer; every split of the patterns we serve is unchanged (#3071). @@ -59,7 +61,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. -- Flash Next at `--tp 2` on CUDA: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows, then gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every summed value is the same IEEE add and the slots are summed in the writeback's order, so the written rows keep their bytes and replies their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (8,192 tokens: 7.12 to 5.62 s, and 3.61 to 2.70 s); prompts of 128 to 512 tokens are 6.0% to 18.3% sooner on the GB10s (#2977). - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). From 40ec17f24a8b9a579de5dd024f61977a4270d35d Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 06:21:44 -0800 Subject: [PATCH 005/186] Remove unused Qwen27 GDN prompt tail scratch --- zig/kernels/metal/core/deltanet.metal | 13 +++---------- zig/src/families/qwen27/gdn.zig | 6 ++---- zig/src/families/qwen27/gdn_contract.zig | 2 -- zig/src/families/qwen27/gdn_pool.zig | 9 +++------ zig/src/families/qwen27/gdn_pool_gpu.zig | 3 +-- zig/src/families/qwen27/gdn_test.zig | 16 ++++++++++++++++ 6 files changed, 25 insertions(+), 24 deletions(-) diff --git a/zig/kernels/metal/core/deltanet.metal b/zig/kernels/metal/core/deltanet.metal index 3d4bf5a3e4f..9e76b801d87 100644 --- a/zig/kernels/metal/core/deltanet.metal +++ b/zig/kernels/metal/core/deltanet.metal @@ -54,7 +54,6 @@ inline void tf_delta_gdn_pre_body(device const bfloat* projection, device bfloat* value, device float* decay, device bfloat* mixing, - device bfloat* tails, uint3 group, uint lane) { const uint row = group.z, head = group.y; @@ -78,11 +77,6 @@ inline void tf_delta_gdn_pre_body(device const bfloat* projection, activated[part] = float(bfloat(rounded * sigmoid)); if (STEP) { for (uint tap = 0; tap < keep; ++tap) next_history[(row_slot[row] * keep + tap) * channels + channel] = bfloat(inputs[tap + 1]); - } else { - for (uint tap = 0; tap < keep; ++tap) { - const uint source = windows[row * p.taps + tap + 1]; - tails[(row * keep + tap) * channels + channel] = source < keep ? history[(row_slot[row] * keep + source) * channels + channel] : projection[(source - keep) * channels + channel]; - } } } if (head < 2 * p.nk) { @@ -126,11 +120,10 @@ kernel void tf_delta_gdn_pre(device const bfloat* projection [[buffer(0)]], device bfloat* value [[buffer(12)]], device float* decay [[buffer(13)]], device bfloat* mixing [[buffer(14)]], - device bfloat* tails [[buffer(15)]], uint3 group [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { - if (p.dk == 128 && p.dv == 128) tf_delta_gdn_pre_body<4, false>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, tails, group, lane); - else tf_delta_gdn_pre_body<0, false>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, tails, group, lane); + if (p.dk == 128 && p.dv == 128) tf_delta_gdn_pre_body<4, false>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, group, lane); + else tf_delta_gdn_pre_body<0, false>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, group, lane); } // tf_delta_gdn_pre_body<4, true> over 4 simdgroups a head; each lane's 4 squares summed in order before simd_sum. @@ -206,7 +199,7 @@ kernel void tf_delta_gdn_pre_step(device const bfloat* projection [[buffer(0)]], uint lane [[thread_index_in_simdgroup]]) { threadgroup float shared[129]; if (p.dk == 128 && p.dv == 128) tf_delta_gdn_pre_step4(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, shared, group, sg, lane); - else if (sg == 0) tf_delta_gdn_pre_body<0, true>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, nullptr, group, lane); + else if (sg == 0) tf_delta_gdn_pre_body<0, true>(projection, a, b, history, weights, decay_weight, time_bias, windows, row_slot, p, query, key, value, decay, mixing, group, lane); } // One tree node from its parent's state: tf_delta_delta_update and tf_delta_delta_read as the snapshot path runs them. inline float4 tf_delta_tree_node(float4 parent_state, device const bfloat* query, device const bfloat* key, device const bfloat* value, diff --git a/zig/src/families/qwen27/gdn.zig b/zig/src/families/qwen27/gdn.zig index 74591e9cc73..003a8b9cdad 100644 --- a/zig/src/families/qwen27/gdn.zig +++ b/zig/src/families/qwen27/gdn.zig @@ -42,7 +42,6 @@ pub const Buffers = struct { value: Span, decay: Span, mixing: Span, - tails: Span, snapshots: Span, recurrent: Span, output: Span, @@ -76,7 +75,6 @@ pub const Buffers = struct { inline for (.{ "query", "key" }) |name| try @field(b, name).check(sizes.query); inline for (.{ "value", "recurrent", "output" }) |name| try @field(b, name).check(sizes.value); try b.decay.check(sizes.decay); - try b.tails.check(sizes.tails); try b.snapshots.check(sizes.snapshots); try b.parents.check(try abi.countBytes(i32, p.rows, 1)); try b.windows.check(try abi.countBytes(u32, p.rows, p.taps)); @@ -90,7 +88,7 @@ pub const Buffers = struct { const reads = [_]Span{ b.qkv, b.a, b.b, b.z, b.conv_weight, b.a_log, b.dt, b.norm, b.history, b.state, b.parents, b.windows, b.row_slots, b.segments, b.keeps, b.kept_rows, b.slot_steps }; // an in-place round's next state is its state: each scan thread reads its own words before it writes them const in_place = b.next_state.buffer.id == b.state.buffer.id and b.next_state.offset == b.state.offset and b.next_state.bytes == b.state.bytes; - const writes = [_]Span{ b.query, b.key, b.value, b.decay, b.mixing, b.tails, b.snapshots, b.recurrent, b.output, if (in_place) .{ .buffer = b.state.buffer, .bytes = 0 } else b.next_state, b.next_history, b.before orelse Span{ .buffer = b.state.buffer, .bytes = 0 } }; + const writes = [_]Span{ b.query, b.key, b.value, b.decay, b.mixing, b.snapshots, b.recurrent, b.output, if (in_place) .{ .buffer = b.state.buffer, .bytes = 0 } else b.next_state, b.next_history, b.before orelse Span{ .buffer = b.state.buffer, .bytes = 0 } }; inline for (writes, 0..) |write, index| { inline for (reads) |read| if (write.overlaps(read)) return error.AliasedGdnState; inline for (writes[0..index]) |other| if (write.overlaps(other)) return error.AliasedGdnState; @@ -123,7 +121,7 @@ pub const Round = struct { e.setPipeline(r.pipelines.pre); bind(e, .{ r.b.qkv, r.b.a, r.b.b, r.b.history, r.b.conv_weight, r.b.a_log, r.b.dt, r.b.windows, r.b.row_slots }, 0); e.setValue(r.p, 9); - bind(e, .{ r.b.query, r.b.key, r.b.value, r.b.decay, r.b.mixing, r.b.tails }, 10); + bind(e, .{ r.b.query, r.b.key, r.b.value, r.b.decay, r.b.mixing }, 10); e.dispatchGroups(.{ .width = 1, .height = 2 * r.p.nk + r.p.nv, .depth = r.p.rows }, .{ .width = 32, .height = 1, .depth = 1 }); e.barrier(); } diff --git a/zig/src/families/qwen27/gdn_contract.zig b/zig/src/families/qwen27/gdn_contract.zig index 62ad1303907..c029a7cba65 100644 --- a/zig/src/families/qwen27/gdn_contract.zig +++ b/zig/src/families/qwen27/gdn_contract.zig @@ -76,7 +76,6 @@ pub const Bytes = struct { state: usize, snapshots: usize, history: usize, - tails: usize, conv_weight: usize, a_log: usize, dt: usize, @@ -97,7 +96,6 @@ pub const Bytes = struct { .state = try countBytes(f32, s.state, p.slots), .snapshots = try countBytes(f32, s.state, @min(p.rows, snapshot_rows)), .history = try countBytes(u16, s.conv, p.slots), - .tails = try countBytes(u16, s.conv, p.rows), .conv_weight = try countBytes(u8, s.qkv * s.c.conv_kernel, byteWidth(k.conv)), .a_log = try countBytes(u8, s.c.v_heads, byteWidth(k.a_log)), .dt = try countBytes(u8, s.c.v_heads, byteWidth(k.dt)), diff --git a/zig/src/families/qwen27/gdn_pool.zig b/zig/src/families/qwen27/gdn_pool.zig index c61623ea663..6d3a407716b 100644 --- a/zig/src/families/qwen27/gdn_pool.zig +++ b/zig/src/families/qwen27/gdn_pool.zig @@ -65,7 +65,6 @@ pub const Budget = struct { decay: usize, mixing: usize, snapshots: usize, - conv_tails: usize, scratch: usize, pub fn init(shape: abi.Shape, layers: usize, slots: usize, rows: usize) !Budget { @@ -77,7 +76,6 @@ pub const Budget = struct { const decay = try abi.countBytes(f32, shape.c.v_heads, layer_rows); const mixing = try abi.countBytes(u16, shape.c.v_heads, layer_rows); const snapshot = try abi.countBytes(f32, shape.state, @min(rows, abi.snapshot_rows)); // trees only: a chain keeps no snapshots - const tails = try abi.countBytes(u16, shape.conv, rows); const outputs = try abi.countBytes(u16, shape.value, rows * 2); return .{ .shape = shape, @@ -93,13 +91,12 @@ pub const Budget = struct { .decay = decay, .mixing = mixing, .snapshots = snapshot, - .conv_tails = tails, - .scratch = try std.math.add(usize, try std.math.add(usize, snapshot, tails), outputs), + .scratch = try std.math.add(usize, snapshot, outputs), }; } - pub fn allocationSizes(b: Budget) [12]usize { - return .{ b.state, b.history, b.history, b.projections, b.query, b.key, b.value, b.decay, b.mixing, b.snapshots, b.conv_tails, b.rows * b.shape.value * 2 }; + pub fn allocationSizes(b: Budget) [11]usize { + return .{ b.state, b.history, b.history, b.projections, b.query, b.key, b.value, b.decay, b.mixing, b.snapshots, b.rows * b.shape.value * 2 }; } pub fn stateOffset(b: Budget, layer: usize, slot: usize) !usize { diff --git a/zig/src/families/qwen27/gdn_pool_gpu.zig b/zig/src/families/qwen27/gdn_pool_gpu.zig index ced5aea0e31..5d8931162d3 100644 --- a/zig/src/families/qwen27/gdn_pool_gpu.zig +++ b/zig/src/families/qwen27/gdn_pool_gpu.zig @@ -12,7 +12,7 @@ const Kernels = @import("gdn_library.zig").Kernels; pub const PromptMarks = @import("gdn_marks_gpu.zig").Bank; -const Field = enum { state, history, next_history, qkv, query, key, value, decay, mixing, snapshots, tails, recurrent }; +const Field = enum { state, history, next_history, qkv, query, key, value, decay, mixing, snapshots, recurrent }; pub const Metadata = struct { parents: Ref, @@ -310,7 +310,6 @@ pub const Pool = struct { .value = try pool.layerSpan(.value, layer), .decay = try pool.layerSpan(.decay, layer), .mixing = try pool.layerSpan(.mixing, layer), - .tails = try refSpan(.{ .buffer = pool.get(.tails) }, bytes.tails), .snapshots = try refSpan(.{ .buffer = pool.get(.snapshots) }, bytes.snapshots), .recurrent = try refSpan(.{ .buffer = pool.get(.recurrent) }, bytes.value), .output = try refSpan(refs.output, bytes.value), diff --git a/zig/src/families/qwen27/gdn_test.zig b/zig/src/families/qwen27/gdn_test.zig index a58faf452ec..8fc83b10764 100644 --- a/zig/src/families/qwen27/gdn_test.zig +++ b/zig/src/families/qwen27/gdn_test.zig @@ -74,6 +74,11 @@ test "actual target widths and snapshot bytes preserve FP32 state" { try std.testing.expectEqual(@as(usize, 4 * 786432 * 4), sizes.state); try std.testing.expectEqual(@as(usize, 48 * 4), sizes.a_log); try std.testing.expectEqual(@as(usize, 48 * 2), sizes.dt); + try std.testing.expectEqual(@as(usize, 8 * shape.qkv * 2), sizes.qkv); + try std.testing.expectEqual(@as(usize, 8 * shape.c.k_heads * shape.c.dk * 2), sizes.query); + try std.testing.expectEqual(@as(usize, 8 * shape.value * 2), sizes.value); + try std.testing.expectEqual(@as(usize, 4 * shape.conv * 2), sizes.history); + try std.testing.expect(!@hasField(abi.Bytes, "tails")); } test "wrong shape flags and dimensions refuse before an encoder exists" { @@ -189,3 +194,14 @@ test "cross-stream conv rows and a broken retained path refuse before encoding" plan.kept_rows[0] = 1; try std.testing.expectError(error.BadGdnPoolPlan, @import("gdn_pool.zig").validatePlan(shape, 2, 32, plan)); } + +// Prompt-tail scratch is three BF16 taps per conv channel and row; live state and logs stay unchanged. +test "GDN prompt budget releases only the unused conv tails at 2048 rows" { + const shape = try abi.Shape.init(config()); + const b = try @import("gdn_pool.zig").Budget.init(shape, 48, 1, 2048); + const released = (shape.c.conv_kernel - 1) * shape.qkv * @sizeOf(u16) * b.rows; + try std.testing.expectEqual(@as(usize, 125829120), released); + try std.testing.expectEqual(@as(usize, 578813952) - released, b.scratch); + try std.testing.expectEqual(@as(usize, 4941545472) - released, b.totalBytes()); + try std.testing.expectEqual(@as(usize, 11), b.allocationSizes().len); +} From 47d08dda1b034c8b6e6448b0e2d872f6d8f6c4b0 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:15:48 +0000 Subject: [PATCH 006/186] changelog: #2915 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5aaa3437a91..298ca3c2cc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B on Metal: prompt passes no longer write the GDN tails buffer, which nothing read. Replies are unchanged; on an M5 Ultra, time to first token is about 0.5% shorter at 2K, 8K and 32K prompts, one and four streams (#2915, for #2905). + - Flash Next on CUDA at `--tp 2`: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows and gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every reduced value is the same single add and the slots are summed in the writeback's order, so replies keep their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (an 8,192-token prompt: 7.12 to 5.62 s and 3.61 to 2.70 s) (#2977). Thanks to @swbiggart. - Tests: `tf-cuda-test flashnext-rank RANK PORT timeout` passes again on two GPUs. NCCL start-up keeps its 30 s bound and only the stalled collective gets 200 ms, and both ranks meet once their communicators are up, so `timeout` and `eof` fail the collective under test, not the start-up (#2976, fixes #2975). Thanks to @swbiggart. From 4c0cfb7bc1821d15da0525be2001d6b5769c0723 Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 19:27:12 +0700 Subject: [PATCH 007/186] core: Admit and Ledger move to engine_admission.zig, unchanged --- CHANGELOG.md | 1 + zig/src/core/engine_admission.zig | 12 ++++++++++++ zig/src/core/engine_api.zig | 11 ++--------- 3 files changed, 15 insertions(+), 9 deletions(-) create mode 100644 zig/src/core/engine_admission.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 298ca3c2cc1..78b877e4c0d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,6 +60,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: the images route test fixture can no longer hang a test run on macOS. Its route closes its end of the socket on every exit and the client's writes are bounded, so a route that answers before reading a large body ends the write at once (#3010). +- Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, so `engine_api.zig` has room for the execution witness's hook; nothing a user sees changes (refs #2741). - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. diff --git a/zig/src/core/engine_admission.zig b/zig/src/core/engine_admission.zig new file mode 100644 index 00000000000..46acdebe1ba --- /dev/null +++ b/zig/src/core/engine_admission.zig @@ -0,0 +1,12 @@ +//! How an engine admits a request beside its running streams: a fit check, or a byte budget it promises from. +const kv_ledger = @import("kv_ledger.zig"); + +pub const Admit = struct { ctx: *anyopaque, fits: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) bool }; + +/// The bytes the lane host may promise at once, and what one request's caches cost at most (core/kv_ledger.zig). +pub const Ledger = struct { + ctx: *anyopaque, + total: u64, + cost: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) u64, + standing: kv_ledger.Standing = .{}, // streams paid at load, in `total`: the engine's own sequence (#1943) +}; diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 27e9551333b..e2ff1527b18 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -248,15 +248,8 @@ pub const Info = struct { } }; -pub const Admit = struct { ctx: *anyopaque, fits: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) bool }; - -/// The bytes the lane host may promise at once, and what one request's caches cost at most (core/kv_ledger.zig). -pub const Ledger = struct { - ctx: *anyopaque, - total: u64, - cost: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) u64, - standing: kv_ledger.Standing = .{}, // streams paid at load, in `total`: the engine's own sequence (#1943) -}; +pub const Admit = @import("engine_admission.zig").Admit; +pub const Ledger = @import("engine_admission.zig").Ledger; pub const PromptCacheCapacity = @import("engine_cache_plan.zig").PromptCacheCapacity; pub const PromptCachePlan = @import("engine_cache_plan.zig").PromptCachePlan; From ff3e811373ec32184cb29d87167659344493e55f Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 19:29:52 +0700 Subject: [PATCH 008/186] changelog: #2896's line names its pull request --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78b877e4c0d..7047fe877f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -61,6 +61,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: the images route test fixture can no longer hang a test run on macOS. Its route closes its end of the socket on every exit and the client's writes are bounded, so a route that answers before reading a large body ends the write at once (#3010). - Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, so `engine_api.zig` has room for the execution witness's hook; nothing a user sees changes (refs #2741). +- Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, so `engine_api.zig` has room for the execution witness's hook; nothing a user sees changes (#2896, refs #2741). Thanks to @kotevcode. - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. From 1ddbd509824b9f8f0f47e39a288dfe30247f867f Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:15:52 +0000 Subject: [PATCH 009/186] changelog: #2896's line from lines/2896.txt in place of its own --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7047fe877f1..f687fcb8802 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, making room for the execution witness's hook; nothing a user sees changes (#2896, refs #2741). Thanks to @kotevcode. + - Qwen3.8-27B on Metal: prompt passes no longer write the GDN tails buffer, which nothing read. Replies are unchanged; on an M5 Ultra, time to first token is about 0.5% shorter at 2K, 8K and 32K prompts, one and four streams (#2915, for #2905). - Flash Next on CUDA at `--tp 2`: a prompt pass wider than 64 rows reduce-scatters its routed experts' FP32 slots by row halves, sums each rank's rows and gathers the sums, where it all-reduced every slot; decode rounds and verify windows keep the all-reduce. Every reduced value is the same single add and the slots are summed in the writeback's order, so replies keep their tokens. Over two GB10s' RoCE link rank 0 sends 3.43 MiB a prompt token instead of 5.81, and cold prefill of 2k to 32k tokens is 26.5% to 26.8% faster there and 30.9% to 33.9% faster on two RTX PRO 6000s (an 8,192-token prompt: 7.12 to 5.62 s and 3.61 to 2.70 s) (#2977). Thanks to @swbiggart. @@ -60,8 +62,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: the images route test fixture can no longer hang a test run on macOS. Its route closes its end of the socket on every exit and the client's writes are bounded, so a route that answers before reading a large body ends the write at once (#3010). -- Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, so `engine_api.zig` has room for the execution witness's hook; nothing a user sees changes (refs #2741). -- Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, so `engine_api.zig` has room for the execution witness's hook; nothing a user sees changes (#2896, refs #2741). Thanks to @kotevcode. - CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823). - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. From 4a21b363573aa4dc93a97be89a11797ef830dd07 Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:26:47 +0000 Subject: [PATCH 010/186] qwen3 cuda: Qwen3 reranker checkpoints on the native CUDA engine A folder with 1_LogitScore opens as a reranker, as on HIP: its yes and no rows of the output head read on the host at load, the score and packed-decision hooks of lane.zig on the CUDA engine, and Info's rerank labels with decides, so /v1/rerank serves it. A decision's trunk runs once a pass: its keys stay in the cache's first rows, and each branch attends to them and to its own rows from the trunk's end. The CUDA GPU check adds the reranker's gate (yes - no and the score against transformers, the packed decision equal to each prompt's run alone). --- zig/build/qwen3_cuda.zig | 2 +- zig/src/families/qwen3/cuda/engine.zig | 8 +- zig/src/families/qwen3/cuda/forward_test.zig | 95 +++++++++++++++++++- zig/src/families/qwen3/cuda/lm.zig | 30 ++++--- zig/src/families/qwen3/cuda/native.zig | 54 +++++++++-- zig/src/families/qwen3/cuda/weights.zig | 6 +- zig/src/families/qwen3/qwen3_cuda.zig | 1 + 7 files changed, 169 insertions(+), 27 deletions(-) diff --git a/zig/build/qwen3_cuda.zig b/zig/build/qwen3_cuda.zig index 145f6b8afd2..b099264135d 100644 --- a/zig/build/qwen3_cuda.zig +++ b/zig/build/qwen3_cuda.zig @@ -29,7 +29,7 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const run = b.addRunArtifact(gpu); run.has_side_effects = true; // the checkpoints and oracle files named by TF_QWEN3_* are not cached inputs b.step("qwen3-cuda-gpu-build", "Compile the Qwen3 dense CUDA checks without running them").dependOn(&gpu.step); - b.step("qwen3-cuda-gpu-test", "Qwen3 embedding vectors on CUDA against transformers (TF_QWEN3_EMBED_DIR, _REF, _BAR)").dependOn(&run.step); + b.step("qwen3-cuda-gpu-test", "Qwen3 embedding vectors and reranker logits on CUDA against transformers (TF_QWEN3_EMBED_*, TF_QWEN3_RERANK_*)").dependOn(&run.step); all.dependOn(&gpu.step); const host = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3/checks.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true }) })); b.step("test-qwen3-checks-host", "Qwen3 checks' oracle pins, references and batching without a GPU").dependOn(&host.step); diff --git a/zig/src/families/qwen3/cuda/engine.zig b/zig/src/families/qwen3/cuda/engine.zig index 5d7539721f5..fb240934c1d 100644 --- a/zig/src/families/qwen3/cuda/engine.zig +++ b/zig/src/families/qwen3/cuda/engine.zig @@ -91,14 +91,12 @@ pub const Engine = struct { e.sized = true; } - /// Each input's final-normed last row as fp32 into `out` (ends.len rows of hidden); inputs end at `ends`. + /// Each input's final-normed last row as fp32 into `out`: after a shared `trunk` of tokens, inputs end at `ends`. pub fn run(e: *Engine, tokens: []const u32, trunk: u32, ends: []const u32, out: []f32) !void { const s = e.m.spec; - // a shared trunk comes with the reranker's packed decisions; embeddings run independent inputs - if (trunk != 0) return error.InvalidBatch; if (!e.sized or tokens.len == 0 or tokens.len > e.window or ends.len == 0 or ends.len > max_inputs or ends[ends.len - 1] != tokens.len or out.len != ends.len * s.hidden) return error.InvalidBatch; - var begin: u32 = 0; + var begin: u32 = trunk; for (ends) |x| { if (x <= begin) return error.InvalidBatch; begin = x; @@ -107,7 +105,7 @@ pub const Engine = struct { const ids = e.ids_host.slice(u32)[0..tokens.len]; @memcpy(ids, tokens); try store.view(&e.d, e.a.ids, 4 * tokens.len).uploadAsync(0, std.mem.sliceAsBytes(ids), e.s.handle); - try lm.Lm.forward(&e.l, &e.m, e.a, e.cap, ends); + try lm.Lm.forward(&e.l, &e.m, e.a, e.cap, trunk, ends); const rows = e.rows_host.slice(u16)[0 .. ends.len * s.hidden]; try store.view(&e.d, e.a.rows, 2 * rows.len).downloadAsync(0, std.mem.sliceAsBytes(rows), e.s.handle); try e.s.synchronize(); diff --git a/zig/src/families/qwen3/cuda/forward_test.zig b/zig/src/families/qwen3/cuda/forward_test.zig index dbbe8eff2ed..e75613f9b61 100644 --- a/zig/src/families/qwen3/cuda/forward_test.zig +++ b/zig/src/families/qwen3/cuda/forward_test.zig @@ -1,11 +1,22 @@ -//! The Qwen3 CUDA engine against transformers: embedding vectors, a batch's rows equal to alone's bits. +//! The Qwen3 CUDA engine against transformers: embedding vectors, the reranker's logits, a batch equal to alone. const std = @import("std"); const engine = @import("engine.zig"); const checks = @import("../checks.zig"); +const score = @import("../score.zig"); const t = std.testing; const gpa = t.allocator; +/// The engine as score.decidePacked drives any runner: a trunk once a run, its branches after it. +const Runner = struct { + e: *engine.Engine, + hidden: usize, + + pub fn run(r: @This(), tokens: []const u32, trunk_len: u32, ends: []const u32, rows: []f32) !void { + try r.e.run(tokens, trunk_len, ends, rows); + } +}; + test "Qwen3-Embedding: the card's and the long inputs on CUDA match transformers, a batch's rows equal alone's bits" { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -80,3 +91,85 @@ test "Qwen3-Embedding: the card's and the long inputs on CUDA match transformers } } } + +test "Qwen3-Reranker: each case's yes - no logit matches transformers, and the packed decision equals the runs alone" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try checks.env(a, "TF_QWEN3_RERANK_DIR") orelse return error.SkipZigTest; + const ref = try checks.readJson(a, try checks.env(a, "TF_QWEN3_RERANK_REF") orelse return error.SkipZigTest); + const cases = ref.object.get("cases").?.array.items; + const yes: u32 = @intCast(ref.object.get("yes").?.integer); + const no: u32 = @intCast(ref.object.get("no").?.integer); + var prompts: std.ArrayList([]const u32) = .empty; + var longest: usize = 0; + for (cases) |c| { + if (c.object.get("diff") == null) continue; // ids-only cases carry no score + const prompt = try checks.tokensOf(a, c.object.get("ids").?); + longest = @max(longest, prompt.len); + try prompts.append(a, prompt); + } + try t.expect(prompts.items.len >= 32); + const window = std.mem.alignForward(usize, @max(engine.least_window, longest), 512); + const e = try engine.Engine.load(gpa, t.io, dir, 0); + defer e.deinit(); + try e.size(window); + const hidden = e.m.spec.hidden; + var heads = score.Heads.of(&e.m, t.io, dir); + defer heads.deinit(gpa); + var worst_diff: f64 = 0; + var worst_score: f64 = 0; + for (cases) |c| { + const want_diff = c.object.get("diff") orelse continue; + const prompt = try checks.tokensOf(a, c.object.get("ids").?); + const row = try a.alloc(f32, hidden); + try e.run(prompt, 0, &.{@intCast(prompt.len)}, row); + var logits: [2]f64 = undefined; + try score.labelLogits(gpa, &heads, row, &.{ yes, no }, &logits); + const diff = logits[0] - logits[1]; + const rank = 1 / (1 + @exp(-diff)); + const ref_score = checks.number(c.object.get("score").?); + // a NaN would slip through @max: every value must be finite before it counts + for ([_]f64{ diff, rank, checks.number(want_diff), ref_score }) |v| try t.expect(std.math.isFinite(v)); + std.debug.print("rerank {s} {d} tokens: diff {d:.4} ref {d:.4} score {d:.6} ref {d:.6}\n", .{ c.object.get("kind").?.string, prompt.len, diff, checks.number(want_diff), rank, ref_score }); + worst_diff = @max(worst_diff, @abs(diff - checks.number(want_diff))); + worst_score = @max(worst_score, @abs(rank - ref_score)); + } + std.debug.print("rerank: {d} cases, worst |diff| {d:.4} (gate 0.25), worst |score| {d:.5} (gate 0.02)\n", .{ prompts.items.len, worst_diff, worst_score }); + try t.expect(worst_diff <= 0.25); + try t.expect(worst_score <= 0.02); + // the decision path: branches packed into few runs, each prompt's logits the bits its own run gives + var trunk_len: usize = std.math.maxInt(usize); + for (prompts.items) |p| trunk_len = @min(trunk_len, p.len); + trunk_len -= 1; + for (prompts.items[1..]) |p| { + var n: usize = 0; + while (n < trunk_len and p[n] == prompts.items[0][n]) n += 1; + trunk_len = n; + } + try t.expect(trunk_len > 0); + const trunk = prompts.items[0][0..trunk_len]; + const branches = try a.alloc([]const u32, prompts.items.len); + const counts = try a.alloc(usize, prompts.items.len); + var labels: std.ArrayList(u32) = .empty; + for (prompts.items, branches, counts, 0..) |p, *b, *n, i| { + b.* = p[trunk_len..]; + // unequal label counts: yes and no, then yes alone + n.* = if (i % 2 == 0) 2 else 1; + try labels.append(a, yes); + if (i % 2 == 0) try labels.append(a, no); + } + const packed_logits = try a.alloc(f64, labels.items.len); + try score.decidePacked(a, Runner{ .e = e, .hidden = hidden }, &heads, trunk, branches, counts, labels.items, packed_logits, window, engine.max_inputs); + var at: usize = 0; + for (prompts.items, 0..) |prompt, i| { + const n = counts[i]; + const row = try a.alloc(f32, hidden); + try e.run(prompt, 0, &.{@intCast(prompt.len)}, row); + var alone: [2]f64 = undefined; + try score.labelLogits(gpa, &heads, row, labels.items[at..][0..n], alone[0..n]); + for (packed_logits[at..][0..n], alone[0..n]) |x, y| try t.expectEqual(@as(u64, @bitCast(y)), @as(u64, @bitCast(x))); + at += n; + } + std.debug.print("rerank decide: {d} branches (trunk {d}) packed at a {d}-token window equal their runs alone\n", .{ prompts.items.len, trunk_len, window }); +} diff --git a/zig/src/families/qwen3/cuda/lm.zig b/zig/src/families/qwen3/cuda/lm.zig index a019fc4c756..6dfb43b56d9 100644 --- a/zig/src/families/qwen3/cuda/lm.zig +++ b/zig/src/families/qwen3/cuda/lm.zig @@ -165,6 +165,14 @@ test "the geometry limits: the attention's head size and groups, and the row ker try std.testing.expectEqual(Geometry.widths, geometry(s)); } +test "the published Qwen3 sizes, embedding and reranker alike, are within the geometry limits" { + const sizes = [_][3]usize{ .{ 1024, 3072, 16 }, .{ 2560, 9728, 32 }, .{ 4096, 12288, 32 } }; // hidden, inter, q heads + for (sizes) |z| { + const spec: config.Spec = .{ .layers = 28, .hidden = z[0], .heads = z[2], .kv_heads = 8, .head_dim = 128, .inter = z[1], .vocab = 151669, .eps = 1e-6, .theta = 1e6, .max_positions = 40960, .tied = z[0] < 4096 }; + try std.testing.expectEqual(Geometry.ok, geometry(spec)); + } +} + /// The shared operators one pass runs, over one stream. pub const Lm = struct { d: *const Driver, @@ -221,11 +229,11 @@ pub const Lm = struct { lm.g.deinit(); } - /// Every input's final-normed last row into the block's `rows`, inputs independent and ending at `ends`. - pub fn forward(lm: *Lm, m: *const weights.Model, a: Addrs, cap: usize, ends: []const u32) !void { + /// Every input's final-normed last row into `rows`: after a shared `trunk` of rows, inputs end at `ends`. + pub fn forward(lm: *Lm, m: *const weights.Model, a: Addrs, cap: usize, trunk: usize, ends: []const u32) !void { const s = m.spec; try lm.emb.bf16(lm.s, a.ids, m.embed, a.x, ends[ends.len - 1], s.hidden); - for (m.layers, 0..) |*l, i| try layerRows(lm, m, l, a, cap, ends, ends[ends.len - 1], i == 0); + for (m.layers, 0..) |*l, i| try layerRows(lm, m, l, a, cap, trunk, ends, ends[ends.len - 1], i == 0); for (ends, 0..) |e, i| { // the last layer's residual add and the final norm, for the one row an input pools const row = (e - 1) * s.hidden * 2; @@ -255,7 +263,7 @@ pub const Lm = struct { }; /// One decoder layer over every row; `first` is layer 0, whose residual is the embedding rows themselves. -fn layerRows(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs, cap: usize, ends: []const u32, total: usize, first: bool) !void { +fn layerRows(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs, cap: usize, trunk: usize, ends: []const u32, total: usize, first: bool) !void { const s = m.spec; const qw = s.heads * s.head_dim; const kvw = s.kv_heads * s.head_dim; @@ -265,10 +273,12 @@ fn layerRows(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs try lm.g.run(lm.s, a.normed, l.q, a.q, total, qw, s.hidden); try lm.g.run(lm.s, a.normed, l.k, a.k, total, kvw, s.hidden); try lm.g.run(lm.s, a.normed, l.v, a.v, total, kvw, s.hidden); - var begin: usize = 0; + // the trunk's keys stay in the cache's first rows; each input after it attends to them and its own, from there + if (trunk > 0) try attention(lm, m, l, a, cap, 0, trunk, 0); + var begin: usize = trunk; for (ends) |e| { defer begin = e; - try attention(lm, m, l, a, cap, begin, e); + try attention(lm, m, l, a, cap, begin, e, trunk); } try lm.g.run(lm.s, a.attn, l.o, a.z, total, s.hidden, qw); try lm.nm.addRmsH(lm.s, a.z, a.x, l.post_norm, .f32, a.x, a.normed, null, total, s.hidden, eps); @@ -279,8 +289,8 @@ fn layerRows(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs try lm.g.run(lm.s, a.gate, l.down, a.z, total, s.hidden, s.inter); } -/// One input's rows [begin, end): QK norm, RoPE and the cache write, then causal attention over that cache alone. -fn attention(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs, cap: usize, begin: usize, end: usize) !void { +/// Rows [begin, end) standing at positions from `at`: QK norm, RoPE and the cache write, then causal attention. +fn attention(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs, cap: usize, begin: usize, end: usize, at: usize) !void { const s = m.spec; const hd = s.head_dim; const qw = s.heads * hd; @@ -300,8 +310,8 @@ fn attention(lm: *Lm, m: *const weights.Model, l: *const weights.Layer, a: Addrs const hd64: i64 = @intCast(hd); const src: cuda.qk_rows.Src = .{ .q = a.q + (begin + p) * qw * 2, .q_ld = ld, .q_hs = hd64, .k = a.k + k, .k_ld = kld, .k_hs = hd64, .v = a.v + k, .v_ld = kld, .v_hs = hd64 }; const dst: cuda.qk_rows.Out = .{ .q = out, .k = a.kb, .v = a.kb + attn_rows * kvw * 2, .kc = a.kc, .vc = a.vc, .len = cap }; - // a block's first row stands at position p: the position table is the meta the chunks read, shifted too - const pos = a.pos + 4 * p; + // a block's first row stands at position at + p: the position table is the meta the chunks read, shifted too + const pos = a.pos + 4 * (at + p); try lm.qk.rowsFrom(lm.s, src, heads, shape, pos, null, dst, rows); try lm.chunks(s, dst.q, a, pos, rows, nch, scale); try lm.merge(s, out, a, pos, rows, nch); diff --git a/zig/src/families/qwen3/cuda/native.zig b/zig/src/families/qwen3/cuda/native.zig index e977918ea58..3675e61f8ab 100644 --- a/zig/src/families/qwen3/cuda/native.zig +++ b/zig/src/families/qwen3/cuda/native.zig @@ -5,6 +5,7 @@ const lanes = @import("lanes"); const api = @import("engine_api"); const engine = @import("engine.zig"); const lane = @import("../lane.zig"); +const score = @import("../score.zig"); const folder = @import("../head.zig"); const Engine = engine.Engine; @@ -45,6 +46,8 @@ const Owned = struct { core: lanes.Engine, host: api.LaneHost, startup: []const u8, + /// The checkpoint folder, kept for the reranker's head rows. + dir: []const u8, }; /// The context this thread's CUDA calls go to: open's thread at startup, the lane host's thread at rounds. @@ -78,6 +81,7 @@ fn release(ptr: *anyopaque) void { o.p.deinit(); o.e.deinit(); o.gpa.free(o.startup); + o.gpa.free(o.dir); o.gpa.destroy(o); } @@ -101,6 +105,18 @@ fn scratchOf(e: *Engine, window: usize) anyerror!usize { return e.scratchBytes(window); } +/// Whether a reranker's yes or no token is past the vocabulary's rows: its head row would not exist. +fn labelOutside(l: api.Rerank, vocab: usize) bool { + return l.yes >= vocab or l.no >= vocab; +} + +test "a reranker's labels must both be rows of the vocabulary" { + try std.testing.expect(!labelOutside(.{ .yes = 9693, .no = 2152 }, 151669)); + try std.testing.expect(!labelOutside(.{ .yes = 0, .no = 2 }, 3)); + try std.testing.expect(labelOutside(.{ .yes = 3, .no = 0 }, 3)); + try std.testing.expect(labelOutside(.{ .yes = 0, .no = 3 }, 3)); +} + /// The engine for `o.dir` on `device`, its head from the folder and its window from the memory budget. fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, device: u32, problem: *[]const u8, comptime budget: type, comptime windows: type, comptime libs: type) anyerror!api.Opened { // `budget` and `windows` are the native module's own file imports; the family cannot open them again. @@ -134,13 +150,9 @@ fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, device: u32 return refuse(a, problem, "{s}{d}: this machine has {d} CUDA device{s} (0 to {d})", .{ named, device, count, if (count == 1) "" else "s", count - 1 }); } const head = try folder.headOf(a, io, o.dir, "CUDA", problem); - const normalize = switch (head) { - .embedding => |h| h.normalize, - .rerank => return refuse(a, problem, "{s} is a Qwen3 reranker (1_LogitScore): the CUDA engine serves Qwen3 embedding checkpoints so far", .{o.dir}), - }; const own = try gpa.create(Owned); errdefer gpa.destroy(own); - own.* = .{ .gpa = gpa, .e = undefined, .p = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .startup = undefined }; + own.* = .{ .gpa = gpa, .e = undefined, .p = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .startup = undefined, .dir = undefined }; own.e = Engine.load(gpa, io, o.dir, @intCast(device)) catch |err| { problem.* = libs.explain(err) orelse try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} ({s})", .{ o.dir, @errorName(err) }); return error.Refused; @@ -168,16 +180,40 @@ fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, device: u32 }; e.size(window) catch |err| return refuse(a, problem, "the native CUDA engine cannot allocate its scratch ({s})", .{@errorName(err)}); + const normalize = switch (head) { + .embedding => |h| h.normalize, + .rerank => false, + }; + const labels: ?api.Rerank = switch (head) { + .embedding => null, + .rerank => |l| l, + }; + if (labels) |l| if (labelOutside(l, e.m.spec.vocab)) + return refuse(a, problem, "1_LogitScore/config.json: tokens {d} and {d} are not both in this {d}-token vocabulary", .{ l.yes, l.no, e.m.spec.vocab }); + own.dir = try gpa.dupe(u8, o.dir); + errdefer gpa.free(own.dir); own.p = Prompt.init(gpa, e, normalize); errdefer own.p.deinit(); + if (labels) |l| { + own.p.heads = score.Heads.of(&e.m, io, own.dir); + own.p.table.score = Prompt.scoreLabels; + own.p.table.decide = Prompt.decideLabels; + // both label rows read at load: a bad head fails startup, and the first request reads nothing + for ([_]u32{ l.yes, l.no }) |id| _ = own.p.heads.?.row(gpa, id) catch |err| + return refuse(a, problem, "the native CUDA engine cannot read lm_head row {d} ({s})", .{ id, @errorName(err) }); + } own.cfg = try lanes.Config.init(gpa, .{}, 1, 0); errdefer own.cfg.deinit(gpa); own.clock = .{ .io = io }; own.core = lanes.Engine.init(gpa, &own.cfg, own.p.backend(), own.clock.clock()); errdefer own.core.deinit(); const shrunk: []const u8 = if (window < asked) " (shrunk from the asked window to fit the budget)" else ""; - own.startup = try std.fmt.allocPrint(gpa, "CUDA sm{d} device {d}: qwen3 embedding, model {d:.2} GiB; {d}-token passes{s}, {d:.2} GiB scratch, of {d:.1} GiB left after a {d:.1} GiB reserve; {s}", .{ - cap, device, gibs(e.weights), window, shrunk, gibs(e.scratchBytes(window) catch 0), gibs(room), gibs(pool.reserve), if (normalize) "vectors unit-normed" else "vectors as pooled", + own.startup = try std.fmt.allocPrint(gpa, "CUDA sm{d} device {d}: qwen3 {s}, model {d:.2} GiB; {d}-token passes{s}, {d:.2} GiB scratch, of {d:.1} GiB left after a {d:.1} GiB reserve; {s}", .{ + cap, device, if (head == .rerank) "reranker" else "embedding", gibs(e.weights), window, shrunk, gibs(e.scratchBytes(window) catch 0), gibs(room), gibs(pool.reserve), + switch (head) { + .embedding => |h| if (h.normalize) "vectors unit-normed" else "vectors as pooled", + .rerank => "label scores from lm_head rows read at load, on the host", + }, }); errdefer gpa.free(own.startup); const spec = e.m.spec; @@ -185,7 +221,9 @@ fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, device: u32 .lanes = 1, .context_window = @intCast(window), .text_generation = false, - .embeddings = .{ .dimensions = @intCast(spec.hidden), .vocabulary = @intCast(spec.vocab), .max_batch_tokens = @intCast(window), .max_inputs = @intCast(engine.max_inputs), .mrl_min = @min(32, @as(u32, @intCast(spec.hidden))) }, + .embeddings = if (head == .rerank) null else .{ .dimensions = @intCast(spec.hidden), .vocabulary = @intCast(spec.vocab), .max_batch_tokens = @intCast(window), .max_inputs = @intCast(engine.max_inputs), .mrl_min = @min(32, @as(u32, @intCast(spec.hidden))) }, + .decides = head == .rerank, + .rerank = labels, .reasoning = reasoning, .startup = own.startup, }); diff --git a/zig/src/families/qwen3/cuda/weights.zig b/zig/src/families/qwen3/cuda/weights.zig index e7417b46a35..b84030b9fd8 100644 --- a/zig/src/families/qwen3/cuda/weights.zig +++ b/zig/src/families/qwen3/cuda/weights.zig @@ -1,4 +1,4 @@ -//! A Qwen3 dense checkpoint on the device: BF16 tensors as stored, norms widened to fp32. +//! A Qwen3 dense checkpoint on the device: BF16 tensors as stored, norms widened to fp32, the output head on the host. const std = @import("std"); const core = @import("core"); const cuda = @import("cuda"); @@ -27,6 +27,8 @@ pub const Layer = struct { pub const Model = struct { gpa: Allocator, spec: config.Spec, + /// The stored name of the output head's rows: a reranker scores them on the host (names.headKey). + head: []const u8, /// The embedding table [vocab, hidden] the token rows come from. embed: u64, final_norm: u64, @@ -39,7 +41,7 @@ pub const Model = struct { var ck = try Checkpoint.openModel(gpa, io, dir); defer ck.close(); ck.scope = names.scopeOf(&ck); - var m: Model = .{ .gpa = gpa, .spec = spec, .embed = 0, .final_norm = 0, .layers = &.{}, .st = .{ .gpa = gpa, .d = d } }; + var m: Model = .{ .gpa = gpa, .spec = spec, .head = names.headKey(spec.tied, ck.scope), .embed = 0, .final_norm = 0, .layers = &.{}, .st = .{ .gpa = gpa, .d = d } }; errdefer m.deinit(); m.embed = try projection(&ck, &m.st, "embed_tokens.weight", spec.vocab, spec.hidden); m.final_norm = try norm(&ck, &m.st, "norm.weight", spec.hidden); diff --git a/zig/src/families/qwen3/qwen3_cuda.zig b/zig/src/families/qwen3/qwen3_cuda.zig index 972fc98247f..87d2fd773e2 100644 --- a/zig/src/families/qwen3/qwen3_cuda.zig +++ b/zig/src/families/qwen3/qwen3_cuda.zig @@ -17,4 +17,5 @@ test { _ = score; _ = gemm; _ = engine; + _ = native; } From d5d31e696d66faf24e71b614a30ef158ca1ade4a Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:00:49 +0000 Subject: [PATCH 011/186] qwen3 cuda: the pass's arguments in a checked function with a host test; the docs and file headers name the reranker Engine.run's argument checks move to checkBatch, so the trunk, strictly rising ends, the window, the vocabulary and the rows are pinned by a host test (as on HIP). docs/api.md says the native HIP and CUDA engines serve the reranker; two comments said embedding only. --- docs/api.md | 4 +-- zig/src/families/qwen3/cuda/engine.zig | 47 +++++++++++++++++++++----- zig/src/families/qwen3/cuda/native.zig | 2 +- zig/src/native/cuda_qwen3.zig | 2 +- 4 files changed, 43 insertions(+), 12 deletions(-) diff --git a/docs/api.md b/docs/api.md index 0f22026e707..bb388574e91 100644 --- a/docs/api.md +++ b/docs/api.md @@ -168,8 +168,8 @@ For decisions, `chat_template_kwargs` may be omitted, null, or an object contain ## Rerank `POST /v1/rerank` (also `/rerank`, and by `PUT` as SGLang takes it) scores each document against a query with a -reranker checkpoint. The native HIP engine serves `Qwen/Qwen3-Reranker-8B`: the engine declares the reranker and -its label tokens (`1_LogitScore/config.json`), and its startup line says `qwen3 reranker`. Every other server answers +reranker checkpoint. The native HIP and CUDA engines serve the `Qwen/Qwen3-Reranker` checkpoints: the engine declares +the reranker and its label tokens (`1_LogitScore/config.json`), and its startup line says `qwen3 reranker`. Every other server answers the route with HTTP 400 "this model is not a reranker", and a reranker server answers chat, completion, Responses, Messages, speech, decision, System One and embedding requests with HTTP 400 "this model is a reranker: use POST /v1/rerank". diff --git a/zig/src/families/qwen3/cuda/engine.zig b/zig/src/families/qwen3/cuda/engine.zig index fb240934c1d..1744331b5d6 100644 --- a/zig/src/families/qwen3/cuda/engine.zig +++ b/zig/src/families/qwen3/cuda/engine.zig @@ -14,6 +14,43 @@ pub const max_inputs: usize = 128; /// The least window the budget may shrink a load to. pub const least_window: usize = 2048; +/// A pass's arguments before any upload: the trunk's tokens, then inputs ending at strictly rising `ends` past it. +pub fn checkBatch(tokens: []const u32, trunk: u32, ends: []const u32, out_len: usize, window: usize, vocab: usize, hidden: usize) error{InvalidBatch}!void { + if (tokens.len == 0 or tokens.len > window or ends.len == 0 or ends.len > max_inputs or + ends[ends.len - 1] != tokens.len or out_len != ends.len * hidden) return error.InvalidBatch; + var begin: u32 = trunk; + for (ends) |x| { + if (x <= begin) return error.InvalidBatch; + begin = x; + } + for (tokens) |t| if (t >= vocab) return error.InvalidBatch; +} + +test "a pass's arguments are checked before any upload: the trunk, the ends, the window, the vocabulary and the rows" { + const tokens = [_]u32{ 1, 2, 3, 4 }; + try checkBatch(&tokens, 0, &.{ 2, 4 }, 6, 8, 10, 3); + try checkBatch(&tokens, 1, &.{ 2, 4 }, 6, 4, 10, 3); + const bad = error.InvalidBatch; + try std.testing.expectError(bad, checkBatch(&tokens, 2, &.{ 2, 4 }, 6, 8, 10, 3)); // the first input is empty + try std.testing.expectError(bad, checkBatch(&tokens, 3, &.{ 2, 4 }, 6, 8, 10, 3)); // the trunk passes it + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 3, 2, 4 }, 9, 8, 10, 3)); // ends that fall + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 2, 2, 4 }, 9, 8, 10, 3)); // an empty input between + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{}, 0, 8, 10, 3)); + try std.testing.expectError(bad, checkBatch(&.{}, 0, &.{0}, 3, 8, 10, 3)); + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 2, 3 }, 6, 8, 10, 3)); // the last end is not the count + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 2, 4 }, 5, 8, 10, 3)); // the rows do not fit `out` + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 2, 4 }, 6, 3, 10, 3)); // past the window + try std.testing.expectError(bad, checkBatch(&tokens, 0, &.{ 2, 4 }, 6, 8, 4, 3)); // a token outside the vocabulary + var many: [max_inputs + 1]u32 = undefined; + var at: u32 = 0; + for (&many) |*x| { + at += 1; + x.* = at; + } + var ids: [max_inputs + 1]u32 = @splat(1); + try std.testing.expectError(bad, checkBatch(&ids, 0, &many, many.len * 3, many.len, 10, 3)); // more inputs than the rows hold +} + pub const Engine = struct { gpa: Allocator, d: cuda.Driver, @@ -94,14 +131,8 @@ pub const Engine = struct { /// Each input's final-normed last row as fp32 into `out`: after a shared `trunk` of tokens, inputs end at `ends`. pub fn run(e: *Engine, tokens: []const u32, trunk: u32, ends: []const u32, out: []f32) !void { const s = e.m.spec; - if (!e.sized or tokens.len == 0 or tokens.len > e.window or ends.len == 0 or ends.len > max_inputs or - ends[ends.len - 1] != tokens.len or out.len != ends.len * s.hidden) return error.InvalidBatch; - var begin: u32 = trunk; - for (ends) |x| { - if (x <= begin) return error.InvalidBatch; - begin = x; - } - for (tokens) |t| if (t >= s.vocab) return error.InvalidBatch; + if (!e.sized) return error.InvalidBatch; + try checkBatch(tokens, trunk, ends, out.len, e.window, s.vocab, s.hidden); const ids = e.ids_host.slice(u32)[0..tokens.len]; @memcpy(ids, tokens); try store.view(&e.d, e.a.ids, 4 * tokens.len).uploadAsync(0, std.mem.sliceAsBytes(ids), e.s.handle); diff --git a/zig/src/families/qwen3/cuda/native.zig b/zig/src/families/qwen3/cuda/native.zig index 3675e61f8ab..906836f7317 100644 --- a/zig/src/families/qwen3/cuda/native.zig +++ b/zig/src/families/qwen3/cuda/native.zig @@ -17,7 +17,7 @@ pub const model_type = "qwen3"; pub const formats: []const []const u8 = &.{"bf16"}; /// The window when neither --context nor the checkpoint's config names one (the Qwen3 8B serving default). pub const default_context: i64 = 32768; -/// An embedding engine writes no reply text, so no reasoning block (Info.reasoning). +/// A prompt-only engine writes no reply text, so no reasoning block (Info.reasoning). pub const reasoning = false; pub const drafter_refusal = "--drafter: Qwen3 dense checkpoints have no drafts to train; serve without --drafter"; pub const record_refusal = "--record: qwen3 checkpoints give no rows to record on CUDA yet"; diff --git a/zig/src/native/cuda_qwen3.zig b/zig/src/native/cuda_qwen3.zig index f0737dba8c8..6b60034d726 100644 --- a/zig/src/native/cuda_qwen3.zig +++ b/zig/src/native/cuda_qwen3.zig @@ -1,4 +1,4 @@ -//! Qwen3 dense embedding checkpoints on the native CUDA server: the flags its family does not take, then its open. +//! Qwen3 dense embedding and reranker checkpoints on the native CUDA server: the family's flag refusals, then its open. const std = @import("std"); const api = @import("engine_api"); const qwen3 = @import("qwen3_cuda"); From e879cf53b51c4fdd53efa9374b5015b6db83227a Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:01:34 +0000 Subject: [PATCH 012/186] qwen3 cuda: the reranker gate packs a decision at the chunk and block edges the cases' short common prefix never reaches A trunk of 700 tokens (past one 512-key chunk, off its multiple) with branches of 600, 513, 1 and 530 rows, cut from the longest case, packed in one decision; each branch's logits must equal its prompt's run alone, bit for bit. --- zig/src/families/qwen3/cuda/forward_test.zig | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/zig/src/families/qwen3/cuda/forward_test.zig b/zig/src/families/qwen3/cuda/forward_test.zig index e75613f9b61..140d72ace4e 100644 --- a/zig/src/families/qwen3/cuda/forward_test.zig +++ b/zig/src/families/qwen3/cuda/forward_test.zig @@ -172,4 +172,35 @@ test "Qwen3-Reranker: each case's yes - no logit matches transformers, and the p at += n; } std.debug.print("rerank decide: {d} branches (trunk {d}) packed at a {d}-token window equal their runs alone\n", .{ prompts.items.len, trunk_len, window }); + // edges the cases' short common prefix never reaches: a trunk off a 512-key chunk, branches of 1 to over 512 rows + var base = prompts.items[0]; + for (prompts.items) |p| { + if (p.len > base.len) base = p; + } + const edge_trunk: usize = 700; + const edge_lens = [_]usize{ 600, 513, 1, 530 }; + var edge_total: usize = edge_trunk; + for (edge_lens) |n| edge_total += n; + try t.expect(base.len >= edge_trunk + 600 and edge_total <= window); + var edge_branches: [edge_lens.len][]const u32 = undefined; + var edge_counts: [edge_lens.len]usize = undefined; + var edge_labels: std.ArrayList(u32) = .empty; + for (edge_lens, &edge_branches, &edge_counts, 0..) |n, *b, *c, i| { + b.* = base[edge_trunk..][0..n]; + c.* = if (i % 2 == 0) 2 else 1; + try edge_labels.append(a, yes); + if (i % 2 == 0) try edge_labels.append(a, no); + } + const edge_packed = try a.alloc(f64, edge_labels.items.len); + try score.decidePacked(a, Runner{ .e = e, .hidden = hidden }, &heads, base[0..edge_trunk], &edge_branches, &edge_counts, edge_labels.items, edge_packed, window, engine.max_inputs); + var edge_at: usize = 0; + for (edge_lens, edge_counts) |n, c| { + const row = try a.alloc(f32, hidden); + try e.run(base[0 .. edge_trunk + n], 0, &.{@intCast(edge_trunk + n)}, row); + var alone: [2]f64 = undefined; + try score.labelLogits(gpa, &heads, row, edge_labels.items[edge_at..][0..c], alone[0..c]); + for (edge_packed[edge_at..][0..c], alone[0..c]) |x, y| try t.expectEqual(@as(u64, @bitCast(y)), @as(u64, @bitCast(x))); + edge_at += c; + } + std.debug.print("rerank decide edges: trunk {d}, branches of {d}, {d}, {d} and {d} rows equal their runs alone\n", .{ edge_trunk, edge_lens[0], edge_lens[1], edge_lens[2], edge_lens[3] }); } From ac81368ed7124c857e0c14d5d5196f3f9f26a6d5 Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:59:54 +0000 Subject: [PATCH 013/186] changelog: Qwen3 rerankers on CUDA under Unreleased (#2995) --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index f687fcb8802..3dff3fd2036 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu. - CUDA: `Qwen/Qwen3-Embedding-0.6B`, `-4B` and `-8B` load on the native engine from the published BF16 weights and serve `/v1/embeddings` (last-token pooling, unit-normed, `dimensions` crops). On one GB10 the worst cosine distance to float32 is 1.55e-4, 1.34e-4 and 1.08e-4 (transformers' own BF16: 1.88e-4, 1.79e-4 and 1.77e-4), a request's vectors equal each input sent alone, and a reranker folder or a flag the family does not serve is refused by name (#2285). Thanks to @Arminova. +- CUDA: `Qwen/Qwen3-Reranker-0.6B`, `-4B` and `-8B` load on the native engine from the published BF16 weights and serve `/v1/rerank`, scoring a request's documents as one decision whose shared start runs once a pass. On one GB10 the worst score from float32 is 5.5e-5, 6.2e-4 and 2.9e-4 on 16 documents of up to 2,579 tokens (transformers' own BF16: 3.4e-5, 2.8e-4 and 2.2e-4), a document scored alone equals its score in a request of 16, and a prompt past the window is refused (#2995, part of #2261). Thanks to @Arminova. - Qwen3.8-27B on Metal: the startup line says how many streams draft ("for both streams", "for all 4 streams", "for 8 of 16 streams"), the smaller of the stream count and `TF_Q27_DRAFT_SLOTS`, instead of "for one of them"; log text only (#2806). - Server: `--profile fast` refuses as not planned, since a speed mode that changes output tokens is not shipped; `--profile fast-exact` now names the token gates it still needs, and an unknown spelling is told that `stable` is the only served profile (#2794). Thanks to @devopsfortunato. From e5943ff3e922c9b2ccead498037a1a6fca0b6fc3 Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:45:42 +0000 Subject: [PATCH 014/186] qwen3 cuda: the rerank docs name the checkpoints each native engine serves; the weights header names the head's stored name --- docs/api.md | 10 +++++----- zig/src/families/qwen3/cuda/weights.zig | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/api.md b/docs/api.md index bb388574e91..65f790f38d4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -168,11 +168,11 @@ For decisions, `chat_template_kwargs` may be omitted, null, or an object contain ## Rerank `POST /v1/rerank` (also `/rerank`, and by `PUT` as SGLang takes it) scores each document against a query with a -reranker checkpoint. The native HIP and CUDA engines serve the `Qwen/Qwen3-Reranker` checkpoints: the engine declares -the reranker and its label tokens (`1_LogitScore/config.json`), and its startup line says `qwen3 reranker`. Every other server answers -the route with HTTP 400 "this model is not a reranker", and a reranker server answers chat, completion, Responses, -Messages, speech, decision, System One and embedding requests with HTTP 400 "this model is a reranker: use POST -/v1/rerank". +reranker checkpoint. The native HIP engine serves `Qwen/Qwen3-Reranker-8B`, and the native CUDA engine `-0.6B`, `-4B` +and `-8B`: the engine declares the reranker and its label tokens (`1_LogitScore/config.json`), and its startup line +says `qwen3 reranker`. Every other server answers the route with HTTP 400 "this model is not a reranker", and a +reranker server answers chat, completion, Responses, Messages, speech, decision, System One and embedding requests +with HTTP 400 "this model is a reranker: use POST /v1/rerank". | Field | Meaning | | --- | --- | diff --git a/zig/src/families/qwen3/cuda/weights.zig b/zig/src/families/qwen3/cuda/weights.zig index b84030b9fd8..5793445d250 100644 --- a/zig/src/families/qwen3/cuda/weights.zig +++ b/zig/src/families/qwen3/cuda/weights.zig @@ -1,4 +1,4 @@ -//! A Qwen3 dense checkpoint on the device: BF16 tensors as stored, norms widened to fp32, the output head on the host. +//! A Qwen3 dense checkpoint: BF16 tensors on the device as stored, norms widened to fp32, the head's stored name. const std = @import("std"); const core = @import("core"); const cuda = @import("cuda"); From 39946c997445d838bfe0e4649a558fe4fdafe973 Mon Sep 17 00:00:00 2001 From: gilby Date: Sun, 1 Nov 2026 16:36:29 +0100 Subject: [PATCH 015/186] hip flashnext: the MTP head absorbs a prompt chunk on the prompt tiles After each prompt chunk the draft head runs over the chunk's rows to fill its own caches. It ran as a decode window, so its projections and its MoE took the stream tiles for up to a whole chunk of rows, re-reading each weight once per four rows. A prompt chunk now runs it on the prompt tiles (o.prefill and moe.run's prompt path), as the main layers' prompt rows already do; decode windows keep the window path. The head's prompt-row caches take the prompt tiles' sums, so its drafts can differ in the last bits; replies are verified as before and do not change. --- CHANGELOG.md | 1 + zig/src/families/flashnext/hip/head.zig | 8 +++++--- zig/src/families/flashnext/hip/lanes.zig | 3 +++ 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3dff3fd2036..a673c3316f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -637,6 +637,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: 15 more lean findings fixed in the CLI, cluster and CUDA comments, with no change in behaviour (#2320). - Internal: 15 more lean findings fixed in the Metal kernels' comments, with no change in behaviour (#2318). +- HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 3% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's drafts are not bit-identical to before (#1878). - Kept prompt states: a state saved to disk is reused only by the same build and the same checkpoint files, so a rewritten checkpoint or a new build starts cold instead of reusing an old state; the first long prompt after an upgrade prefills again (#2163, #2131). diff --git a/zig/src/families/flashnext/hip/head.zig b/zig/src/families/flashnext/hip/head.zig index b7308a13c6f..ca0f0ecd763 100644 --- a/zig/src/families/flashnext/hip/head.zig +++ b/zig/src/families/flashnext/hip/head.zig @@ -24,9 +24,10 @@ pub fn caches( return cache.Caches.init(gpa, d, h.dims, &.{false}, cap, act, eos); } -/// `tokens` (each the token after its row) over residual `streams` from `pos`; `ones` a one a stream a row. +/// `tokens` (each the token after its row) over `streams` from `pos`; `prompt`: a prompt chunk's rows on prompt tiles. pub fn run( ops: Ops, + prompt: bool, m: *const model.Model, host: *forward.Host, c: *cache.Caches, @@ -36,7 +37,8 @@ pub fn run( ones: hip.DeviceBuffer, ) !Out { var o = ops; - o.window = true; + o.window = !prompt; + o.prefill = prompt; const w = &m.mtp.?; const d = w.dims; const rows = tokens.len; @@ -82,7 +84,7 @@ pub fn run( const branch = try mixers.sparse(o, m, d, l.mixer.sparse, &seg, 0, mixed, rows, dev_pos, host, null); try forward.hcNorm(o, m, h, branch, inject, x, l.mhc.scale, rows); const mixed2 = try forward.hcPre(o, m, l.mhc, x, rows, inject); - const y = try hip.moe.run(o, m.act, l.moe, mixed2, rows, false); + const y = try hip.moe.run(o, m.act, l.moe, mixed2, rows, prompt); try forward.hcNorm(o, m, h, y, inject, x, w.mix.scale, rows); c.pooled = (pos + rows) / d.index_ratio; return .{ .mixed = try forward.hcPre(o, m, w.mix, x, rows, null), .streams = h }; diff --git a/zig/src/families/flashnext/hip/lanes.zig b/zig/src/families/flashnext/hip/lanes.zig index 2829f99497e..1a807fe3ef3 100644 --- a/zig/src/families/flashnext/hip/lanes.zig +++ b/zig/src/families/flashnext/hip/lanes.zig @@ -357,6 +357,7 @@ pub const Fx = struct { const n = if (last) rows - 1 else rows; if (n > 0) _ = try head_mod.run( e.ops(&e.arena), + true, m, &e.host, hc, @@ -540,6 +541,7 @@ pub const Fx = struct { const o = e.ops(&e.arena); var out = try head_mod.run( o, + false, m, &e.host, hc, @@ -564,6 +566,7 @@ pub const Fx = struct { if (d + 1 == depth or prob[0] < stop_under) break; out = try head_mod.run( o, + false, m, &e.host, hc, From d1fecd771dfcac15272bc8e751462e5be7a5572f Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 12 Nov 2026 08:38:10 +0100 Subject: [PATCH 016/186] flashnext hip: the head's run takes the prefill flag as an argument, not from the options --- zig/src/families/flashnext/hip/engine.zig | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/zig/src/families/flashnext/hip/engine.zig b/zig/src/families/flashnext/hip/engine.zig index 5196a05937f..6f597278648 100644 --- a/zig/src/families/flashnext/hip/engine.zig +++ b/zig/src/families/flashnext/hip/engine.zig @@ -304,9 +304,7 @@ pub const Engine = struct { defer hc.deinit(e.gpa); var ones = try hip.DeviceBuffer.alloc(&e.driver, rows * e.cfg.hc_count * e.act.size()); defer ones.free(); - var ho = o; - ho.prefill = true; - _ = try head_mod.run(ho, &e.weights.model, &e.host, &hc, out.streams, tokens, pos, ones); + _ = try head_mod.run(o, true, &e.weights.model, &e.host, &hc, out.streams, tokens, pos, ones); } _ = try e.logits(o, out.mixed, verify_rows); // the peak, not what is left at the end: each layer hands its scratch back when it is done From c2ee9273a55f981d06ce35a7bfc6e601db90b7ba Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 12 Nov 2026 08:38:10 +0100 Subject: [PATCH 017/186] changelog: the MTP head's prompt-tile line moves to the top of Unreleased --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a673c3316f9..df461708d7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,6 +67,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). +- HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. @@ -637,7 +638,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: 15 more lean findings fixed in the CLI, cluster and CUDA comments, with no change in behaviour (#2320). - Internal: 15 more lean findings fixed in the Metal kernels' comments, with no change in behaviour (#2318). -- HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 3% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's drafts are not bit-identical to before (#1878). - Kept prompt states: a state saved to disk is reused only by the same build and the same checkpoint files, so a rewritten checkpoint or a new build starts cold instead of reusing an old state; the first long prompt after an upgrade prefills again (#2163, #2131). From 52ee8059489aba327eda8c292eac80f50c74ffa2 Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 12 Nov 2026 15:34:35 +0100 Subject: [PATCH 018/186] changelog: the MTP head's prompt-tile line carries the cold first-token gain of its receipt --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index df461708d7a..e4c3df20d79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -68,6 +68,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. +- HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. From f9df8bfe7dc8ccb5b7431d25de3f456f242d8afc Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:16:19 +0100 Subject: [PATCH 019/186] dsv41 metal: the lane target, the prompt cache and the native server over Metal target.zig is the lanes' verify target: a slot's prompt, its windows of rows for several slots in one launch, keeps and rewinds, and kept prompt states (the prompt cache) whose resumed rows the reply reports as cached_tokens. boot.zig loads a release checkpoint folder and entry.zig serves it as tensorfold-dsv41 (dsv41-metal-serve). Tests serve the tiny model through LaneHost: concurrent equals solo equals serial, drafted equals draft:false, cache hits equal fresh runs, seeds and cancels. --- zig/build/dsv41_metal.zig | 32 +- zig/src/families/deepseek_v41/draft/iface.zig | 7 + zig/src/families/deepseek_v41/draft/lanes.zig | 1 + zig/src/families/deepseek_v41/metal/boot.zig | 144 ++++++ zig/src/families/deepseek_v41/metal/entry.zig | 44 ++ zig/src/families/deepseek_v41/metal/metal.zig | 8 +- zig/src/families/deepseek_v41/metal/pass.zig | 57 +++ zig/src/families/deepseek_v41/metal/serve.zig | 121 +++++ .../families/deepseek_v41/metal/snapshots.zig | 154 +++++++ .../families/deepseek_v41/metal/target.zig | 418 ++++++++++++++++++ .../families/deepseek_v41/phases_timed.zig | 6 +- zig/tests/dsv41_metal/boot_test.zig | 171 +++++++ zig/tests/dsv41_metal/fixture_serve.zig | 69 +++ zig/tests/dsv41_metal/layer_test.zig | 1 + zig/tests/dsv41_metal/multi_test.zig | 219 +++++++++ zig/tests/dsv41_metal/serve_test.zig | 267 +++++++++++ 16 files changed, 1716 insertions(+), 3 deletions(-) create mode 100644 zig/src/families/deepseek_v41/metal/boot.zig create mode 100644 zig/src/families/deepseek_v41/metal/entry.zig create mode 100644 zig/src/families/deepseek_v41/metal/pass.zig create mode 100644 zig/src/families/deepseek_v41/metal/serve.zig create mode 100644 zig/src/families/deepseek_v41/metal/snapshots.zig create mode 100644 zig/src/families/deepseek_v41/metal/target.zig create mode 100644 zig/tests/dsv41_metal/boot_test.zig create mode 100644 zig/tests/dsv41_metal/fixture_serve.zig create mode 100644 zig/tests/dsv41_metal/multi_test.zig create mode 100644 zig/tests/dsv41_metal/serve_test.zig diff --git a/zig/build/dsv41_metal.zig b/zig/build/dsv41_metal.zig index 1ee5d6bd149..3eb2207ab1a 100644 --- a/zig/build/dsv41_metal.zig +++ b/zig/build/dsv41_metal.zig @@ -4,7 +4,7 @@ const std = @import("std"); /// The package over the family's neutral host module; `metal` and `sources` are the engine's own modules. pub fn package(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, optimize: std.builtin.OptimizeMode) struct { pkg: *std.Build.Module, host: *std.Build.Module } { const host = @import("dsv41_shared.zig").moduleAt(b, b.graph.host, optimize, "neutral.zig"); - // the lanes and the engine API the backend's modules import + // the lanes, the engine API and the served engine (serve_lanes) the backend's target and pass plug into @import("dsv41_draft.zig").addImportsAt(b, host, b.graph.host, optimize); const pkg = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/metal/metal.zig"), @@ -16,6 +16,20 @@ pub fn package(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Modu return .{ .pkg = pkg, .host = host }; } +/// The native server's engine over the package: metal/entry.zig, the family's own serving modules beside it. +pub fn engine(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const p = package(b, metal, sources, optimize); + const m = b.createModule(.{ + .root_source_file = b.path("zig/src/families/deepseek_v41/metal/entry.zig"), + .target = b.graph.host, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "dsv41", .module = p.host }, .{ .name = "metal", .module = metal } }, + }); + for ([_][]const u8{ "engine_api", "core", "tokenizer", "dsv41_serve", "lanes", "json" }) |name| m.addImport(name, p.host.import_table.get(name).?); + return m; +} + /// `test-dsv41-metal`: package, layer and forward checks on this GPU; `test-dsv41-metal-real`: a TF_DSV41_REAL layer. pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module, test_step: *std.Build.Step) void { const p = package(b, metal, sources, .Debug); @@ -32,6 +46,19 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module const run = b.addRunArtifact(b.addTest(.{ .root_module = tests })); run.has_side_effects = true; // a device run: never cached step.dependOn(&run.step); + // the served engine over the tiny model: LaneHost replies, drafts, the prompt cache, seeds, cancels + const serve = b.createModule(.{ + .root_source_file = b.path("zig/tests/dsv41_metal/serve_test.zig"), + .target = b.graph.host, + .optimize = .Debug, + .link_libc = true, + .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "metal", .module = metal }, .{ .name = "dsv41", .module = p.host } }, + }); + serve.addImport("dsv41_serve_fixtures", p.host.import_table.get("dsv41_serve").?.import_table.get("dsv41_serve_fixtures").?); + const serve_run = b.addRunArtifact(b.addTest(.{ .root_module = serve })); + serve_run.has_side_effects = true; + step.dependOn(&serve_run.step); + b.step("test-dsv41-metal-serve", "DeepSeek-V4.1 Metal served through LaneHost: concurrent, drafted, cached, seeded, cancelled").dependOn(&serve_run.step); test_step.dependOn(pkg_run); // the package's host checks only; the device run stays in its own step const fast = package(b, metal, sources, .ReleaseSafe); const real = b.createModule(.{ @@ -66,4 +93,7 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module const ptime_run = b.addRunArtifact(b.addTest(.{ .root_module = ptime })); ptime_run.has_side_effects = true; b.step("test-dsv41-metal-prompt-time", "DeepSeek-V4.1 Metal: ms a stage of one release layer's prompt chunks (TF_DSV41_REAL)").dependOn(&ptime_run.step); + // tensorfold-dsv41 over Metal: the native server with this family's lane engine (the CUDA pair's macOS twin) + const exe = @import("dsv41_draft.zig").serverOn(b, b.graph.host, .ReleaseSafe, engine(b, metal, sources, .ReleaseFast)); + b.step("dsv41-metal-serve", "Install tensorfold-dsv41 over Metal (DeepSeek-V4.1 release checkpoints on this Mac)").dependOn(&b.addInstallArtifact(exe, .{}).step); } diff --git a/zig/src/families/deepseek_v41/draft/iface.zig b/zig/src/families/deepseek_v41/draft/iface.zig index 39aa01f4717..f8bfae834e6 100644 --- a/zig/src/families/deepseek_v41/draft/iface.zig +++ b/zig/src/families/deepseek_v41/draft/iface.zig @@ -46,8 +46,15 @@ pub const Target = struct { replays: ?*const fn (ptr: *anyopaque, list: []const Final) anyerror!void = null, /// Engram reads ahead of the next window (TF_DSV41_ENGRAM_WARM); reads only, window bits are unaffected. warm: ?*const fn (ptr: *anyopaque, asks: []const Warm) void = null, + /// After `prefill`: the prompt rows the slot took from kept state (its first row computed is this one). + resumed: ?*const fn (ptr: *anyopaque, slot: u32) u64 = null, }; + /// The prompt rows the slot's last `prefill` took from kept state (0: none, or a target that keeps none). + pub fn resumedRows(t: Target, slot: u32) u64 { + return if (t.vtable.resumed) |f| f(t.ptr, slot) else 0; + } + /// A slot's whole prompt whose rows but the last are in (its `finish` / `tail` follows). pub const Final = struct { slot: u32, ids: []const u32 }; diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 2f875c9715e..dec67d5e17a 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -188,6 +188,7 @@ pub const Lanes = struct { const ids = s.prompt(); x.target.admit(slot, s.max_new); st.first = try x.target.prefill(slot, ids, s.sampling, ids.len); + s.cached = @intCast(x.target.resumedRows(slot)); st.len = ids.len; _ = x.scratch.reset(.retain_capacity); // the taps hold the prompt's last rows (after a verify-tail prompt only that row), at their positions diff --git a/zig/src/families/deepseek_v41/metal/boot.zig b/zig/src/families/deepseek_v41/metal/boot.zig new file mode 100644 index 00000000000..157552c6d26 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/boot.zig @@ -0,0 +1,144 @@ +//! A release checkpoint folder on this Mac: config, kernels, mapped weights, DSpark, Engram rows; the served model. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const pipelines = @import("pipelines.zig"); +const model_weights = @import("model_weights.zig"); +const dspark = @import("dspark.zig"); +const rope = @import("rope.zig"); +const scratch = @import("scratch.zig"); +const forward = @import("forward.zig"); +const serve = @import("serve.zig"); +const target_ = @import("target.zig"); + +const Config = host.config.Config; +const eh = host.engram_host; +const er = host.engram_rows; +const Allocator = std.mem.Allocator; + +/// The files a boot reads beside the shards: the Engram token map (here or under TF_DSV41_ASSETS). +pub const token_map = "engram-token-map.bin"; + +pub const Options = struct { + slots: u32 = 4, + /// positions a slot holds (null: 262,144, at most the checkpoint's) + context: ?u64 = null, + drafts: bool = true, + cache_bytes: u64 = 4 << 30, + /// a checkpoint of the release's format at other shapes (TF_DSV41_SMALL=1: the smoke test's tiny model) + small: bool = false, + /// Engram row shards (null: TF_DSV41_ENGRAM_DIR, else the config's engram_table_dir beside the checkpoint) + engram_dir: ?[]const u8 = null, +}; + +/// Everything a served model borrows, owned here and closed after the model's target and pass. +pub const Boot = struct { + a: Allocator, + c: Config, + device: mtl.Device, + queue: mtl.Queue, + p: pipelines.Pipelines, + w: *model_weights.Model, + ds: ?dspark.Weights = null, + tables: rope.Tables, + s: scratch.Scratch, + map_bytes: []align(4) u8 = &.{}, + lbuf: [eh.max_layers]u32 = undefined, + eng: eh.Host = undefined, + rows: ?*er.Rows = null, + + /// The checkpoint in `dir` loaded and served by a `serve.Metal` (the caller's `serve_lanes.Served` closes it). + pub fn open(a: Allocator, io: std.Io, dir: []const u8, o: Options) !*serve.Metal { + const b = try a.create(Boot); + errdefer a.destroy(b); + b.a = a; + b.ds = null; + b.map_bytes = &.{}; + b.rows = null; + b.c = try readConfig(a, io, dir, o.small); + const c = &b.c; + const limit: u32 = @intCast(@min(o.context orelse 262144, c.rope.max_positions)); + b.device = try mtl.Device.init(); + errdefer b.device.deinit(); + b.queue = try b.device.queue(); + errdefer b.queue.deinit(); + b.p = try pipelines.Pipelines.load(a, b.device, c); + errdefer b.p.deinit(); + var pack = try host.pack.Pack.open(a, io, dir); + defer pack.deinit(); + b.w = try model_weights.Model.load(a, io, b.device, b.queue, &pack, c, null); + errdefer b.w.deinit(); + if (o.drafts and c.mtp_layers > 0) b.ds = try dspark.loadWeights(a, io, b.device, &pack, c); + errdefer if (b.ds) |*w| w.deinit(); + b.tables = try rope.Tables.init(b.device, c, limit); + errdefer b.tables.deinit(); + b.s = try scratch.Scratch.init(b.device, c, limit); + errdefer b.s.deinit(); + const lookup = try b.engram(io, dir, o); + errdefer b.closeEngram(); + const r = b.w.report(); + std.log.info("dsv41 metal: {d} layers, {d} shards ({d} MiB mapped, {d} MiB copied), {d} slots of {d} positions, DSpark {s}", .{ c.layers, r.shards, r.shard_bytes >> 20, r.arena_bytes >> 20, o.slots, limit, if (b.ds != null) "on" else "off" }); + const env: target_.Env = .{ .device = b.device, .queue = b.queue, .c = c, .p = &b.p, .tables = &b.tables, .s = &b.s, .model = b.w.view(), .lookup = lookup }; + return serve.Metal.open(a, .{ .env = env, .top = &b.w.top, .dspark = if (b.ds) |*w| w else null, .owner = .{ .ctx = b, .close = close } }, .{ .slots = o.slots, .limit = limit, .drafts = o.drafts, .cache_bytes = o.cache_bytes }); + } + + fn readConfig(a: Allocator, io: std.Io, dir: []const u8, small: bool) !Config { + if (!small) return Config.read(a, io, dir); + const path = try std.fs.path.join(a, &.{ dir, "config.json" }); + defer a.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 24)); + defer a.free(bytes); + return Config.parseShapes(a, bytes); + } + + /// The Engram hashes (the canonical token map) and the packed row shards; null without Engram layers. + fn engram(b: *Boot, io: std.Io, dir: []const u8, o: Options) !?forward.Lookup { + const c = &b.c; + if (c.engram_layers.len == 0) return null; + const assets = if (std.c.getenv("TF_DSV41_ASSETS")) |v| std.mem.span(v) else dir; + const path = try std.fs.path.join(b.a, &.{ assets, token_map }); + defer b.a.free(path); + b.map_bytes = std.Io.Dir.cwd().readFileAllocOptions(io, path, b.a, .limited(1 << 26), .@"4", null) catch |e| { + std.log.err("dsv41 metal: the Engram token map {s} ({t}): put it beside the checkpoint or under TF_DSV41_ASSETS", .{ path, e }); + return error.EngramTokenMap; + }; + errdefer b.a.free(b.map_bytes); + b.eng = try eh.Host.load(b.map_bytes, c.engram_heads, c.engram_pad, c.engram_vocab, &b.lbuf); + const rows_dir = try b.rowsDir(dir, o); + defer b.a.free(rows_dir); + const cap: u32 = if (std.c.getenv("TF_DSV41_ENGRAM_CACHE")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else 65536; + b.rows = try er.Rows.open(b.a, rows_dir, 0, 1, b.eng.layer_ids, cap); + std.log.info("dsv41 metal: Engram rows from {s} ({d} records a layer cached)", .{ rows_dir, cap }); + return .{ .a = b.a, .host = &b.eng, .source = b.rows.?.source() }; + } + + /// The row shards' folder: the option, TF_DSV41_ENGRAM_DIR, else engram_table_dir (relative to the checkpoint). + fn rowsDir(b: *Boot, dir: []const u8, o: Options) ![]u8 { + if (o.engram_dir) |d| return b.a.dupe(u8, d); + if (std.c.getenv("TF_DSV41_ENGRAM_DIR")) |v| return b.a.dupe(u8, std.mem.span(v)); + const t = b.c.engram_table_dir.items(); + if (t.len == 0) return error.EngramRowsUnset; + if (std.fs.path.isAbsolute(t)) return b.a.dupe(u8, t); + return std.fs.path.join(b.a, &.{ dir, t }); + } + + fn closeEngram(b: *Boot) void { + if (b.rows) |r| r.close(); + b.rows = null; + if (b.map_bytes.len > 0) b.a.free(b.map_bytes); + b.map_bytes = &.{}; + } + + fn close(ctx: *anyopaque) void { + const b: *Boot = @ptrCast(@alignCast(ctx)); + b.closeEngram(); + b.s.deinit(); + b.tables.deinit(); + if (b.ds) |*w| w.deinit(); + b.w.deinit(); + b.p.deinit(); + b.queue.deinit(); + b.device.deinit(); + b.a.destroy(b); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/entry.zig b/zig/src/families/deepseek_v41/metal/entry.zig new file mode 100644 index 00000000000..7eb864254b5 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/entry.zig @@ -0,0 +1,44 @@ +//! DeepSeek-V4.1's Metal entry point for the native server: a release checkpoint served through the lane core. +const std = @import("std"); +const api = @import("engine_api"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); + +const serve_lanes = host.serve_lanes; +const Allocator = std.mem.Allocator; + +pub const backends: []const []const u8 = &.{"metal"}; +pub const families: []const api.Family = &.{.{ .model_type = "deepseek_v41", .formats = &.{"fp8"}, .raw_text = false }}; + +pub fn chip(_: Allocator) ?[]const u8 { + return null; +} + +/// The engine for `o.dir`, or null with `problem` set (the native_engines contract). +pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (@import("sys").experimental.refusal(o.model_type, @import("sys").getenv("TF_EXPERIMENTAL")) != null) { + problem.* = "DeepSeek-V4.1 on Metal is experimental: its served lanes have no whole-model gate yet; set TF_EXPERIMENTAL=deepseek to opt in"; + return null; + } + if (!std.mem.eql(u8, o.model_type, "deepseek_v41")) { + problem.* = try std.fmt.allocPrint(a, "this build serves deepseek_v41 checkpoints only, not {s}", .{o.model_type}); + return null; + } + const opts: pkg.boot.Options = .{ + .slots = @max(1, @min(o.lanes, 16)), + .context = if (o.context) |c| @intCast(@max(c, 1)) else null, + .drafts = o.drafts, + .cache_bytes = if (o.prompt_cache_gib) |g| @intFromFloat(@max(g, 0) * (1 << 30)) else 4 << 30, + .small = if (std.c.getenv("TF_DSV41_SMALL")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false, + }; + const m = pkg.boot.Boot.open(gpa, io, o.dir, opts) catch |e| { + problem.* = try std.fmt.allocPrint(a, "the DeepSeek-V4.1 Metal engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); + return null; + }; + const s = serve_lanes.Served.open(gpa, io, m.model(), o.dir, o.context) catch |e| { + m.model().close(m); + problem.* = try std.fmt.allocPrint(a, "the DeepSeek-V4.1 Metal engine cannot start its lanes ({s})", .{@errorName(e)}); + return null; + }; + return .{ .engine = s.engine(), .close = serve_lanes.Served.close, .ctx = s, .halt = serve_lanes.Served.halt }; +} diff --git a/zig/src/families/deepseek_v41/metal/metal.zig b/zig/src/families/deepseek_v41/metal/metal.zig index 55f4e533334..bec87c2930a 100644 --- a/zig/src/families/deepseek_v41/metal/metal.zig +++ b/zig/src/families/deepseek_v41/metal/metal.zig @@ -1,4 +1,4 @@ -//! DeepSeek V4.1 on Metal, the backend package: the decode forward (1-16 rows), the prompt path and DSpark. +//! DeepSeek V4.1 on Metal, the backend package: the decode forward (1-16 rows), the prompt path, the served lanes. const std = @import("std"); pub const host = @import("dsv41"); @@ -25,6 +25,11 @@ pub const dspark_plan = @import("dspark_plan.zig"); pub const dspark_stage = @import("dspark_stage.zig"); pub const dspark_weights = @import("dspark_weights.zig"); pub const engram = pipelines.engram; +pub const snapshots = @import("snapshots.zig"); +pub const target = @import("target.zig"); +pub const pass = @import("pass.zig"); +pub const serve = @import("serve.zig"); +pub const boot = @import("boot.zig"); test { std.testing.refAllDecls(@This()); @@ -34,4 +39,5 @@ test { _ = prompt_scratch; _ = dspark; _ = dspark_plan; + _ = pass; } diff --git a/zig/src/families/deepseek_v41/metal/pass.zig b/zig/src/families/deepseek_v41/metal/pass.zig new file mode 100644 index 00000000000..520be10a3e8 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/pass.zig @@ -0,0 +1,57 @@ +//! DSpark's pass for the lanes (iface.Pass) over the Metal drafter: the target's taps in, blocks of drafts out. +const std = @import("std"); +const host = @import("dsv41"); +const dspark = @import("dspark.zig"); +const target_ = @import("target.zig"); + +const iface = host.iface; + +pub const Pass = struct { + x: *dspark.Pass, + /// the target whose taps the ingests name (by address) + t: *const target_.Target, + passes: u64 = 0, + ingests: u64 = 0, + + pub fn pass(p: *Pass) iface.Pass { + return .{ .ptr = p, .vtable = &.{ .ingest = ingest, .propose = propose, .reset = reset, .markov = markov } }; + } + + fn self(ptr: *anyopaque) *Pass { + return @ptrCast(@alignCast(ptr)); + } + + fn ingest(ptr: *anyopaque, slot: u32, start: u64, taps: iface.Taps, rows: []const u32) anyerror!void { + const p = self(ptr); + if (rows.len == 0) return; + const at = p.t.bufferOf(taps.address) orelse return error.UnknownTaps; + p.ingests += 1; + try p.x.ingest(slot, start, .{ .buf = at.buf, .off = at.off, .rows = taps.rows }, rows); + } + + fn propose(ptr: *anyopaque, asks: []const iface.Ask, out: []iface.Proposal) anyerror!void { + const p = self(ptr); + var xs: [16]dspark.Ask = undefined; + var ys: [16]dspark.Proposal = undefined; + if (asks.len > xs.len or out.len != asks.len) return error.TooManySlots; + for (asks, out, xs[0..asks.len], ys[0..asks.len]) |a, o, *x, *y| { + x.* = .{ .slot = a.slot, .anchor = a.anchor, .start = a.start, .params = a.params }; + y.* = .{ .drafts = o.drafts, .conf = o.conf, .cand = o.cand, .base = o.base }; + } + p.passes += 1; + try p.x.propose(xs[0..asks.len], ys[0..asks.len]); + } + + fn reset(ptr: *anyopaque, slot: u32) void { + self(ptr).x.reset(slot); + } + + fn markov(ptr: *anyopaque) ?host.dspark.Markov { + return self(ptr).x.markov(); + } +}; + +test "the pass's asks and proposals are the drafter's, field for field" { + try std.testing.expectEqual(@sizeOf(iface.Ask), @sizeOf(dspark.Ask)); + try std.testing.expectEqual(@sizeOf(iface.Proposal), @sizeOf(dspark.Proposal)); +} diff --git a/zig/src/families/deepseek_v41/metal/serve.zig b/zig/src/families/deepseek_v41/metal/serve.zig new file mode 100644 index 00000000000..b0d03fc2f11 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/serve.zig @@ -0,0 +1,121 @@ +//! DeepSeek-V4.1 on Metal as the lanes serve it (serve_lanes.Model): the target, DSpark's pass, costs and hooks. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const weights = @import("weights.zig"); +const dspark = @import("dspark.zig"); +const prompt = @import("prompt.zig"); +const target_ = @import("target.zig"); +const pass_ = @import("pass.zig"); +const pipelines = @import("pipelines.zig"); + +const serve_lanes = host.serve_lanes; +const iface = host.iface; +const api = host.api; +const Allocator = std.mem.Allocator; + +/// The parts a served model borrows; `owner` (boot.zig's model, or a test's rig) is closed after them. +pub const Parts = struct { + env: target_.Env, + top: *const weights.Top, + /// the drafter's weights (null: serial windows, no drafts) + dspark: ?*const dspark.Weights, + owner: ?Owner = null, +}; + +pub const Owner = struct { ctx: *anyopaque, close: *const fn (ctx: *anyopaque) void }; + +pub const Options = struct { + slots: u32 = 4, + limit: u32, + /// DSpark's pass when the checkpoint has one (false: serial windows through the same lanes) + drafts: bool = true, + cache_bytes: u64 = 0, + prompt: prompt.Opts = .{}, +}; + +/// A position's draft candidates the pass keeps (the shape the lanes size their buffers by). +const candidates = 64; + +pub const Metal = struct { + a: Allocator, + parts: Parts, + opt: Options, + t: *target_.Target, + dp: ?*dspark.Pass = null, + ad: pass_.Pass = undefined, + + pub fn open(a: Allocator, parts: Parts, opt: Options) !*Metal { + const m = try a.create(Metal); + errdefer a.destroy(m); + m.* = .{ .a = a, .parts = parts, .opt = opt, .t = undefined }; + m.t = try target_.Target.init(a, parts.env, .{ .slots = opt.slots, .limit = opt.limit, .prompt = opt.prompt, .cache_bytes = opt.cache_bytes }); + errdefer m.t.deinit(); + const e = parts.env; + if (opt.drafts) if (parts.dspark) |w| { + m.dp = try dspark.Pass.init(a, e.device, e.queue, e.c, e.p, e.s, e.tables, parts.top, w, .{ .slots = opt.slots, .candidates = @min(candidates, e.c.vocab) }); + m.ad = .{ .x = m.dp.?, .t = m.t }; + }; + return m; + } + + /// The lanes' view: one slot a stream, rows of 16 a forward, the release's DSpark shape on one device. + pub fn model(m: *Metal) serve_lanes.Model { + const c = m.parts.env.c; + const shape: host.dspark.Shape = .{ .block = c.dspark_block, .noise = c.dspark_noise_token, .window = c.window, .rank = c.dspark_markov_rank, .hidden = c.hidden, .candidates = @min(candidates, c.vocab) }; + return .{ + .ptr = m, + .name = "dsv41-metal", + .slots = m.opt.slots, + .limit = m.opt.limit, + .world = 1, + .prefill_step = m.opt.prompt.chunk, + .forward_rows = pipelines.max_rows, + .shape = shape, + .target = targetFn, + .pass = passFn, + .calibrate = calibrate, + .memory = .{ .ctx = m, .read = memory }, + .health = .{ .ctx = m, .read = healthy }, + .close = close, + // the default cost table: TF_DSV41_CALIB has nothing to time on Metal + .calib = .{ .mode = "default" }, + }; + } + + fn self(ctx: *anyopaque) *Metal { + return @ptrCast(@alignCast(ctx)); + } + + fn targetFn(ctx: *anyopaque) iface.Target { + return self(ctx).t.target(); + } + + fn passFn(ctx: *anyopaque) ?iface.Pass { + const m = self(ctx); + return if (m.dp != null) m.ad.pass() else null; + } + + fn calibrate(_: *anyopaque, _: Allocator, _: std.Io, _: []const u8, _: host.calib_env.Knobs, _: host.calib_env.Shape) anyerror!host.calib_lanes.Boot { + return error.Unmeasured; // the mode is "default": never asked to measure + } + + /// The device's allocated bytes now (the weights, slots, prompt buffers and kept-state host copies aside). + fn memory(ctx: ?*anyopaque, _: bool) ?api.Memory { + const m = self(ctx.?); + const used = m.parts.env.device.allocated(); + return .{ .active = used, .cache = m.t.cache.held, .peak = used }; + } + + fn healthy(ctx: *anyopaque) bool { + return !self(ctx).t.failed; + } + + fn close(ctx: *anyopaque) void { + const m = self(ctx); + if (m.dp) |p| p.deinit(); + m.t.deinit(); + if (m.parts.owner) |o| o.close(o.ctx); + m.a.destroy(m); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/snapshots.zig b/zig/src/families/deepseek_v41/metal/snapshots.zig new file mode 100644 index 00000000000..dc14dd0b0d4 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/snapshots.zig @@ -0,0 +1,154 @@ +//! Kept prompt states on the host: a slot's regions up to its rows, its ids and DSpark taps, restored into any slot. +const std = @import("std"); +const mtl = @import("metal"); +const Config = @import("dsv41").config.Config; +const kv = @import("kv.zig"); + +/// One kept state: the slot after `ids` (every row run), and the choice's logits when it ended a prompt. +pub const Entry = struct { + ids: []u32, + /// each region's window, pool rows, index key rows and carry, in layer order + bytes: []u8, + /// the taps ring: position p's row at p % ring rows + taps: []u16, + logits: ?[]f32, + tick: u64 = 0, + + fn size(e: *const Entry) u64 { + return e.ids.len * 4 + e.bytes.len + e.taps.len * 2 + if (e.logits) |l| l.len * 4 else 0; + } + + fn free(e: *Entry, a: std.mem.Allocator) void { + a.free(e.ids); + a.free(e.bytes); + a.free(e.taps); + if (e.logits) |l| a.free(l); + a.destroy(e); + } +}; + +/// The pool and key rows a state of `len` positions has written at a KV source (and the row it may be filling). +pub fn poolRowsOf(c: *const Config, slot: *const kv.Slot, layer: u32, len: u64) usize { + const all = kv.poolRows(c, layer, slot.limit); + return @min(all, @as(usize, @intCast(len / c.compressRatio(layer))) + 1); +} + +/// Every region's bytes a state of `len` positions reads, as (buffer, bytes) runs in order. +fn runs(c: *const Config, slot: *const kv.Slot, len: u64, out: *[4 * 64]Run) []Run { + var n: usize = 0; + for (slot.regions, 0..) |r, l| { + out[n] = .{ .buf = r.window, .len = r.window.length() }; + n += 1; + if (r.pool) |p| { + const rows = poolRowsOf(c, slot, @intCast(l), len); + out[n] = .{ .buf = p, .len = rows * c.head_dim * 2 }; + out[n + 1] = .{ .buf = r.keys.?, .len = rows * kv.index_key_bytes }; + n += 2; + } + if (r.carry) |k| { + out[n] = .{ .buf = k, .len = k.length() }; + n += 1; + } + } + return out[0..n]; +} + +const Run = struct { buf: mtl.Buffer, len: usize }; + +/// Zeroes what a state of up to `high` positions may have written: a reused slot reads as a new one. +pub fn clear(c: *const Config, slot: *const kv.Slot, high: u64) void { + var buf: [4 * 64]Run = undefined; + for (runs(c, slot, high, &buf)) |r| @memset(r.buf.contents()[0..r.len], 0); + for (slot.regions) |*r| kv.resetCarry(r, c); +} + +pub const Store = struct { + a: std.mem.Allocator, + /// bytes the entries may hold (0: no prompt cache) + cap: u64, + held: u64 = 0, + tick: u64 = 0, + entries: std.ArrayList(*Entry) = .empty, + + pub fn deinit(st: *Store) void { + for (st.entries.items) |e| e.free(st.a); + st.entries.deinit(st.a); + } + + pub fn on(st: *const Store) bool { + return st.cap > 0; + } + + /// The longest entry that starts `ids`: shorter than them, or as long with the choice's logits. + pub fn best(st: *Store, ids: []const u32) ?*Entry { + var out: ?*Entry = null; + for (st.entries.items) |e| { + if (e.ids.len > ids.len or !std.mem.eql(u32, e.ids, ids[0..e.ids.len])) continue; + if (e.ids.len == ids.len and e.logits == null) continue; + if (out == null or e.ids.len > out.?.ids.len) out = e; + } + if (out) |e| { + st.tick += 1; + e.tick = st.tick; + } + return out; + } + + /// The slot's state after `ids` kept (the oldest entries make room); one entry a distinct `ids`. + pub fn save(st: *Store, c: *const Config, slot: *const kv.Slot, ids: []const u32, ring: []const u16, logits: ?[]const f32) !void { + if (!st.on() or ids.len == 0) return; + st.tick += 1; + for (st.entries.items) |e| if (std.mem.eql(u32, e.ids, ids)) { + e.tick = st.tick; + if (e.logits != null or logits == null) return; + e.logits = try st.a.dupe(f32, logits.?); + st.held += logits.?.len * 4; + return; + }; + var buf: [4 * 64]Run = undefined; + const rs = runs(c, slot, ids.len, &buf); + var total: usize = 0; + for (rs) |r| total += r.len; + const size = ids.len * 4 + total + ring.len * 2 + if (logits) |l| l.len * 4 else 0; + if (size > st.cap) return; + while (st.held + size > st.cap) st.evict(); + const e = try st.a.create(Entry); + errdefer st.a.destroy(e); + e.* = .{ .ids = try st.a.dupe(u32, ids), .bytes = undefined, .taps = undefined, .logits = null, .tick = st.tick }; + errdefer st.a.free(e.ids); + e.bytes = try st.a.alloc(u8, total); + errdefer st.a.free(e.bytes); + var at: usize = 0; + for (rs) |r| { + @memcpy(e.bytes[at..][0..r.len], r.buf.contents()[0..r.len]); + at += r.len; + } + e.taps = try st.a.dupe(u16, ring); + errdefer st.a.free(e.taps); + if (logits) |l| e.logits = try st.a.dupe(f32, l); + try st.entries.append(st.a, e); + st.held += e.size(); + } + + /// The entry's state into `slot` (cleared first up to `high`, its old rows) and its taps into `ring`. + pub fn restore(_: *const Store, c: *const Config, e: *const Entry, slot: *const kv.Slot, high: u64, ring: []u16) void { + clear(c, slot, high); + var buf: [4 * 64]Run = undefined; + var at: usize = 0; + for (runs(c, slot, e.ids.len, &buf)) |r| { + @memcpy(r.buf.contents()[0..r.len], e.bytes[at..][0..r.len]); + at += r.len; + } + @memcpy(ring, e.taps); + } + + fn evict(st: *Store) void { + var old: usize = 0; + for (st.entries.items, 0..) |e, i| if (e.tick < st.entries.items[old].tick) { + old = i; + }; + const e = st.entries.swapRemove(old); + st.held -= e.size(); + e.free(st.a); + } +}; diff --git a/zig/src/families/deepseek_v41/metal/target.zig b/zig/src/families/deepseek_v41/metal/target.zig new file mode 100644 index 00000000000..2e83b64b236 --- /dev/null +++ b/zig/src/families/deepseek_v41/metal/target.zig @@ -0,0 +1,418 @@ +//! The lanes' verify target on Metal: a slot's prompt (prompt.zig), its windows of rows (forward.zig), keeps and reuse. +const std = @import("std"); +const mtl = @import("metal"); +const host = @import("dsv41"); +const kv = @import("kv.zig"); +const rope = @import("rope.zig"); +const ctx_ = @import("ctx.zig"); +const scratch = @import("scratch.zig"); +const forward = @import("forward.zig"); +const prompt = @import("prompt.zig"); +const pipelines = @import("pipelines.zig"); +const snapshots = @import("snapshots.zig"); + +const iface = host.iface; +const lanes = host.lanes; +const Config = host.config.Config; +const Allocator = std.mem.Allocator; +const R = pipelines.max_rows; + +/// What the target reads and borrows: the device, kernels, RoPE tables, layer scratch, weights and Engram rows. +pub const Env = struct { + device: mtl.Device, + queue: mtl.Queue, + c: *const Config, + p: *const pipelines.Pipelines, + tables: *const rope.Tables, + s: *const scratch.Scratch, + model: forward.Model, + lookup: ?forward.Lookup, +}; + +pub const Options = struct { + slots: u32, + /// positions a slot holds (the context limit) + limit: u32, + prompt: prompt.Opts = .{}, + /// bytes of kept prompt states (0: no prompt cache) + cache_bytes: u64 = 0, +}; + +/// What the target did, for the tests and the logs. +pub const Stats = struct { prompts: u64 = 0, hits: u64 = 0, cached_rows: u64 = 0, windows: u64 = 0, rows: u64 = 0, rewinds: u64 = 0, launches: u64 = 0 }; + +/// A slot: its device state, prompt tail, ids (rows run, the last window's included until its keep) and taps. +const Slot = struct { + kv: kv.Slot, + tail: prompt.Tail, + hist: std.ArrayList(u32) = .empty, + /// the last window: its first position and rows; `open` until its keep + base: u64 = 0, + rows: u32 = 0, + open: bool = false, + /// the taps of the last `ring rows` positions, position p at row p % ring rows + ring: []u16, + /// the taps `taps` names after a prompt: its last rows in order + ptaps: mtl.Buffer, + ptaps_rows: u32 = 0, + after_prompt: bool = false, + /// the most positions any request wrote here since the slot was last cleared + high: u64 = 0, + live: bool = false, + /// every launch since the request began completed (a release keeps only a whole state) + whole: bool = false, + /// the rows the last prompt took from kept state + resumed: u32 = 0, +}; + +pub const Target = struct { + a: Allocator, + env: Env, + opt: Options, + b: forward.Bufs, + pr: prompt.Prompt, + slots: []Slot, + cache: snapshots.Store, + stats: Stats = .{}, + /// a command buffer failed: the device state is no longer trusted (the health hook) + failed: bool = false, + + pub fn init(a: Allocator, env: Env, opt: Options) !*Target { + const c = env.c; + if (opt.slots == 0 or opt.limit == 0 or opt.limit > env.tables.positions) return error.BadOptions; + _ = try prompt.check(c); + const t = try a.create(Target); + errdefer a.destroy(t); + t.* = .{ .a = a, .env = env, .opt = opt, .b = undefined, .pr = undefined, .slots = &.{}, .cache = .{ .a = a, .cap = opt.cache_bytes } }; + t.b = try forward.Bufs.init(env.device, c); + errdefer t.b.deinit(); + t.pr = try prompt.Prompt.init(env.device, c, opt.prompt, opt.limit); + errdefer t.pr.deinit(); + const slots = try a.alloc(Slot, opt.slots); + errdefer a.free(slots); + var made: usize = 0; + errdefer for (slots[0..made]) |*s| freeSlot(a, s); + const cap = try prompt.tailRows(c, opt.prompt); + const ptb = @max(4, @as(usize, ringRows(c)) * tapWidth(c) * 2); + for (slots) |*s| { + var st = try kv.Slot.init(a, env.device, c, opt.limit); + errdefer st.deinit(); + var tail = try prompt.Tail.init(env.device, c, cap); + errdefer tail.deinit(); + const ptaps = try env.device.buffer(ptb, mtl.ResourceOptions.shared); + errdefer ptaps.deinit(); + const ring = try a.alloc(u16, @as(usize, ringRows(c)) * tapWidth(c)); + @memset(ring, 0); + s.* = .{ .kv = st, .tail = tail, .ring = ring, .ptaps = ptaps }; + made += 1; + } + t.slots = slots; + return t; + } + + fn freeSlot(a: Allocator, s: *Slot) void { + s.kv.deinit(); + s.tail.deinit(); + s.ptaps.deinit(); + a.free(s.ring); + s.hist.deinit(a); + } + + pub fn deinit(t: *Target) void { + for (t.slots) |*s| freeSlot(t.a, s); + t.a.free(t.slots); + t.cache.deinit(); + t.pr.deinit(); + t.b.deinit(); + t.a.destroy(t); + } + + /// The taps ring's rows: the drafter's context window. + fn ringRows(c: *const Config) u32 { + return c.window; + } + + /// Values in a taps row: every target layer's hidden. + fn tapWidth(c: *const Config) usize { + return c.dspark_targets.len * c.hidden; + } + + pub fn target(t: *Target) iface.Target { + return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .resumed = resumedFn } }; + } + + fn self(ptr: *anyopaque) *Target { + return @ptrCast(@alignCast(ptr)); + } + + /// The device buffer and offset of an address a `taps` named (the pass's ingest reads it there). + pub fn bufferOf(t: *const Target, address: u64) ?struct { buf: mtl.Buffer, off: usize } { + for (t.slots) |s| { + const bufs = [_]?mtl.Buffer{ s.ptaps, s.kv.taps }; + for (bufs) |ob| if (ob) |b| { + const lo = @intFromPtr(b.contents()); + if (address >= lo and address < lo + b.length()) return .{ .buf = b, .off = @intCast(address - lo) }; + }; + } + return null; + } + + // prompts ------------------------------------------------------------------------------------------------------ + + fn resumedFn(ptr: *anyopaque, slot: u32) u64 { + const t = self(ptr); + return if (slot < t.slots.len) t.slots[slot].resumed else 0; + } + + fn prefill(ptr: *anyopaque, slot: u32, ids: []const u32, smp: ?lanes.Sampling, draw: u64) anyerror!u32 { + const t = self(ptr); + if (slot >= t.slots.len) return error.NoSlot; + const s = &t.slots[slot]; + const c = t.env.c; + if (s.live) return error.SlotLive; + if (ids.len == 0) return error.EmptyPrompt; + if (ids.len >= t.opt.limit) return error.PastLimit; + for (ids) |id| if (id >= c.vocab) return error.BadToken; + s.whole = false; + s.open = false; + s.rows = 0; + t.stats.prompts += 1; + // a kept state that starts this prompt, else a cleared slot + const hit = if (t.cache.on()) t.cache.best(ids) else null; + const at: u32 = if (hit) |e| @intCast(e.ids.len) else 0; + if (hit) |e| { + t.cache.restore(c, e, &s.kv, s.high, s.ring); + t.stats.hits += 1; + t.stats.cached_rows += at; + } else snapshots.clear(c, &s.kv, s.high); + s.resumed = at; + s.tail.end = at; + s.tail.done = at; + s.hist.clearRetainingCapacity(); + try s.hist.appendSlice(t.a, ids); + s.high = @max(at, ids.len); + const n: u32 = @intCast(ids.len); + const V: usize = c.vocab; + var logits: []const f32 = undefined; + if (at == n) logits = hit.?.logits.? else { + t.pr.run(t.promptEnv(), &s.kv, &s.tail, ids) catch |e| return t.fail(e); + const all = t.pr.logits(c); + logits = all[(t.pr.head - 1) * V ..][0..V]; + const tr = t.pr.tapRows(c); + for (0..t.pr.taps) |i| t.ringPut(s, n - t.pr.taps + i, tr[i * tapWidth(c) ..][0..tapWidth(c)]); + } + t.promptTaps(s, n); + const choice = try t.choose(logits, draw, smp); + s.base = n; + s.live = true; + s.whole = true; + if (t.cache.on() and at < n) t.cache.save(c, &s.kv, ids, s.ring, logits) catch {}; + return choice; + } + + fn promptEnv(t: *const Target) prompt.Env { + const e = t.env; + return .{ .c = e.c, .p = e.p, .tables = e.tables, .model = e.model, .lookup = e.lookup, .queue = e.queue }; + } + + fn ringPut(t: *const Target, s: *Slot, pos: u64, row: []const u16) void { + const w = tapWidth(t.env.c); + if (w == 0) return; + const at: usize = @intCast(pos % ringRows(t.env.c)); + @memcpy(s.ring[at * w ..][0..w], row); + } + + /// The taps of the prompt's last rows (the ring's, up to `ring rows`) in order, for `taps` after a prefill. + fn promptTaps(t: *const Target, s: *Slot, n: u32) void { + const c = t.env.c; + const w = tapWidth(c); + s.after_prompt = true; + s.ptaps_rows = if (w == 0) 0 else @min(n, ringRows(c)); + const out = s.ptaps.slice(u16, @as(usize, s.ptaps_rows) * w); + for (0..s.ptaps_rows) |i| { + const pos = n - s.ptaps_rows + i; + @memcpy(out[i * w ..][0..w], s.ring[(pos % ringRows(c)) * w ..][0..w]); + } + } + + /// The keyed choice of a logits row over the whole vocabulary (greedy: the first maximum, ties to the lower id). + fn choose(t: *const Target, row: []const f32, draw: u64, smp: ?lanes.Sampling) !u32 { + const s = smp orelse return argmax(row); + if (s.temperature <= 0) return argmax(row); + const vals = try t.a.alloc(f64, row.len); + defer t.a.free(vals); + const ids = try t.a.alloc(u64, row.len); + defer t.a.free(ids); + for (vals, ids, row, 0..) |*v, *id, x, i| { + v.* = x; + id.* = i; + } + return @intCast(try lanes.sampling.choose(t.a, vals, ids, draw, s)); + } + + fn argmax(row: []const f32) u32 { + var best: usize = 0; + for (row, 0..) |x, i| if (x > row[best]) { + best = i; + }; + return @intCast(best); + } + + fn fail(t: *Target, e: anyerror) anyerror { + if (e == error.CommandBufferFailed) t.failed = true; + return e; + } + + // windows ------------------------------------------------------------------------------------------------------ + + /// A round's windows: slots' rows share a launch (up to R rows, a slot once), each launch one forward. + fn window(ptr: *anyopaque, segments: []const iface.Segment, choices: [][]u32) anyerror!void { + const t = self(ptr); + var from: usize = 0; + while (from < segments.len) { + var rows: usize = 0; + var to = from; + while (to < segments.len) : (to += 1) { + const g = segments[to]; + if (g.slot >= t.slots.len) return error.NoSlot; + if (to > from and (rows + g.tokens.len > R or dup(segments[from..to], g.slot))) break; + rows += g.tokens.len; + } + try t.windowGroup(segments[from..to], choices[from..to]); + from = to; + } + } + + fn dup(done: []const iface.Segment, slot: u32) bool { + for (done) |o| if (o.slot == slot) return true; + return false; + } + + /// One launch over the rows of these segments (distinct slots, at most R rows in all). + fn windowGroup(t: *Target, group: []const iface.Segment, choices: [][]u32) !void { + const c = t.env.c; + const V: usize = c.vocab; + var total: u32 = 0; + for (group, choices) |g, out| { + const s = &t.slots[g.slot]; + const n = g.tokens.len; + if (!s.live) return error.SlotNotLive; + if (g.start != s.hist.items.len) return error.PositionMismatch; + if (n == 0 or total + n > R or g.draws.len != n or out.len < n) return error.BadRows; + if (g.start + n > t.opt.limit) return error.PastLimit; + // the decode forward runs consecutive positions: a tree's rows must be one chain + if (g.parents) |ps| for (ps, 0..) |q, r| if (q != @as(i32, @intCast(r)) - 1) return error.TreeWindow; + for (g.tokens) |id| if (id >= c.vocab) return error.BadToken; + total += @intCast(n); + } + for (group) |g| { + const s = &t.slots[g.slot]; + s.whole = false; + try s.hist.appendSlice(t.a, g.tokens); + s.high = @max(s.high, s.hist.items.len); + } + t.launch(group) catch |e| return t.fail(e); + const lg = t.b.logitsOf(c, total); + const w = tapWidth(c); + var r0: usize = 0; + for (group, choices) |g, out| { + const s = &t.slots[g.slot]; + const n = g.tokens.len; + for (0..n) |r| out[r] = try t.choose(lg[(r0 + r) * V ..][0..V], g.draws[r], g.sampling); + if (s.kv.taps) |tb| for (0..n) |r| t.ringPut(s, g.start + r, tb.slice(u16, n * w)[r * w ..][0..w]); + s.base = g.start; + s.rows = @intCast(n); + s.open = true; + s.after_prompt = false; + s.tail.end = @intCast(s.hist.items.len); + s.tail.done = s.tail.end; + s.whole = true; + t.stats.windows += 1; + t.stats.rows += n; + r0 += n; + } + t.stats.launches += 1; + } + + /// One launch of the decode forward over the segments' rows, each slot's ids at its positions from `start`. + fn launch(t: *Target, group: []const iface.Segment) !void { + const e = t.env; + const c = e.c; + const look = c.engram_max_ngram - 1; + var pieces: [R]forward.Piece = undefined; + var spans: [R]scratch.Span = undefined; + var segs: [R]ctx_.Seg = undefined; + var r0: u32 = 0; + for (group, 0..) |g, i| { + const s = &t.slots[g.slot]; + const n: u32 = @intCast(g.tokens.len); + const start: u32 = @intCast(g.start); + pieces[i] = .{ .ids = g.tokens, .lookback = s.hist.items[start -| look..start] }; + spans[i] = .{ .p0 = start, .n = n }; + segs[i] = .{ .slot = &s.kv, .p0 = start, .r0 = r0, .n = n }; + r0 += n; + } + try forward.prepareSegs(&t.b, c, pieces[0..group.len], e.lookup); + try e.s.writeSpans(c, spans[0..group.len]); + const x: ctx_.Ctx = .{ .c = c, .p = e.p, .s = e.s, .slot = segs[0].slot, .tables = e.tables, .p0 = segs[0].p0, .rows = r0, .segs = segs[0..group.len] }; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + try forward.run(&x, e.queue, e.model, &t.b); + } + + /// Keep the window's first rows (a chain from row 0): later rows' ids go, the compressor's carry rewinds. + fn keep(ptr: *anyopaque, slot: u32, path: []const u32) anyerror!void { + const t = self(ptr); + if (slot >= t.slots.len) return error.NoSlot; + const s = &t.slots[slot]; + if (!s.open) return error.NoWindow; + if (path.len == 0 or path.len > s.rows) return error.NotAPath; + for (path, 0..) |r, i| if (r != i) return error.NotAPath; + s.open = false; + const m: u32 = @intCast(path.len); + if (m == s.rows) return; + const len = s.base + m; + s.hist.shrinkRetainingCapacity(@intCast(len)); + for (s.kv.regions) |*r| kv.rewindCarry(r, t.env.c, s.base, m); + s.tail.end = @intCast(len); + s.tail.done = s.tail.end; + t.stats.rewinds += 1; + } + + fn taps(ptr: *anyopaque, slot: u32) iface.Taps { + const t = self(ptr); + const s = &t.slots[slot]; + const w: u32 = @intCast(tapWidth(t.env.c)); + if (s.after_prompt) return .{ .address = @intFromPtr(s.ptaps.contents()), .rows = s.ptaps_rows, .stride = w }; + const tb = s.kv.taps orelse return .{ .address = 0, .rows = 0, .stride = w }; + return .{ .address = @intFromPtr(tb.contents()), .rows = s.rows, .stride = w }; + } + + /// The request ended: its state is kept (the prompt cache) when every launch of it completed. + fn release(ptr: *anyopaque, slot: u32) void { + const t = self(ptr); + if (slot >= t.slots.len) return; + const s = &t.slots[slot]; + if (!s.live) return; + s.live = false; + s.open = false; + if (s.whole and t.cache.on()) t.cache.save(t.env.c, &s.kv, s.hist.items, s.ring, null) catch {}; + } + + /// The serial reference without lanes: the prompt, then one row a step, each drawn at its own position. + pub fn serial(t: *Target, slot: u32, ids: []const u32, n: usize, smp: ?lanes.Sampling) ![]u32 { + var out: std.ArrayList(u32) = .empty; + errdefer out.deinit(t.a); + var tok = try prefill(t, slot, ids, smp, ids.len); + while (out.items.len < n) { + try out.append(t.a, tok); + const at: u64 = ids.len + out.items.len - 1; + var ch: [1]u32 = undefined; + var cs = [_][]u32{&ch}; + try window(t, &.{.{ .slot = slot, .start = at, .tokens = &.{tok}, .parents = null, .draws = &.{at + 1}, .sampling = smp }}, &cs); + tok = ch[0]; + } + release(t, slot); + return out.toOwnedSlice(t.a); + } +}; diff --git a/zig/src/families/deepseek_v41/phases_timed.zig b/zig/src/families/deepseek_v41/phases_timed.zig index 5c195fbdbd2..b2c38fadaeb 100644 --- a/zig/src/families/deepseek_v41/phases_timed.zig +++ b/zig/src/families/deepseek_v41/phases_timed.zig @@ -17,7 +17,11 @@ pub const TimedTarget = struct { return .{ .ptr = t, .vtable = &vtable }; } - const vtable: iface.Target.VTable = .{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .admit = admit, .begin = begin, .piece = piece, .finish = finish, .tail = tail, .pieces = pieces, .replays = replays, .warm = warm }; + const vtable: iface.Target.VTable = .{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .admit = admit, .begin = begin, .piece = piece, .finish = finish, .tail = tail, .pieces = pieces, .replays = replays, .warm = warm, .resumed = resumed }; + + fn resumed(p: *anyopaque, slot: u32) u64 { + return self(p).inner.resumedRows(slot); + } fn admit(p: *anyopaque, slot: u32, max_new: u64) void { self(p).inner.admit(slot, max_new); diff --git a/zig/tests/dsv41_metal/boot_test.zig b/zig/tests/dsv41_metal/boot_test.zig new file mode 100644 index 00000000000..62f45a6ca43 --- /dev/null +++ b/zig/tests/dsv41_metal/boot_test.zig @@ -0,0 +1,171 @@ +//! A small release-format checkpoint folder booted the way the server boots one (boot.zig), and served a reply. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const fixtures = @import("dsv41_serve_fixtures"); + +const api = host.api; +const Served = host.serve_lanes.Served; +const Config = host.config.Config; +const testing = std.testing; + +/// The mini tokenizer's special tokens move from ids 1000 .. 1010 to right after its 266 byte-level ones. +const special0 = 266; + +/// The tiny drafted model without Engram, its vocabulary the mini tokenizer's (277 ids in 384; BOS 266, EOS 267). +pub fn config() Config { + var c = fixture.dsparkConfig(); + c.vocab = 384; + c.engram_layers = host.config.Ids.of(&.{}); + c.engram_rows = host.config.List(u64, 8).of(&.{}); + c.bos = special0; + c.eos = special0 + 1; + return c; +} + +/// The mini tokenizer with its special tokens at special0 .. +fn tokenizer(a: std.mem.Allocator) ![]u8 { + var text = try a.dupe(u8, fixtures.mini_tokenizer); + for (0..11) |i| { + var from: [16]u8 = undefined; + var to: [16]u8 = undefined; + const next = try std.mem.replaceOwned(u8, a, text, try std.fmt.bufPrint(&from, "\"id\": {d},", .{1000 + i}), try std.fmt.bufPrint(&to, "\"id\": {d},", .{special0 + i})); + a.free(text); + text = next; + } + return text; +} + +fn list(w: *std.Io.Writer, xs: anytype) !void { + try w.writeByte('['); + for (xs, 0..) |x, i| try w.print("{s}{d}", .{ if (i == 0) "" else ",", x }); + try w.writeByte(']'); +} + +/// config.json as a release checkpoint writes it, at `c`'s shapes. +fn configJson(a: std.mem.Allocator, c: *const Config) ![]u8 { + var out: std.Io.Writer.Allocating = .init(a); + defer out.deinit(); + const w = &out.writer; + try w.print("{{\"model_type\":\"deepseek_v41\",\"bos_token_id\":{d},\"eos_token_id\":{d},", .{ c.bos, c.eos }); + try w.writeAll("\"quantization_config\":{\"quant_method\":\"fp8\",\"weight_block_size\":[32,32],\"scale_fmt\":\"ue8m0\",\"expert_dtype\":\"fp4\"},"); + const ints = .{ + .{ "vocab_size", c.vocab }, .{ "hidden_size", c.hidden }, .{ "num_hidden_layers", c.layers }, + .{ "num_attention_heads", c.heads }, .{ "head_dim", c.head_dim }, .{ "qk_rope_head_dim", c.rope_dim }, + .{ "q_lora_rank", c.q_lora }, .{ "o_lora_rank", c.o_lora }, .{ "o_groups", c.o_groups }, + .{ "sliding_window", c.window }, .{ "n_routed_experts", c.experts }, .{ "n_shared_experts", c.shared_experts }, + .{ "num_experts_per_tok", c.topk }, .{ "moe_intermediate_size", c.expert_width }, .{ "index_n_heads", c.index_heads }, + .{ "index_head_dim", c.index_dim }, .{ "index_topk", c.index_topk }, .{ "candidate_source_layer_id", c.candidate_source }, + .{ "candidate_topk_blocks", c.candidate_blocks }, .{ "candidate_block_size", c.candidate_block_size }, .{ "hc_mult", c.hc_mult }, + .{ "hc_sinkhorn_iters", c.hc_sinkhorn_iters }, .{ "engram_max_ngram_size", c.engram_max_ngram }, .{ "engram_vocab_size", c.engram_vocab }, + .{ "engram_n_heads", c.engram_heads }, .{ "engram_head_dim", c.engram_head_dim }, .{ "engram_pad_token_id", c.engram_pad }, + .{ "engram_compressed_vocab_size", c.engram_compressed_vocab }, .{ "num_nextn_predict_layers", c.mtp_layers }, .{ "dspark_block_size", c.dspark_block }, + .{ "dspark_noise_token_id", c.dspark_noise_token }, .{ "dspark_markov_rank", c.dspark_markov_rank }, .{ "dspark_n_routed_experts", c.dspark_experts }, + .{ "dspark_num_experts_per_tok", c.dspark_topk }, .{ "num_key_value_heads", 1 }, .{ "max_position_embeddings", c.rope.max_positions }, + }; + inline for (ints) |p| try w.print("\"{s}\":{d},", .{ p[0], p[1] }); + const floats = .{ + .{ "rms_norm_eps", c.eps }, .{ "routed_scaling_factor", c.routed_scale }, .{ "swiglu_limit", c.swiglu_limit }, + .{ "hc_eps", c.hc_eps }, .{ "hc_post_alpha", c.hc_post_alpha }, .{ "engram_gate_clamp", c.engram_gate_clamp }, + .{ "rope_theta", c.rope.theta }, .{ "compress_rope_theta", c.rope.compress_theta }, + }; + inline for (floats) |p| try w.print("\"{s}\":{e},", .{ p[0], p[1] }); + const lists = .{ + .{ "compress_ratios", c.compress_ratios.items() }, .{ "kv_source_layer_ids", c.kv_sources.items() }, + .{ "index_source_layer_ids", c.index_sources.items() }, .{ "engram_layer_ids", c.engram_layers.items() }, + .{ "dspark_target_layer_ids", c.dspark_targets.items() }, .{ "engram_num_embeddings", c.engram_rows.items() }, + }; + inline for (lists) |p| { + try w.print("\"{s}\":", .{p[0]}); + try list(w, p[1]); + try w.writeByte(','); + } + try w.print("\"norm_topk_prob\":{},\"scoring_func\":\"sqrtsoftplus\",\"topk_method\":\"noaux_tc\",", .{c.norm_topk}); + try w.print("\"rope_scaling\":{{\"rope_type\":\"yarn\",\"factor\":{e},\"beta_fast\":{e},\"beta_slow\":{e},\"original_max_position_embeddings\":{d}}}}}", .{ c.rope.factor, c.rope.beta_fast, c.rope.beta_slow, c.rope.original_max_positions }); + return out.toOwnedSlice(); +} + +/// The folder: config.json, the weights (3 shards and their index), the mini tokenizer. +pub fn write(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, seed: u64) !void { + var c = config(); + try c.validate(); + const json = try configJson(a, &c); + defer a.free(json); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = json }); + const img = try fixture.image(a, &c, seed); + defer a.free(img); + try @import("shards.zig").write(a, io, dir, img, 3, 0); + const tok = try tokenizer(a); + defer a.free(tok); + try dir.writeFile(io, .{ .sub_path = "tokenizer.json", .data = tok }); +} + +test "dsv41 metal boot: a small checkpoint folder parses back to its config, boots and replies, drafted == draft:false" { + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try write(a, io, tmp.dir, 81); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + // TF_DSV41_SMOKE_DIR: the same folder written there too, for a server's smoke boot + if (std.c.getenv("TF_DSV41_SMOKE_DIR")) |d| { + var out = try std.Io.Dir.cwd().openDir(io, std.mem.span(d), .{}); + defer out.close(io); + try write(a, io, out, 81); + } + const want = config(); + const m = try pkg.boot.Boot.open(a, io, dir, .{ .slots = 2, .context = 1024, .small = true, .cache_bytes = 0 }); + const got = m.parts.env.c; + try testing.expectEqual(want.vocab, got.vocab); + try testing.expectEqualSlices(u32, want.compress_ratios.items(), got.compress_ratios.items()); + try testing.expectEqual(want.eos, got.eos); + const s = try Served.open(a, io, m.model(), dir, null); + defer Served.close(s); + const e = s.engine(); + const prompt = [_]u32{ special0, special0 + 5, 70, 71, 72, 73, special0 + 6 }; + var replies: [2][24]u32 = undefined; + for ([_]bool{ true, false }, &replies, 0..) |drafts, *r, i| { + var box: Box = .{}; + const req: api.Request = .{ .prompt = &prompt, .max_tokens = 24, .drafts = drafts }; + try e.submit(@intCast(900 + i), &req, .{ .ctx = &box, .event = Box.event }); + try box.wait(io); + try testing.expectEqual(@as(usize, 24), box.n); + r.* = box.out; + } + std.debug.print("{s} dsv41 metal boot: {d} layers from 3 shards, a reply of 24 tokens, drafted == draft:false\n", .{ if (std.mem.eql(u32, &replies[0], &replies[1])) "PASS" else "FAIL", got.layers }); + try testing.expectEqualSlices(u32, &replies[1], &replies[0]); +} + +const Box = struct { + mutex: std.Io.Mutex = .init, + out: [24]u32 = undefined, + n: usize = 0, + done: bool = false, + + fn event(ctx: *anyopaque, _: api.Id, e: *const api.Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(testing.io); + defer b.mutex.unlock(testing.io); + switch (e.*) { + .tokens => |ts| for (ts) |t| if (b.n < b.out.len) { + b.out[b.n] = t; + b.n += 1; + }, + .finished => b.done = true, + else => {}, + } + } + + fn wait(b: *Box, io: std.Io) !void { + for (0..300_000) |_| { + b.mutex.lockUncancelable(io); + const done = b.done; + b.mutex.unlock(io); + if (done) return; + std.Io.sleep(io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } +}; diff --git a/zig/tests/dsv41_metal/fixture_serve.zig b/zig/tests/dsv41_metal/fixture_serve.zig new file mode 100644 index 00000000000..8f276cc7c5d --- /dev/null +++ b/zig/tests/dsv41_metal/fixture_serve.zig @@ -0,0 +1,69 @@ +//! The tiny model served on Metal: prompt split, Engram before it, a DSpark drafter; its serve.Metal over the rig. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const ff = @import("fixture_fwd.zig"); + +pub const limit = 1280; + +/// The prompt model (ced's split at layer 4, Engram at 1 and 3) with fixture.dsparkConfig's drafter. +pub fn config() host.config.Config { + var c = fixture.promptConfig(); + const d = fixture.dsparkConfig(); + c.mtp_layers = d.mtp_layers; + c.dspark_block = d.dspark_block; + c.dspark_noise_token = d.dspark_noise_token; + c.dspark_markov_rank = d.dspark_markov_rank; + c.dspark_experts = d.dspark_experts; + c.dspark_topk = d.dspark_topk; + return c; +} + +/// The rig (weights, Engram rows and hashes) and the drafter's weights the served model borrows. +const Held = struct { + a: std.mem.Allocator, + f: *ff.Fwd, + w: pkg.dspark.Weights, + + fn close(ctx: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(ctx)); + h.w.deinit(); + h.f.deinit(); + h.a.destroy(h); + } +}; + +pub const Opts = struct { slots: u32 = 4, drafts: bool = true, cache_bytes: u64 = 0, chunk: u32 = 64 }; + +/// The tiny model from `seed` served on this GPU; the returned model owns the rig (closed with it). +pub fn open(a: std.mem.Allocator, io: std.Io, seed: u64, o: Opts) !*pkg.serve.Metal { + const h = try a.create(Held); + errdefer a.destroy(h); + h.a = a; + h.f = try ff.Fwd.initWith(a, io, config(), limit, seed); + errdefer h.f.deinit(); + const r = h.f.r; + { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const img = try fixture.image(a, &r.c, seed); + defer a.free(img); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = img }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var pack = try host.pack.Pack.open(a, io, dir); + defer pack.deinit(); + h.w = try pkg.dspark.loadWeights(a, io, r.device, &pack, &r.c); + } + errdefer h.w.deinit(); + const env: pkg.target.Env = .{ .device = r.device, .queue = r.queue, .c = &r.c, .p = &r.p, .tables = &r.tables, .s = &r.s, .model = h.f.model(), .lookup = h.f.lookup() }; + const parts: pkg.serve.Parts = .{ .env = env, .top = &r.top.?, .dspark = &h.w, .owner = .{ .ctx = h, .close = Held.close } }; + return pkg.serve.Metal.open(a, parts, .{ .slots = o.slots, .limit = limit, .drafts = o.drafts, .cache_bytes = o.cache_bytes, .prompt = .{ .chunk = o.chunk } }); +} + +/// Closes a model `open` made (the served engine closes its own through serve_lanes). +pub fn close(m: *pkg.serve.Metal) void { + const md = m.model(); + md.close(md.ptr); +} diff --git a/zig/tests/dsv41_metal/layer_test.zig b/zig/tests/dsv41_metal/layer_test.zig index 0cb448c3720..83ed237f9bb 100644 --- a/zig/tests/dsv41_metal/layer_test.zig +++ b/zig/tests/dsv41_metal/layer_test.zig @@ -17,6 +17,7 @@ test { _ = @import("zerocopy_test.zig"); _ = @import("prompt_test.zig"); _ = @import("dspark_test.zig"); + _ = @import("multi_test.zig"); } test "dsv41 metal: every stage of every layer kind within its float64 bound" { diff --git a/zig/tests/dsv41_metal/multi_test.zig b/zig/tests/dsv41_metal/multi_test.zig new file mode 100644 index 00000000000..4819c786c6d --- /dev/null +++ b/zig/tests/dsv41_metal/multi_test.zig @@ -0,0 +1,219 @@ +//! Rows of several slots in one decode launch equal each slot's rows launched alone, bit for bit. +const std = @import("std"); +const mtl = @import("metal"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fixture = @import("fixture.zig"); +const ff = @import("fixture_fwd.zig"); +const fs = @import("fixture_serve.zig"); + +const testing = std.testing; +const fwd = pkg.forward; +const limit = 1280; +const look = 3; + +/// One slot's share of a launch: where it stands and the ids its rows carry (after `look` ids of lookback). +const Part = struct { p0: u32, hist: []u32, n: usize }; + +fn finish(r: *fixture.Rig, f: *ff.Fwd, x: *const pkg.ctx.Ctx, rows: u32) !void { + const pass = fixture.Pass.begin(r.queue); + try fwd.encode(x, pass.e, f.model(), &f.b); + try pass.finish(); + _ = rows; +} + +/// Each part's rows alone on its slot (the single-stream path): logits [rows][vocab] and taps per part, concatenated. +fn alone(a: std.mem.Allocator, f: *ff.Fwd, slots: []pkg.kv.Slot, parts: []const Part, logits: []f32, taps: []u16) !void { + const r = f.r; + const c = &r.c; + const V: usize = c.vocab; + const TH = c.dspark_targets.len * c.hidden; + _ = a; + var r0: usize = 0; + for (parts, slots) |pt, *sl| { + try fwd.prepare(&f.b, c, pt.hist[look..][0..pt.n], pt.hist[0..look], f.lookup()); + try r.s.writeRows(c, pt.p0, @intCast(pt.n)); + const x = r.ctx(sl, pt.p0, @intCast(pt.n)); + try finish(r, f, &x, @intCast(pt.n)); + @memcpy(logits[r0 * V ..][0 .. pt.n * V], f.b.logitsOf(c, pt.n)); + @memcpy(taps[r0 * TH ..][0 .. pt.n * TH], sl.taps.?.slice(u16, pt.n * TH)); + r0 += pt.n; + } +} + +/// All parts' rows in one launch; the taps are read back from each slot. +fn together(f: *ff.Fwd, slots: []pkg.kv.Slot, parts: []const Part, logits: []f32, taps: []u16) !void { + const r = f.r; + const c = &r.c; + const V: usize = c.vocab; + const TH = c.dspark_targets.len * c.hidden; + var pieces: [16]fwd.Piece = undefined; + var spans: [16]pkg.scratch.Span = undefined; + var segs: [16]pkg.ctx.Seg = undefined; + var r0: u32 = 0; + for (parts, slots, 0..) |pt, *sl, i| { + pieces[i] = .{ .ids = pt.hist[look..][0..pt.n], .lookback = pt.hist[0..look] }; + spans[i] = .{ .p0 = pt.p0, .n = @intCast(pt.n) }; + segs[i] = .{ .slot = sl, .p0 = pt.p0, .r0 = r0, .n = @intCast(pt.n) }; + r0 += @intCast(pt.n); + } + const k = parts.len; + try fwd.prepareSegs(&f.b, c, pieces[0..k], f.lookup()); + try r.s.writeSpans(c, spans[0..k]); + var x = r.ctx(&slots[0], parts[0].p0, r0); + x.segs = segs[0..k]; + try finish(r, f, &x, r0); + @memcpy(logits[0 .. r0 * V], f.b.logitsOf(c, r0)); + var at: usize = 0; + for (parts, slots) |pt, *sl| { + @memcpy(taps[at * TH ..][0 .. pt.n * TH], sl.taps.?.slice(u16, pt.n * TH)); + at += pt.n; + } +} + +fn reset(r: *fixture.Rig, to: *const pkg.kv.Slot, from: *const pkg.kv.Slot) !void { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = r.queue.commandBuffer(); + to.copyFrom(cb, from); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.CommandBufferFailed; +} + +/// Bytes that differ between two slots' every region: window, pool, keys, carry and the carry's inputs. +fn stateDiff(x: *const pkg.kv.Slot, y: *const pkg.kv.Slot) usize { + var n: usize = 0; + for (x.regions, y.regions) |p, q| { + const pairs = [_][2]?mtl.Buffer{ .{ p.window, q.window }, .{ p.pool, q.pool }, .{ p.keys, q.keys }, .{ p.carry, q.carry }, .{ p.inputs, q.inputs } }; + for (pairs) |pr| if (pr[0]) |bx| { + const len = bx.length(); + for (bx.contents()[0..len], pr[1].?.contents()[0..len]) |s, t| n += @intFromBool(s != t); + }; + } + return n; +} + +/// Slots at `starts` with `rows` rows each: one launch of all vs one launch each, bit for bit. +fn compare(a: std.mem.Allocator, f: *ff.Fwd, starts: []const u32, rows: []const usize, seed: u64) !usize { + const r = f.r; + const c = &r.c; + const k = starts.len; + var pristine: [4]pkg.kv.Slot = undefined; + var one: [4]pkg.kv.Slot = undefined; + var many: [4]pkg.kv.Slot = undefined; + var parts: [4]Part = undefined; + var made: usize = 0; + defer for (0..made) |i| { + pristine[i].deinit(); + one[i].deinit(); + many[i].deinit(); + a.free(parts[i].hist); + }; + var total: usize = 0; + for (0..k) |i| { + pristine[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + one[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + many[i] = try pkg.kv.Slot.init(a, r.device, c, limit); + parts[i] = .{ .p0 = starts[i], .hist = try ff.ids(a, c, look + rows[i], seed + 10 * i), .n = rows[i] }; + made = i + 1; + fixture.fillPast(c, &pristine[i], starts[i], seed + 100 * i); + try reset(r, &one[i], &pristine[i]); + try reset(r, &many[i], &pristine[i]); + total += rows[i]; + } + const TH = c.dspark_targets.len * c.hidden; + const l1 = try a.alloc(f32, total * c.vocab); + defer a.free(l1); + const ln = try a.alloc(f32, total * c.vocab); + defer a.free(ln); + const t1 = try a.alloc(u16, total * TH); + defer a.free(t1); + const tn = try a.alloc(u16, total * TH); + defer a.free(tn); + try alone(a, f, one[0..k], parts[0..k], l1, t1); + try together(f, many[0..k], parts[0..k], ln, tn); + var bad: usize = 0; + for (l1, ln) |s, t| bad += @intFromBool(@as(u32, @bitCast(s)) != @as(u32, @bitCast(t))); + var tb: usize = 0; + for (t1, tn) |s, t| tb += @intFromBool(s != t); + var st: usize = 0; + for (0..k) |i| st += stateDiff(&one[i], &many[i]); + const word = if (bad + tb + st == 0) "PASS" else "FAIL"; + std.debug.print("{s} dsv41 metal multi-slot forward: {d} slots, {d} rows in one launch vs one launch a slot: {d} logits, {d} tap values, {d} state bytes differ\n", .{ word, k, total, bad, tb, st }); + return bad + tb + st; +} + +test "dsv41 metal multi-slot forward: rows of 2, 3 and 4 slots in one launch equal each slot's launch alone" { + const a = testing.allocator; + const f = try ff.Fwd.init(a, testing.io, limit, 61); + defer f.deinit(); + var bad: usize = 0; + // positions of both parities, a slot near the start and one past the window; ratio 2 pools complete on odd rows + bad += try compare(a, f, &.{ 1101, 1104 }, &.{ 3, 5 }, 300); + bad += try compare(a, f, &.{ 1103, 640, 1100 }, &.{ 1, 4, 2 }, 310); + bad += try compare(a, f, &.{ 1101, 1104, 1099, 97 }, &.{ 4, 4, 4, 4 }, 320); + bad += try compare(a, f, &.{ 1101, 1102, 1103, 1104 }, &.{ 1, 1, 1, 1 }, 330); + bad += try compare(a, f, &.{ 7, 128, 1101, 1279 - 6 }, &.{ 2, 6, 1, 7 }, 340); + try testing.expectEqual(@as(usize, 0), bad); +} + +/// A model with the tiny drafter's config and prompts of different lengths on its slots. +fn rounds(m: *pkg.serve.Metal, packed_: bool, out: *std.ArrayList(u32)) !void { + const a = testing.allocator; + const tg = m.t.target(); + const lens = [4]usize{ 150, 211, 64, 333 }; + var next: [4]u32 = undefined; + var prompts: [4][]u32 = undefined; + for (&prompts, 0..) |*p, i| p.* = try ff.ids(a, &fs.config(), lens[i], 400 + i); + defer for (prompts) |p| a.free(p); + for (0..4) |i| next[i] = try tg.prefill(@intCast(i), prompts[i], null, lens[i]); + var pos = lens; + // verify rounds of one anchor and three drafts a slot; the keep cuts back a different number each time + for (0..4) |round| { + var toks: [4][4]u32 = undefined; + var draws: [4][4]u64 = undefined; + var chs: [4][4]u32 = undefined; + var segs: [4]host.iface.Segment = undefined; + var cs: [4][]u32 = undefined; + for (0..4) |i| { + toks[i][0] = next[i]; + for (1..4) |j| toks[i][j] = @intCast((next[i] * 7 + j * 13 + round * 5 + i) % 200); + for (&draws[i], 0..) |*d, j| d.* = pos[i] + 1 + j; + segs[i] = .{ .slot = @intCast(i), .start = pos[i], .tokens = &toks[i], .parents = null, .draws = &draws[i], .sampling = null }; + cs[i] = &chs[i]; + } + if (packed_) try tg.window(&segs, &cs) else for (0..4) |i| try tg.window(segs[i .. i + 1], cs[i .. i + 1]); + for (0..4) |i| { + const kept: u32 = @intCast(1 + (i + round) % 4); + const path = [_]u32{ 0, 1, 2, 3 }; + try tg.keep(@intCast(i), path[0..kept]); + try out.appendSlice(a, &chs[i]); + pos[i] += kept; + next[i] = chs[i][kept - 1]; + } + } + for (0..4) |i| tg.release(@intCast(i)); +} + +test "dsv41 metal multi-slot forward: a 4-slot DSpark verify round in one launch equals per-slot launches" { + const a = testing.allocator; + const m1 = try fs.open(a, testing.io, 71, .{}); + defer fs.close(m1); + const m2 = try fs.open(a, testing.io, 71, .{}); + defer fs.close(m2); + var o1: std.ArrayList(u32) = .empty; + defer o1.deinit(a); + var o2: std.ArrayList(u32) = .empty; + defer o2.deinit(a); + try rounds(m1, true, &o1); + try rounds(m2, false, &o2); + var st: usize = 0; + for (m1.t.slots, m2.t.slots) |*x, *y| st += stateDiff(&x.kv, &y.kv); + const same = std.mem.eql(u32, o1.items, o2.items); + const word = if (same and st == 0) "PASS" else "FAIL"; + std.debug.print("{s} dsv41 metal multi-slot verify rounds: {d} choices, {d} launches in one vs {d} apart, {d} state bytes differ\n", .{ word, o1.items.len, m1.t.stats.launches, m2.t.stats.launches, st }); + try testing.expect(same); + try testing.expectEqual(@as(usize, 0), st); + try testing.expect(m1.t.stats.launches < m2.t.stats.launches); +} diff --git a/zig/tests/dsv41_metal/serve_test.zig b/zig/tests/dsv41_metal/serve_test.zig new file mode 100644 index 00000000000..9bf82fd51a5 --- /dev/null +++ b/zig/tests/dsv41_metal/serve_test.zig @@ -0,0 +1,267 @@ +//! The tiny model served on Metal through serve_lanes and LaneHost: replies equal the target's own serial decode. +const std = @import("std"); +const pkg = @import("dsv41_metal"); +const host = @import("dsv41"); +const fs = @import("fixture_serve.zig"); +const ff = @import("fixture_fwd.zig"); + +const api = host.api; +const Served = host.serve_lanes.Served; +const testing = std.testing; +const gpa = testing.allocator; +const io = testing.io; +const seed = 71; + +test { + _ = @import("boot_test.zig"); +} + +/// A reply's tokens and finish, written on LaneHost's thread and read on the test's. +const Box = struct { + mutex: std.Io.Mutex = .init, + out: std.ArrayList(u32) = .empty, + done: bool = false, + reason: ?api.Reason = null, + drafted: u64 = 0, + cached: u32 = 0, + + fn event(ctx: *anyopaque, _: api.Id, e: *const api.Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(io); + defer b.mutex.unlock(io); + switch (e.*) { + .tokens => |ts| b.out.appendSlice(gpa, ts) catch {}, + .finished => |fin| { + b.done = true; + b.reason = fin.reason; + b.drafted = fin.stats.drafted; + }, + .prefilled => |n| b.cached = n, + else => {}, + } + } + + fn wait(b: *Box) !void { + for (0..600_000) |_| { + b.mutex.lockUncancelable(io); + const done = b.done; + b.mutex.unlock(io); + if (done) return; + std.Io.sleep(io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; + } + + fn deinit(b: *Box) void { + b.out.deinit(gpa); + } +}; + +/// One request to send: its prompt, length, drafts and sampling. +const Ask = struct { prompt: []const u32, max: u32 = 24, drafts: bool = true, sampling: ?api.Sampling = null }; + +var next_id: api.Id = 1; +/// The prompt rows each reply of the last `serve` call took from kept state (its `prefilled` event), in ask order. +var cached_rows: [8]u32 = undefined; + +/// The asks sent together (or each after the last finished); each reply returned gpa-owned, the drafted rows summed. +fn serve(e: api.Engine, asks: []const Ask, together: bool, out: [][]u32, drafted: *u64) !void { + const boxes = try gpa.alloc(Box, asks.len); + defer gpa.free(boxes); + @memset(boxes, .{}); + defer for (boxes) |*b| b.deinit(); + const reqs = try gpa.alloc(api.Request, asks.len); + defer gpa.free(reqs); + for (reqs, asks) |*r, k| r.* = .{ .prompt = k.prompt, .max_tokens = k.max, .drafts = k.drafts, .sampling = k.sampling }; + for (reqs, boxes) |*r, *b| { + next_id += 1; + try e.submit(next_id, r, .{ .ctx = b, .event = Box.event }); + if (!together) try b.wait(); + } + for (boxes, out, asks, 0..) |*b, *o, k, i| { + try b.wait(); + try testing.expectEqual(api.Reason.length, b.reason.?); + try testing.expectEqual(@as(usize, k.max), b.out.items.len); + drafted.* += b.drafted; + cached_rows[i] = b.cached; + o.* = try b.out.toOwnedSlice(gpa); + } +} + +fn freeAll(xs: [][]u32) void { + for (xs) |x| gpa.free(x); +} + +/// The reference: the target's serial decode of each ask on a model of its own (no lanes, no cache, no drafts). +fn reference(asks: []const Ask, out: [][]u32) !void { + const m = try fs.open(gpa, io, seed, .{ .slots = 1, .drafts = false }); + defer fs.close(m); + for (asks, out) |k, *o| o.* = try m.t.serial(0, k.prompt, k.max, k.sampling); +} + +fn ids(n: usize, s: u64) ![]u32 { + return ff.ids(gpa, &fs.config(), n, s); +} + +fn say(ok: bool, comptime fmt: []const u8, args: anytype) void { + std.debug.print("{s} dsv41 metal serve: ", .{if (ok) "PASS" else "FAIL"}); + std.debug.print(fmt ++ "\n", args); +} + +fn same(want: []const []u32, got: []const []u32) usize { + var bad: usize = 0; + for (want, got) |w, g| bad += @intFromBool(!std.mem.eql(u32, w, g)); + return bad; +} + +test "dsv41 metal serve: (a) concurrent == solo == serial, (b) drafted == draft:false, through LaneHost" { + const p0 = try ids(150, 1); + defer gpa.free(p0); + const p1 = try ids(620, 2); + defer gpa.free(p1); + const asks = [_]Ask{ .{ .prompt = p0 }, .{ .prompt = p1 } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{}); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var solo: [2][]u32 = undefined; + try serve(e, &asks, false, &solo, &drafted); + defer freeAll(&solo); + var both: [2][]u32 = undefined; + try serve(e, &asks, true, &both, &drafted); + defer freeAll(&both); + const plain_asks = [_]Ask{ .{ .prompt = p0, .drafts = false }, .{ .prompt = p1, .drafts = false } }; + var plain_drafted: u64 = 0; + var plain: [2][]u32 = undefined; + try serve(e, &plain_asks, true, &plain, &plain_drafted); + defer freeAll(&plain); + const bad = [3]usize{ same(&want, &solo), same(&want, &both), same(&want, &plain) }; + say(bad[0] == 0, "(a) solo vs serial: {d} of 2 replies differ", .{bad[0]}); + say(bad[1] == 0, "(a) together vs serial: {d} of 2 replies differ ({d} windows of {d} rows in all)", .{ bad[1], m.t.stats.windows, m.t.stats.rows }); + say(bad[2] == 0 and drafted > 0 and plain_drafted == 0, "(b) draft:false vs serial: {d} differ; drafted rows {d} (drafts on), {d} (off), {d} pass launches, {d} windows cut back", .{ bad[2], drafted, plain_drafted, m.ad.passes, m.t.stats.rewinds }); + try testing.expectEqual([3]usize{ 0, 0, 0 }, bad); + try testing.expect(drafted > 0 and m.t.stats.rewinds > 0); + try testing.expectEqual(@as(u64, 0), plain_drafted); +} + +test "dsv41 metal serve: (c) prompt-cache hits (the same prompt, a later turn) == fresh runs" { + const p = try ids(300, 3); + defer gpa.free(p); + const m = try fs.open(gpa, io, seed, .{ .cache_bytes = 1 << 30 }); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var first: [1][]u32 = undefined; + try serve(e, &.{.{ .prompt = p, .max = 40 }}, false, &first, &drafted); + defer freeAll(&first); + const first_cached = cached_rows[0]; + // a later turn: the prompt, the reply but its last token, new ids; then the first prompt again + const extra = try ids(50, 4); + defer gpa.free(extra); + const turn = try std.mem.concat(gpa, u32, &.{ p, first[0][0..39], extra }); + defer gpa.free(turn); + const branch = try std.mem.concat(gpa, u32, &.{ p, extra }); + defer gpa.free(branch); + const asks = [_]Ask{ .{ .prompt = turn, .max = 30 }, .{ .prompt = p, .max = 40 }, .{ .prompt = branch, .max = 30 } }; + const hits0 = m.t.stats.hits; + var got: [3][]u32 = undefined; + try serve(e, &asks, false, &got, &drafted); + defer freeAll(&got); + const hits = m.t.stats.hits - hits0; + const reported = cached_rows[0..3].*; + var want: [3][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const bad = same(&want, &got); + say(bad == 0 and hits == 3, "(c) 3 cached prompts ({d} hits, {d} rows restored in all) vs fresh: {d} replies differ", .{ hits, m.t.stats.cached_rows, bad }); + try testing.expectEqual(@as(usize, 0), bad); + try testing.expectEqual(@as(u64, 3), hits); + // each reply reports the rows it resumed (prompt_tokens_details.cached_tokens) + say(reported[0] > p.len and reported[1] >= p.len and reported[2] >= p.len, "(c) cached rows reported {any} for prompts of {d}, {d} and {d} tokens", .{ reported, turn.len, p.len, branch.len }); + for (reported, [3]usize{ turn.len, p.len, branch.len }) |r, n| try testing.expect(r > 0 and r <= n); + try testing.expect(reported[0] > p.len and reported[1] >= p.len and reported[2] >= p.len); + try testing.expectEqual(@as(u32, 0), first_cached); + // the later turn resumed past the first prompt: its reply's rows came from the cache too + try testing.expect(m.t.stats.cached_rows >= 3 * p.len + 30); +} + +test "dsv41 metal serve: (d) sampled replies are their seed's: again, drafted or not, together, serial" { + const p0 = try ids(200, 5); + defer gpa.free(p0); + const p1 = try ids(90, 6); + defer gpa.free(p1); + const hot: api.Sampling = .{ .seed = 1234, .temperature = 0.9, .top_k = 0, .top_p = 1.0 }; + const nucleus: api.Sampling = .{ .seed = 99, .temperature = 0.7, .top_k = 40, .top_p = 0.9 }; + const asks = [_]Ask{ .{ .prompt = p0, .sampling = hot }, .{ .prompt = p1, .sampling = nucleus } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{}); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + var drafted: u64 = 0; + var a1: [2][]u32 = undefined; + try serve(e, &asks, true, &a1, &drafted); + defer freeAll(&a1); + var a2: [2][]u32 = undefined; + try serve(e, &asks, false, &a2, &drafted); + defer freeAll(&a2); + const plain_asks = [_]Ask{ .{ .prompt = p0, .sampling = hot, .drafts = false }, .{ .prompt = p1, .sampling = nucleus, .drafts = false } }; + var a3: [2][]u32 = undefined; + try serve(e, &plain_asks, true, &a3, &drafted); + defer freeAll(&a3); + var other_seed = hot; + other_seed.seed = 1235; + var a4: [1][]u32 = undefined; + try serve(e, &.{.{ .prompt = p0, .sampling = other_seed }}, false, &a4, &drafted); + defer freeAll(&a4); + const bad = [3]usize{ same(&want, &a1), same(&want, &a2), same(&want, &a3) }; + const moved = !std.mem.eql(u32, a4[0], want[0]); + say(bad[0] + bad[1] + bad[2] == 0 and moved, "(d) together / again / draft:false vs serial: {d} / {d} / {d} differ; another seed's reply {s}", .{ bad[0], bad[1], bad[2], if (moved) "differs" else "is the same" }); + try testing.expectEqual([3]usize{ 0, 0, 0 }, bad); + try testing.expect(moved); +} + +test "dsv41 metal serve: (e) a request cancelled during its prompt leaves the slot and the next request valid" { + const long = try ids(1100, 7); + defer gpa.free(long); + const p0 = try ids(120, 8); + defer gpa.free(p0); + const p1 = try ids(140, 9); + defer gpa.free(p1); + const asks = [_]Ask{ .{ .prompt = p0, .max = 40 }, .{ .prompt = p1 } }; + var want: [2][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{ .slots = 2, .cache_bytes = 1 << 30 }); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + const e = s.engine(); + // a decoding stream, then the long prompt (cancelled as soon as it is submitted), then a new request + var bg: Box = .{}; + defer bg.deinit(); + const r0: api.Request = .{ .prompt = p0, .max_tokens = 40 }; + try e.submit(500, &r0, .{ .ctx = &bg, .event = Box.event }); + var gone: Box = .{}; + defer gone.deinit(); + const rl: api.Request = .{ .prompt = long, .max_tokens = 64 }; + try e.submit(501, &rl, .{ .ctx = &gone, .event = Box.event }); + e.cancel(501); + try gone.wait(); + var drafted: u64 = 0; + var after: [1][]u32 = undefined; + try serve(e, asks[1..], false, &after, &drafted); + defer freeAll(&after); + try bg.wait(); + const ok0 = std.mem.eql(u32, want[0], bg.out.items); + const ok1 = std.mem.eql(u32, want[1], after[0]); + say(ok0 and ok1, "(e) cancelled ({t}, {d} tokens out): the decoding reply {s}, the next request's {s} the serial one", .{ gone.reason.?, gone.out.items.len, if (ok0) "equals" else "differs from", if (ok1) "equals" else "differs from" }); + try testing.expect(gone.reason.? == .cancelled); + try testing.expect(ok0 and ok1); +} From 24a1f2f88e66a2845f2133a4849595aa84457e06 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Tue, 10 Nov 2026 21:21:46 +0100 Subject: [PATCH 020/186] dsv41 metal: decode plain at four or more active streams, where drafts lose --- CHANGELOG.md | 1 + .../families/deepseek_v41/draft/calib_env.zig | 6 ++ zig/src/families/deepseek_v41/draft/depth.zig | 1 + zig/src/families/deepseek_v41/draft/draft.zig | 1 + .../deepseek_v41/draft/lane_owner.zig | 9 ++- zig/src/families/deepseek_v41/draft/lanes.zig | 1 + .../deepseek_v41/draft/plain_streams_test.zig | 80 +++++++++++++++++++ .../families/deepseek_v41/draft/policy.zig | 4 + zig/src/families/deepseek_v41/serve_lanes.zig | 2 +- 9 files changed, 103 insertions(+), 2 deletions(-) create mode 100644 zig/src/families/deepseek_v41/draft/plain_streams_test.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index e4c3df20d79..022096db823 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -82,6 +82,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Qwen3.8-27B on Metal: three glue kernels nothing launched (`q27_norm`, `q27_mlp`, `q27_gdn_post`) and their host mirrors are removed; saved prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#2816, part of #2251). Thanks to @jschmied. - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. +- DeepSeek-V4.1 on Metal: rounds with four or more active streams decode plain instead of drafting, since drafts lose there (#1720, part of #1676). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. diff --git a/zig/src/families/deepseek_v41/draft/calib_env.zig b/zig/src/families/deepseek_v41/draft/calib_env.zig index 9505e344ac3..187987d3ea8 100644 --- a/zig/src/families/deepseek_v41/draft/calib_env.zig +++ b/zig/src/families/deepseek_v41/draft/calib_env.zig @@ -225,6 +225,7 @@ pub fn depthSettings(get: *const fn ([*:0]const u8) ?[]const u8) !depth.Settings s.fair = try float(get("TF_DSV41_DEPTH_FAIR"), joint.fair, 0.0, 4.0); s.host_ms = try float(get("TF_DSV41_DEPTH_HOST_MS"), joint.host_ms, 0.0, 100.0); if (trimmed(get("TF_DSV41_GRAPH_ROWS_MAX"))) |v| s.max_rows = @max(1, @min(try std.fmt.parseInt(u64, v, 10), costs_mod.max_rows)); + if (trimmed(get("TF_DSV41_DRAFT_PLAIN_STREAMS"))) |v| s.plain_streams = try std.fmt.parseInt(usize, v, 10); return s; } @@ -245,6 +246,11 @@ pub fn env(name: [*:0]const u8) ?[]const u8 { return if (std.c.getenv(name)) |v| std.mem.span(v) else null; } +/// TF_DSV41_DRAFT_PLAIN_STREAMS as depthSettings reads it (set and non-blank). +pub fn plainStreamsSet() bool { + return trimmed(env("TF_DSV41_DRAFT_PLAIN_STREAMS")) != null; +} + fn trimmed(v: ?[]const u8) ?[]const u8 { const t = std.mem.trim(u8, v orelse return null, " \t\r\n"); return if (t.len == 0) null else t; diff --git a/zig/src/families/deepseek_v41/draft/depth.zig b/zig/src/families/deepseek_v41/draft/depth.zig index f1953f6cfa2..e07abfad2f3 100644 --- a/zig/src/families/deepseek_v41/draft/depth.zig +++ b/zig/src/families/deepseek_v41/draft/depth.zig @@ -22,6 +22,7 @@ pub const Settings = struct { max_rows: u64 = 64, code_accept: bool = false, // TF_DSV41_CODE_ACCEPT: expand high-acceptance short-context chains context_bucket: u64 = 2048, // TF_DSV41_GRAPH_BUCKET; short contexts only use its linear buckets + plain_streams: usize = 0, // TF_DSV41_DRAFT_PLAIN_STREAMS: a round with this many active streams or more decodes plain; 0 = off }; /// [E(0) .. E(len)]: tokens a round verifying k drafts commits on average, into `out` (len + 1). diff --git a/zig/src/families/deepseek_v41/draft/draft.zig b/zig/src/families/deepseek_v41/draft/draft.zig index 1cb2f6eee43..69559911c81 100644 --- a/zig/src/families/deepseek_v41/draft/draft.zig +++ b/zig/src/families/deepseek_v41/draft/draft.zig @@ -38,6 +38,7 @@ test { _ = branches; _ = calib_env; _ = @import("calib_env_test.zig"); + _ = @import("plain_streams_test.zig"); _ = calib_measure; _ = @import("calib_measure_test.zig"); _ = branches_oracle; diff --git a/zig/src/families/deepseek_v41/draft/lane_owner.zig b/zig/src/families/deepseek_v41/draft/lane_owner.zig index f9846c91977..d0957002833 100644 --- a/zig/src/families/deepseek_v41/draft/lane_owner.zig +++ b/zig/src/families/deepseek_v41/draft/lane_owner.zig @@ -5,12 +5,19 @@ const iface = @import("iface.zig"); const costs = @import("costs.zig"); const Policy = @import("policy.zig").Policy; const Spec = @import("spec.zig").Spec; -pub fn init(gpa: std.mem.Allocator, target: iface.Target, pass: iface.Pass, prices: costs.Costs, opts: ln.Options) !*ln.Lanes { +pub fn init(gpa: std.mem.Allocator, target: iface.Target, pass: iface.Pass, prices: costs.Costs, opts_in: ln.Options) !*ln.Lanes { + var opts = opts_in; const b = opts.shape.block; const c = opts.shape.candidates; if (b == 0 or b > ln.max_block or c == 0 or c > 1024 or opts.slots == 0 or opts.slots > 16 or opts.siblings > @import("tree.zig").max_siblings) return error.BadOptions; const x = try gpa.create(ln.Lanes); errdefer gpa.destroy(x); + const set = blk: { + var s = opts.depth; + if (!@import("calib_env.zig").plainStreamsSet() and s.plain_streams == 0) s.plain_streams = opts.plain_streams_default; + break :blk s; + }; + opts.depth = set; var policy = try Policy.init(gpa, prices, opts); errdefer policy.deinit(); const drafts = try gpa.alloc(u32, opts.slots * b); diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index dec67d5e17a..df344d2b76f 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -24,6 +24,7 @@ pub const Options = struct { dup: tree.Dup = .{}, // sibling pricing: TF_DSV41_TREE_DUP_MS a sibling, its pending row counted slots: u32 = 4, // streams at once spec: spec_mod.Settings = .{}, // the speculative pass (TF_DSV41_SPEC_DRAFT / _SPEC_NUCLEUS) + plain_streams_default: usize = 0, // depth.plain_streams when TF_DSV41_DRAFT_PLAIN_STREAMS is unset }; /// A stream's slot and what the backend holds for it between rounds. diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig new file mode 100644 index 00000000000..19f6be10864 --- /dev/null +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -0,0 +1,80 @@ +//! TF_DSV41_DRAFT_PLAIN_STREAMS: the env knob, and a round with enough active windows decodes plain. +const std = @import("std"); +const lanes = @import("lanes"); +const twin = @import("twin.zig"); +const ln = @import("lanes.zig"); +const costs_mod = @import("costs.zig"); +const depth = @import("depth.zig"); +const calib_env = @import("calib_env.zig"); + +const gpa = std.testing.allocator; + +const Env = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: [*:0]const u8) ?[]const u8 { + for (pairs) |kv| if (std.mem.eql(u8, kv[0], std.mem.span(name))) return kv[1]; + return null; + } +}; + +test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, bad values error" { + Env.pairs = &.{}; + try std.testing.expectEqual(@as(usize, 0), (try calib_env.depthSettings(&Env.get)).plain_streams); + Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "4" }}; + try std.testing.expectEqual(@as(usize, 4), (try calib_env.depthSettings(&Env.get)).plain_streams); + Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "oops" }}; + try std.testing.expectError(error.InvalidCharacter, calib_env.depthSettings(&Env.get)); + Env.pairs = &.{}; +} + +fn runPlain(cases: []const []const u32, setup: Setup) !Run { + const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; // a confident drafter (lanes_test's): drafts are chosen + const t = try twin.Twin.init(gpa, dims, cases.len + 1); + defer t.deinit(); + const d = try twin.Drafter.init(gpa, t, cases.len + 1); + defer d.deinit(); + var c = try costs_mod.defaults(gpa, 64); + defer c.deinit(gpa); + const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = setup }); + defer x.deinit(); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + var cfg = try lanes.Config.init(gpa, try x.model(arena.allocator(), 64), 16, 15); + defer cfg.deinit(gpa); + var clock: lanes.fake.FixedClock = .{}; + var e = lanes.Engine.init(gpa, &cfg, x.backend(), clock.clock()); + defer e.deinit(); + x.attach(&e); + const streams = try gpa.alloc(lanes.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |p, *s| s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = 40 }); + defer for (streams) |*s| s.deinit(gpa); + for (streams) |*s| try e.addStream(s); + while (e.activeCount() > 0) try e.step(); + var ok = true; + for (streams) |*s| ok = ok and s.emitted().len > 0; + return .{ .drafted = e.drafted, .ok = ok }; +} + +const Setup = depth.Settings; + +const Run = struct { drafted: u64, ok: bool }; + +const a1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; + +test "a round with plain_streams or more active windows decodes plain; below it drafts; 0 keeps drafts" { + // identical prompts and greedy sampling keep every round a full 4-window round until the streams end together + const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; + // 4 active, threshold 4: every drafted window is forced to count 0 + const plain = try runPlain(&four, .{ .plain_streams = 4 }); + try std.testing.expect(plain.ok); + try std.testing.expectEqual(@as(u64, 0), plain.drafted); + // 3 active, threshold 4: below it, drafts are kept + const below = try runPlain(four[0..3], .{ .plain_streams = 4 }); + try std.testing.expect(below.ok); + try std.testing.expect(below.drafted > 0); + // 4 active, threshold 0 (off): drafts are kept + const off = try runPlain(&four, .{}); + try std.testing.expect(off.ok); + try std.testing.expect(off.drafted > 0); +} diff --git a/zig/src/families/deepseek_v41/draft/policy.zig b/zig/src/families/deepseek_v41/draft/policy.zig index dff9d8efd6d..b403291194d 100644 --- a/zig/src/families/deepseek_v41/draft/policy.zig +++ b/zig/src/families/deepseek_v41/draft/policy.zig @@ -86,6 +86,10 @@ pub const Policy = struct { try asks.append(arena, .{ .slot = st.?.slot, .conf = conf, .most = main, .first = st.?.held[0], .position = st.?.start }); try at.append(arena, i); } + if (p.depth.set.plain_streams > 0 and windows.len >= p.depth.set.plain_streams) { + for (at.items) |i| out[i] = .{ .count = 0 }; + return; + } if (asks.items.len == 0) return; const single = alone and asks.items.len == 1 and other_windows == 0; if (single) { diff --git a/zig/src/families/deepseek_v41/serve_lanes.zig b/zig/src/families/deepseek_v41/serve_lanes.zig index 72e5f4408f0..3a49e2ec576 100644 --- a/zig/src/families/deepseek_v41/serve_lanes.zig +++ b/zig/src/families/deepseek_v41/serve_lanes.zig @@ -121,7 +121,7 @@ pub const Served = struct { } // draft trees (TF_DSV41_TREE / _DUP_MS / _ROWS; off by default): first-position siblings const trees = if (dspark_on) try tree.Settings.fromEnv() else tree.Settings{}; - s.x = try ln.Lanes.init(gpa, tgt, pass, s.costs, .{ .shape = backend.shape, .slots = backend.slots, .siblings = trees.siblings, .sib_rows = trees.sib_rows, .dup = trees.dup, .depth = try calib_env.depthSettings(&calib_env.env), .spec = if (dspark_on) try spec.Settings.fromEnv() else .{} }); + s.x = try ln.Lanes.init(gpa, tgt, pass, s.costs, .{ .shape = backend.shape, .slots = backend.slots, .siblings = trees.siblings, .sib_rows = trees.sib_rows, .dup = trees.dup, .depth = try calib_env.depthSettings(&calib_env.env), .spec = if (dspark_on) try spec.Settings.fromEnv() else .{}, .plain_streams_default = 4 }); s.arena = std.heap.ArenaAllocator.init(gpa); s.cfg = try lanes.Config.init(gpa, try s.x.model(s.arena.allocator(), 64), 16, 15); s.x.policy.depth.forward_rows = backend.forward_rows; From 3f8beecbf5a6de914f13ecb17f26e00590d925e0 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 05:19:48 +0100 Subject: [PATCH 021/186] dsv41 metal: decode plain from three active streams, not four --- .../deepseek_v41/draft/plain_streams_test.zig | 16 ++++++++++------ zig/src/families/deepseek_v41/serve_lanes.zig | 2 +- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 19f6be10864..44abbbfa8dd 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -20,8 +20,8 @@ const Env = struct { test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, bad values error" { Env.pairs = &.{}; try std.testing.expectEqual(@as(usize, 0), (try calib_env.depthSettings(&Env.get)).plain_streams); - Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "4" }}; - try std.testing.expectEqual(@as(usize, 4), (try calib_env.depthSettings(&Env.get)).plain_streams); + Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "3" }}; + try std.testing.expectEqual(@as(usize, 3), (try calib_env.depthSettings(&Env.get)).plain_streams); Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "oops" }}; try std.testing.expectError(error.InvalidCharacter, calib_env.depthSettings(&Env.get)); Env.pairs = &.{}; @@ -65,12 +65,16 @@ const a1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; test "a round with plain_streams or more active windows decodes plain; below it drafts; 0 keeps drafts" { // identical prompts and greedy sampling keep every round a full 4-window round until the streams end together const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; - // 4 active, threshold 4: every drafted window is forced to count 0 - const plain = try runPlain(&four, .{ .plain_streams = 4 }); + // 4 active, threshold 3: every drafted window is forced to count 0 + const plain = try runPlain(&four, .{ .plain_streams = 3 }); try std.testing.expect(plain.ok); try std.testing.expectEqual(@as(u64, 0), plain.drafted); - // 3 active, threshold 4: below it, drafts are kept - const below = try runPlain(four[0..3], .{ .plain_streams = 4 }); + // 3 active, threshold 3: at it, plain too + const at = try runPlain(four[0..3], .{ .plain_streams = 3 }); + try std.testing.expect(at.ok); + try std.testing.expectEqual(@as(u64, 0), at.drafted); + // 2 active, threshold 3: below it, drafts are kept + const below = try runPlain(four[0..2], .{ .plain_streams = 3 }); try std.testing.expect(below.ok); try std.testing.expect(below.drafted > 0); // 4 active, threshold 0 (off): drafts are kept diff --git a/zig/src/families/deepseek_v41/serve_lanes.zig b/zig/src/families/deepseek_v41/serve_lanes.zig index 3a49e2ec576..2b526fc8eb2 100644 --- a/zig/src/families/deepseek_v41/serve_lanes.zig +++ b/zig/src/families/deepseek_v41/serve_lanes.zig @@ -121,7 +121,7 @@ pub const Served = struct { } // draft trees (TF_DSV41_TREE / _DUP_MS / _ROWS; off by default): first-position siblings const trees = if (dspark_on) try tree.Settings.fromEnv() else tree.Settings{}; - s.x = try ln.Lanes.init(gpa, tgt, pass, s.costs, .{ .shape = backend.shape, .slots = backend.slots, .siblings = trees.siblings, .sib_rows = trees.sib_rows, .dup = trees.dup, .depth = try calib_env.depthSettings(&calib_env.env), .spec = if (dspark_on) try spec.Settings.fromEnv() else .{}, .plain_streams_default = 4 }); + s.x = try ln.Lanes.init(gpa, tgt, pass, s.costs, .{ .shape = backend.shape, .slots = backend.slots, .siblings = trees.siblings, .sib_rows = trees.sib_rows, .dup = trees.dup, .depth = try calib_env.depthSettings(&calib_env.env), .spec = if (dspark_on) try spec.Settings.fromEnv() else .{}, .plain_streams_default = 3 }); s.arena = std.heap.ArenaAllocator.init(gpa); s.cfg = try lanes.Config.init(gpa, try s.x.model(s.arena.allocator(), 64), 16, 15); s.x.policy.depth.forward_rows = backend.forward_rows; From f470802170c5a6d5eaac38f53de087d2f00342ff Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 06:47:04 +0100 Subject: [PATCH 022/186] dsv41 metal: plain rounds skip the draft pass, and do not launch one for the next A shared round with the plain switch at it asked no drafts and ran no head pass (core Rule.plain_streams, threaded from the family's resolved setting: 0 on CUDA, so its rounds are unchanged), and the family's speculation no longer launches a pass for a round that wide (armSpec/speculate plainRound). --- zig/src/core/lanes/config.zig | 3 + zig/src/core/lanes/depth.zig | 9 ++- zig/src/core/lanes/engine.zig | 2 +- zig/src/families/deepseek_v41/draft/lanes.zig | 9 ++- .../deepseek_v41/draft/plain_streams_test.zig | 79 ++++++++++++++++--- 5 files changed, 88 insertions(+), 14 deletions(-) diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 3d4470c510f..5dd0a1fe0ee 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -46,6 +46,7 @@ pub const Model = struct { granted_depth: bool = false, // with draft_probabilities: every stream drafts its last grant plus a slack draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own) draft_streams: bool = false, // the head drafts every stream of a shared round in one batch + plain_streams: u32 = 0, // a shared round with this many streams or more drafts none (0: off) copy_first: bool = false, // a copy ahead takes the round: the head holds no drafts for it (a drafter's own pass) hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // a drafter's own hold pass, timed at load head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them @@ -80,6 +81,7 @@ pub const Config = struct { batch_streams: u32, node_probabilities: bool, granted_depth: bool = false, // Model.granted_depth + plain_streams: u32 = 0, // Model.plain_streams fill_lanes: u32 = 0, // rows a window may grow to with suffix-match branches (0: off; lanes/fill.zig) fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases) fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round @@ -167,6 +169,7 @@ pub const Config = struct { .batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1, .node_probabilities = m.draft_probabilities or m.draft_confidence != null, .granted_depth = m.granted_depth and m.draft_probabilities, + .plain_streams = m.plain_streams, .draft_streams = m.draft_streams, .head_trees = m.head_trees, .join_tail = m.join_tail, diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index fd65b614991..7a7ecaa7c7d 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -4,9 +4,7 @@ const Allocator = std.mem.Allocator; const Table = @import("table.zig").Table; const alloc = @import("allocate.zig"); pub const pricedBatched = alloc.pricedBatched; - pub const Constants = @import("depth_constants.zig").Constants; - pub const HoldCost = @import("hold_price.zig").HoldCost; const holdMs = @import("hold_price.zig").holdMs; @@ -59,8 +57,9 @@ pub const Rule = struct { node_probabilities: bool, granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too confidence_floor: ?f64 = null, + plain_streams: u32 = 0, // a shared round with this many streams or more drafts none (0: off; a family's switch) batch_rows: u32, - batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's) + batched_heads: bool = false, // one batch drafts every stream of a shared round (steps: the deepest stream's) row_cost: @import("row_cost.zig").RowCost = .{}, // what a row costs past shared_costs here, under k.learn_rows round_ms: Table = .{}, // a round's wall time by drafts, every stream's: what sampled streams price by overhead: std.ArrayList(Overhead) = .empty, // a shared round's ms past its forward, by streams (insertion order) @@ -298,6 +297,10 @@ pub const Rule = struct { /// Head drafts for the streams of a shared round, sized by allocation over their chain chances. pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void { + if (r.plain_streams > 0 and whos.len >= r.plain_streams) { // drafts lose here: none asked, no head pass + @memset(out, 0); + return; + } if (r.confidence_floor != null) { for (whos, out) |*w, *o| o.* = try r.headDepth(w, null, probe); return; diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 018ae4a72eb..ced3f315aeb 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -53,7 +53,6 @@ pub const Engine = struct { grammar_workers: @import("grammar_parallel.zig").State = .{}, arena: std.heap.ArenaAllocator, // a round's temporaries unread: std.ArrayList(Unread) = .empty, // queued draws the log reads at the step's end - pub fn init(gpa: Allocator, cfg: *const Config, backend: be.Backend, clock: be.Clock) Engine { return .{ .gpa = gpa, @@ -74,6 +73,7 @@ pub const Engine = struct { .node_probabilities = cfg.node_probabilities, .granted_depth = cfg.granted_depth, .confidence_floor = cfg.draft_confidence, + .plain_streams = cfg.plain_streams, .batch_rows = cfg.batch_rows, .batched_heads = cfg.draft_streams, }, diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index df344d2b76f..0c63b82bace 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -120,7 +120,7 @@ pub const Lanes = struct { for (window, 0..) |*w, i| w.* = .{ .width = @intCast(i + 1), .ms = v[i] }; const shared = try arena.alloc(lanes.config.Cost, @min(v.len, batch_rows)); for (shared, 0..) |*w, i| w.* = .{ .width = @intCast(i + 1), .ms = v[i] }; - return .{ .exact_width = costs_mod.slot_rows, .mtp = true, .speculate = true, .speculate_early = false, .drafts = x.opts.shape.block, .window_costs = window, .mtp_step_ms = 0.0, .hidden_rows = true, .batch_rows = batch_rows, .max_streams = x.opts.slots, .shared_costs = shared, .draft_probabilities = true, .draft_streams = true, .join_tail = x.target.vtable.tail != null and try joinTail(), .piece_runs = try envOn("TF_DSV41_PIECE_RUNS"), .join_drafts = try envOn("TF_DSV41_JOIN_DRAFTS"), .replay_runs = try envOn("TF_DSV41_REPLAY_RUNS") }; + return .{ .exact_width = costs_mod.slot_rows, .mtp = true, .speculate = true, .speculate_early = false, .drafts = x.opts.shape.block, .window_costs = window, .mtp_step_ms = 0.0, .hidden_rows = true, .batch_rows = batch_rows, .max_streams = x.opts.slots, .shared_costs = shared, .draft_probabilities = true, .draft_streams = true, .plain_streams = @intCast(x.opts.depth.plain_streams), .join_tail = x.target.vtable.tail != null and try joinTail(), .piece_runs = try envOn("TF_DSV41_PIECE_RUNS"), .join_drafts = try envOn("TF_DSV41_JOIN_DRAFTS"), .replay_runs = try envOn("TF_DSV41_REPLAY_RUNS") }; } /// TF_DSV41_TAIL_JOIN=1: a prompt in pieces leaves its last token to its first round's shared window (default off). @@ -362,6 +362,7 @@ pub const Lanes = struct { nucleus = nucleus or spec_mod.isNucleus(w.stream.sampling); } if (!sp.roundAllowed(tree_round, nucleus)) return; + if (x.plainRound(windows)) return; // the plain switch: no pass runs for a round this wide for (windows, segs) |w, g| { const st = x.get(w.stream) catch continue; if (!dsparkRound(st, w) or !sp.wouldWant(st.slot)) continue; @@ -372,6 +373,11 @@ pub const Lanes = struct { } } + /// The plain switch fired: every drafted window of the round was cut to the pending row alone. + fn plainRound(x: *Lanes, windows: []const be.Window) bool { + return x.opts.depth.plain_streams > 0 and windows.len >= x.opts.depth.plain_streams; + } + /// A DSpark round: the window verifies the pass's drafts only (a prefix of the held chain counts too). fn dsparkRound(st: *const St, w: be.Window) bool { if (st.n == 0) return false; @@ -394,6 +400,7 @@ pub const Lanes = struct { nucleus = nucleus or spec_mod.isNucleus(w.stream.sampling); } if (!sp.roundAllowed(tree_round, nucleus)) return sp.clear(); + if (x.plainRound(windows)) return sp.clear(); // the plain switch: this round drafted none, the next will not want the drafts const cands = try a.alloc(spec_mod.Candidate, windows.len); var n: usize = 0; for (windows, segs, choices) |w, g, c| { diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 44abbbfa8dd..e5e2f4cfec2 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -27,7 +27,7 @@ test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, bad values error" { Env.pairs = &.{}; } -fn runPlain(cases: []const []const u32, setup: Setup) !Run { +fn runPlain(cases: []const []const u32, setup: Setup, max_news: []const u32) !Run { const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; // a confident drafter (lanes_test's): drafts are chosen const t = try twin.Twin.init(gpa, dims, cases.len + 1); defer t.deinit(); @@ -35,7 +35,7 @@ fn runPlain(cases: []const []const u32, setup: Setup) !Run { defer d.deinit(); var c = try costs_mod.defaults(gpa, 64); defer c.deinit(gpa); - const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = setup }); + const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = setup, .spec = .{ .on = true } }); defer x.deinit(); var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -47,38 +47,99 @@ fn runPlain(cases: []const []const u32, setup: Setup) !Run { x.attach(&e); const streams = try gpa.alloc(lanes.Stream, cases.len); defer gpa.free(streams); - for (cases, streams) |p, *s| s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = 40 }); + for (cases, streams, max_news) |p, *s, *n| s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = n.* }); defer for (streams) |*s| s.deinit(gpa); for (streams) |*s| try e.addStream(s); while (e.activeCount() > 0) try e.step(); var ok = true; for (streams) |*s| ok = ok and s.emitted().len > 0; - return .{ .drafted = e.drafted, .ok = ok }; + return .{ .drafted = e.drafted, .ok = ok, .passes = d.passes, .out = blk: { + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + break :blk out; + } }; +} + +fn freeRun(r: *Run) void { + for (r.out) |o| gpa.free(o); + gpa.free(r.out); } const Setup = depth.Settings; -const Run = struct { drafted: u64, ok: bool }; +const Run = struct { drafted: u64, ok: bool, passes: u64, out: [][]u32 }; const a1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; test "a round with plain_streams or more active windows decodes plain; below it drafts; 0 keeps drafts" { // identical prompts and greedy sampling keep every round a full 4-window round until the streams end together const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; + const fours = [_]u32{ 40, 40, 40, 40 }; // 4 active, threshold 3: every drafted window is forced to count 0 - const plain = try runPlain(&four, .{ .plain_streams = 3 }); + const plain = try runPlain(&four, .{ .plain_streams = 3 }, &fours); + defer freeRun(@constCast(&plain)); try std.testing.expect(plain.ok); try std.testing.expectEqual(@as(u64, 0), plain.drafted); // 3 active, threshold 3: at it, plain too - const at = try runPlain(four[0..3], .{ .plain_streams = 3 }); + const at = try runPlain(four[0..3], .{ .plain_streams = 3 }, fours[0..3]); + defer freeRun(@constCast(&at)); try std.testing.expect(at.ok); try std.testing.expectEqual(@as(u64, 0), at.drafted); // 2 active, threshold 3: below it, drafts are kept - const below = try runPlain(four[0..2], .{ .plain_streams = 3 }); + const below = try runPlain(four[0..2], .{ .plain_streams = 3 }, fours[0..2]); + defer freeRun(@constCast(&below)); try std.testing.expect(below.ok); try std.testing.expect(below.drafted > 0); // 4 active, threshold 0 (off): drafts are kept - const off = try runPlain(&four, .{}); + const off = try runPlain(&four, .{}, &fours); + defer freeRun(@constCast(&off)); try std.testing.expect(off.ok); try std.testing.expect(off.drafted > 0); } + +test "plain rounds skip the draft pass; the replies stay the serial ones; drafting resumes below the threshold" { + const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; + const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; + const fours = [_]u32{ 40, 40, 40, 40 }; + // off, 4 streams: every round drafts, every round runs the pass + const off = try runPlain(&four, .{}, &fours); + defer freeRun(@constCast(&off)); + try std.testing.expect(off.passes > 0); + // at the threshold: shared rounds run no pass (the engine asks depth 0); the four opening asks stay + const on = try runPlain(&four, .{ .plain_streams = 3 }, &fours); + defer freeRun(@constCast(&on)); + try std.testing.expectEqual(@as(u64, 0), on.drafted); + try std.testing.expectEqual(@as(u64, 4), on.passes); + // tokens identical drafted or plain + for (off.out, on.out) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // serial reference + { + const t = try twin.Twin.init(gpa, dims, 1); + defer t.deinit(); + for (on.out, four) |got, p| { + const want = try t.serial(0, p, 40, null, null); + defer gpa.free(want); + try std.testing.expectEqualSlices(u32, want, got); + } + } + // three streams, the first stops early: its last rounds run with 2 active, drafts resume (a pass again) + const threes = [_]u32{ 6, 40, 40 }; + const tail = [_][]const u32{ &a1, &a1, &a1 }; + const mixed = try runPlain(&tail, .{ .plain_streams = 3 }, &threes); + defer freeRun(@constCast(&mixed)); + try std.testing.expect(mixed.ok); + try std.testing.expect(mixed.drafted > 0); // after a stream finished, 2 active: drafts asked again + try std.testing.expect(mixed.passes > 4); // the three opening asks plus the resumed rounds + // stream 0 stops at 6: its 6 tokens must equal the off run's opening prefix; the others match in full + try std.testing.expectEqualSlices(u32, off.out[0][0..6], mixed.out[0]); + for (off.out[1..3], mixed.out[1..]) |a, b| try std.testing.expectEqualSlices(u32, a, b); + { + const t = try twin.Twin.init(gpa, dims, 1); + defer t.deinit(); + for (mixed.out, tail, &threes) |got, p, *n| { + const want = try t.serial(0, p, n.*, null, null); + defer gpa.free(want); + try std.testing.expectEqualSlices(u32, want, got); + } + } +} From 8317e6bd9e2d511fb4621a7c6835810904146551 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 06:47:30 +0100 Subject: [PATCH 023/186] dsv41 metal: the plain switch line says three streams and the skipped pass --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 022096db823..f7dbd790269 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -83,6 +83,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. - DeepSeek-V4.1 on Metal: rounds with four or more active streams decode plain instead of drafting, since drafts lose there (#1720, part of #1676). +- DeepSeek-V4.1 on Metal: rounds with three or more active streams decode plain and skip the draft pass, since drafts lose there (#1720, part of #1676). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. From cd3115613d2c5466aa563a6b229789d95b16ef42 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 10:31:19 +0100 Subject: [PATCH 024/186] dsv41 metal: a stream's first draft ask waits for its first round --- zig/src/families/deepseek_v41/draft/lanes.zig | 6 ++- .../deepseek_v41/draft/plain_streams_test.zig | 40 +++++++++++++------ 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 0c63b82bace..8ed2270ff9e 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -43,6 +43,7 @@ pub const St = struct { n_sib: u32 = 0, shown: std.ArrayList(u32) = .empty, // the drafts this round's window verifies (node indices, window order) shown_set: bool = false, + asked: bool = false, // a draft ask is made (the plain switch: the first ask stays plain) /// Node i of the held tree: the main chain's drafts, then each sibling's rows. pub fn nodes(st: *const St) u32 { @@ -466,7 +467,10 @@ pub const Lanes = struct { st.n = 0; st.n_sib = 0; st.start = at; + const first_ask = !st.asked; + st.asked = true; if (r.depth == 0) continue; + if (first_ask and x.opts.depth.plain_streams > 0) continue; // the plain switch: the stream's first round runs plain, the next ask drafts if (n == x.opts.slots) return error.TooManySlots; asks[n] = .{ .slot = st.slot, .anchor = anchor, .start = at, .params = dspark.Params.of(r.stream.sampling, at) }; whose[n] = st; @@ -526,7 +530,7 @@ pub const Lanes = struct { fn heldTree(ptr: *anyopaque, s: *Stream, gpa: Allocator) anyerror!?lanes.stream.Held { const x = self(ptr); const st = try x.get(s); - if (st.n == 0) return null; + if (st.n == 0) return .{ .count = 0 }; // the ask was skipped (the plain switch's first round): the engine plans a plain window const total = st.nodes(); const tokens = try gpa.alloc(u32, total); errdefer gpa.free(tokens); diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index e5e2f4cfec2..47ca916ca92 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -97,7 +97,7 @@ test "a round with plain_streams or more active windows decodes plain; below it try std.testing.expect(off.drafted > 0); } -test "plain rounds skip the draft pass; the replies stay the serial ones; drafting resumes below the threshold" { +test "a stream's first draft ask waits for its first round; the replies stay the serial ones" { const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; const fours = [_]u32{ 40, 40, 40, 40 }; @@ -105,11 +105,11 @@ test "plain rounds skip the draft pass; the replies stay the serial ones; drafti const off = try runPlain(&four, .{}, &fours); defer freeRun(@constCast(&off)); try std.testing.expect(off.passes > 0); - // at the threshold: shared rounds run no pass (the engine asks depth 0); the four opening asks stay + // at the threshold: shared rounds run no pass and the opening asks are skipped, so no pass runs at all const on = try runPlain(&four, .{ .plain_streams = 3 }, &fours); defer freeRun(@constCast(&on)); try std.testing.expectEqual(@as(u64, 0), on.drafted); - try std.testing.expectEqual(@as(u64, 4), on.passes); + try std.testing.expectEqual(@as(u64, 0), on.passes); // tokens identical drafted or plain for (off.out, on.out) |a, b| try std.testing.expectEqualSlices(u32, a, b); // serial reference @@ -122,21 +122,35 @@ test "plain rounds skip the draft pass; the replies stay the serial ones; drafti try std.testing.expectEqualSlices(u32, want, got); } } - // three streams, the first stops early: its last rounds run with 2 active, drafts resume (a pass again) - const threes = [_]u32{ 6, 40, 40 }; - const tail = [_][]const u32{ &a1, &a1, &a1 }; - const mixed = try runPlain(&tail, .{ .plain_streams = 3 }, &threes); +} + +test "below the threshold a stream drafts on its first ask; drafting resumes once streams finish" { + const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; + const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; + const fours = [_]u32{ 40, 40, 40, 40 }; + // the threshold not reached (1-2 active): the first ask drafts as usual + const one = try runPlain(four[0..1], .{ .plain_streams = 3 }, fours[0..1]); + defer freeRun(@constCast(&one)); + try std.testing.expect(one.ok); + try std.testing.expect(one.drafted > 0); + try std.testing.expect(one.passes > 0); + const two = try runPlain(four[0..2], .{ .plain_streams = 3 }, fours[0..2]); + defer freeRun(@constCast(&two)); + try std.testing.expect(two.ok); + try std.testing.expect(two.drafted > 0); + try std.testing.expect(two.passes > 0); + // 4 streams, two finish early: the first asks are skipped, drafting resumes at 2 active + const twos = [_]u32{ 6, 6, 40, 40 }; + const mixed = try runPlain(&four, .{ .plain_streams = 3 }, &twos); defer freeRun(@constCast(&mixed)); try std.testing.expect(mixed.ok); - try std.testing.expect(mixed.drafted > 0); // after a stream finished, 2 active: drafts asked again - try std.testing.expect(mixed.passes > 4); // the three opening asks plus the resumed rounds - // stream 0 stops at 6: its 6 tokens must equal the off run's opening prefix; the others match in full - try std.testing.expectEqualSlices(u32, off.out[0][0..6], mixed.out[0]); - for (off.out[1..3], mixed.out[1..]) |a, b| try std.testing.expectEqualSlices(u32, a, b); + try std.testing.expect(mixed.drafted > 0); // after two streams finished, 2 active: drafts asked again + try std.testing.expect(mixed.passes > 0); // every pass is from the resumed rounds + // serial reference, each stream's own length { const t = try twin.Twin.init(gpa, dims, 1); defer t.deinit(); - for (mixed.out, tail, &threes) |got, p, *n| { + for (mixed.out, four, &twos) |got, p, *n| { const want = try t.serial(0, p, n.*, null, null); defer gpa.free(want); try std.testing.expectEqualSlices(u32, want, got); From 873a4bc0bc6954b03d3878c1581ad2aaf5390e9c Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 19:32:12 +0100 Subject: [PATCH 025/186] dsv41 metal: a stream drafts nothing until its burst has formed --- zig/src/families/deepseek_v41/draft/lanes.zig | 18 +++-- .../deepseek_v41/draft/plain_streams_test.zig | 67 ++++++++++++++++++- 2 files changed, 77 insertions(+), 8 deletions(-) diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 8ed2270ff9e..831f937e71b 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -43,7 +43,8 @@ pub const St = struct { n_sib: u32 = 0, shown: std.ArrayList(u32) = .empty, // the drafts this round's window verifies (node indices, window order) shown_set: bool = false, - asked: bool = false, // a draft ask is made (the plain switch: the first ask stays plain) + burst_seen: bool = false, // the plain switch: an ask batch reached plain_streams, the normal rules apply from then on + wide_asks: u8 = 0, // the plain switch: asks made while never wide, a lone stream drafts from its second /// Node i of the held tree: the main chain's drafts, then each sibling's rows. pub fn nodes(st: *const St) u32 { @@ -74,7 +75,6 @@ pub const Lanes = struct { by_slot: []?*St, retired: [16]bool = @splat(false), release_requested: [16]bool = @splat(false), - pub const init = @import("lane_owner.zig").init; pub const deinit = @import("lane_owner.zig").deinit; @@ -435,11 +435,16 @@ pub const Lanes = struct { } } + /// The hold before a stream's first draft (the plain switch): until its burst has formed or its second ask. + const lone_asks_before_draft = 2; // a stream admitted alone drafts from its second ask + fn draft(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { const x = self(ptr); try x.settle(); _ = x.scratch.reset(.retain_capacity); const a = x.scratch.allocator(); + const switch_on = x.opts.depth.plain_streams > 0; + const burst = switch_on and requests.len >= x.opts.depth.plain_streams; // the round's burst: the normal rules apply const asks = try a.alloc(iface.Ask, requests.len); const whose = try a.alloc(*St, requests.len); const depths = try a.alloc(u32, requests.len); @@ -453,7 +458,7 @@ pub const Lanes = struct { if (r.rows) |path| { if (r.start != st.base) return error.PositionMismatch; const done = if (x.spec) |*sp| sp.ingested(st.slot, r.start, path) else false; - if (!done) try x.ingest(st, r.start, path); // a shared round drafts before its keep: the path, not st.len + if (!done) try x.ingest(st, r.start, path); // the drafter's ring holds every committed row, drafted or not // a shared round drafts before its keep: the path, not st.len at = r.start + path.len; } if (r.position != at + 1) return error.PositionMismatch; @@ -467,10 +472,11 @@ pub const Lanes = struct { st.n = 0; st.n_sib = 0; st.start = at; - const first_ask = !st.asked; - st.asked = true; if (r.depth == 0) continue; - if (first_ask and x.opts.depth.plain_streams > 0) continue; // the plain switch: the stream's first round runs plain, the next ask drafts + if (switch_on) { + if (burst) st.burst_seen = true else st.wide_asks += 1; + if (!st.burst_seen and st.wide_asks < lone_asks_before_draft) continue; // the plain switch: the hold until its burst has formed or its second ask + } if (n == x.opts.slots) return error.TooManySlots; asks[n] = .{ .slot = st.slot, .anchor = anchor, .start = at, .params = dspark.Params.of(r.stream.sampling, at) }; whose[n] = st; diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 47ca916ca92..3d0fa7239e7 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -124,11 +124,74 @@ test "a stream's first draft ask waits for its first round; the replies stay the } } +test "a stream drafts no round until its burst has formed or its second ask while never wide" { + // staggered arrival, plain_streams 3: the burst forms at the third stream, from which every stream drafts + const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; + { + const t = try twin.Twin.init(gpa, dims, 5); + defer t.deinit(); + const d = try twin.Drafter.init(gpa, t, 5); + defer d.deinit(); + var c = try costs_mod.defaults(gpa, 64); + defer c.deinit(gpa); + const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = .{ .plain_streams = 3 }, .spec = .{ .on = true } }); + defer x.deinit(); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + var cfg = try lanes.Config.init(gpa, try x.model(arena.allocator(), 64), 16, 15); + defer cfg.deinit(gpa); + var clock: lanes.fake.FixedClock = .{}; + var e = lanes.Engine.init(gpa, &cfg, x.backend(), clock.clock()); + defer e.deinit(); + x.attach(&e); + const streams = try gpa.alloc(lanes.Stream, 4); + defer gpa.free(streams); + var i: usize = 0; + var steps: u64 = 0; + while (e.activeCount() > 0 or i < 4) { + try e.step(); + steps += 1; + if (i < 4) { + const s = &streams[i]; + s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = &a1, .max_new = 40 }); + try e.addStream(s); + i += 1; + } + if (steps > 400) return error.TooManySteps; + } + defer for (streams) |*s| s.deinit(gpa); + var ok = true; + for (streams) |*s| ok = ok and s.emitted().len > 0; + try std.testing.expect(ok); + try std.testing.expect(d.passes > 0); // the burst formed at three streams: the later rounds draft + try std.testing.expect(e.drafted > 0); + // tokens equal to the serial reference + for (streams) |*s| { + const want = try t.serial(0, &a1, 40, null, null); + defer gpa.free(want); + try std.testing.expectEqualSlices(u32, want, s.emitted()); + } + } + // the truly-lone control: one stream admitted alone drafts from its second ask + const lone = try runPlain(&.{&a1}, .{ .plain_streams = 3 }, &.{40}); + defer freeRun(@constCast(&lone)); + try std.testing.expect(lone.ok); + try std.testing.expect(lone.drafted > 0); + try std.testing.expect(lone.passes > 0); + { + const t = try twin.Twin.init(gpa, dims, 1); + defer t.deinit(); + const want = try t.serial(0, &a1, 40, null, null); + defer gpa.free(want); + try std.testing.expectEqualSlices(u32, want, lone.out[0]); + } +} + test "below the threshold a stream drafts on its first ask; drafting resumes once streams finish" { const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; const fours = [_]u32{ 40, 40, 40, 40 }; - // the threshold not reached (1-2 active): the first ask drafts as usual + // the threshold not reached (1-2 active): the first ask is held, the second drafts const one = try runPlain(four[0..1], .{ .plain_streams = 3 }, fours[0..1]); defer freeRun(@constCast(&one)); try std.testing.expect(one.ok); @@ -139,7 +202,7 @@ test "below the threshold a stream drafts on its first ask; drafting resumes onc try std.testing.expect(two.ok); try std.testing.expect(two.drafted > 0); try std.testing.expect(two.passes > 0); - // 4 streams, two finish early: the first asks are skipped, drafting resumes at 2 active + // 4 streams, two finish early: admitted four-deep, the burst has formed from round one, so every ask drafts const twos = [_]u32{ 6, 6, 40, 40 }; const mixed = try runPlain(&four, .{ .plain_streams = 3 }, &twos); defer freeRun(@constCast(&mixed)); From 4a5f5d405b9f17560637a6eed1e8427a81cbd5c0 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 21:36:26 +0100 Subject: [PATCH 026/186] dsv41 metal: restore the engram manifest check, the serve-lane test steps and the forward try's --- zig/build/dsv41_metal.zig | 21 +++++++++++++++---- zig/src/families/deepseek_v41/metal/boot.zig | 15 +++++++++++++ .../families/deepseek_v41/metal/forward.zig | 4 ++-- zig/tests/dsv41_metal/timing_real.zig | 2 +- 4 files changed, 35 insertions(+), 7 deletions(-) diff --git a/zig/build/dsv41_metal.zig b/zig/build/dsv41_metal.zig index 3eb2207ab1a..1de16b75335 100644 --- a/zig/build/dsv41_metal.zig +++ b/zig/build/dsv41_metal.zig @@ -40,9 +40,18 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module .link_libc = true, .imports = &.{ .{ .name = "dsv41_metal", .module = p.pkg }, .{ .name = "metal", .module = metal }, .{ .name = "dsv41", .module = p.host } }, }); + // the package's own checks run with the default test step; its device suites run from their named steps + test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = p.pkg })).step); const step = b.step("test-dsv41-metal", "DeepSeek-V4.1 Metal layer and forward: float64 stages, rows at once == one at a time"); - const pkg_run = &b.addRunArtifact(b.addTest(.{ .root_module = p.pkg })).step; - step.dependOn(pkg_run); + // serve_lanes over the CPU twin, checked here so no CUDA step has to depend on the neutral served engine + const lanes_test = b.createModule(.{ + .root_source_file = b.path("zig/src/families/deepseek_v41/serve_lanes_test.zig"), + .target = b.graph.host, + .optimize = .Debug, + .link_libc = true, + .imports = &.{ .{ .name = "core", .module = p.host.import_table.get("core").? }, .{ .name = "engine_api", .module = p.host.import_table.get("engine_api").? }, .{ .name = "lanes", .module = p.host.import_table.get("lanes").? }, .{ .name = "tokenizer", .module = p.host.import_table.get("tokenizer").? } }, + }); + test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes_test })).step); const run = b.addRunArtifact(b.addTest(.{ .root_module = tests })); run.has_side_effects = true; // a device run: never cached step.dependOn(&run.step); @@ -58,8 +67,12 @@ pub fn steps(b: *std.Build, metal: *std.Build.Module, sources: *std.Build.Module const serve_run = b.addRunArtifact(b.addTest(.{ .root_module = serve })); serve_run.has_side_effects = true; step.dependOn(&serve_run.step); - b.step("test-dsv41-metal-serve", "DeepSeek-V4.1 Metal served through LaneHost: concurrent, drafted, cached, seeded, cancelled").dependOn(&serve_run.step); - test_step.dependOn(pkg_run); // the package's host checks only; the device run stays in its own step + const serve_step = b.step("test-dsv41-metal-serve", "DeepSeek-V4.1 Metal served through LaneHost: concurrent, drafted, cached, seeded, cancelled"); + serve_step.dependOn(&serve_run.step); + // the release-checkpoint steps (TF_DSV41_REAL) stay opt-in, so the umbrella holds the default device suites + const all = b.step("test-dsv41-metal-all", "DeepSeek-V4.1 Metal: every device test step on this GPU (layer, served)"); + all.dependOn(step); + all.dependOn(serve_step); const fast = package(b, metal, sources, .ReleaseSafe); const real = b.createModule(.{ .root_source_file = b.path("zig/tests/dsv41_metal/real_test.zig"), diff --git a/zig/src/families/deepseek_v41/metal/boot.zig b/zig/src/families/deepseek_v41/metal/boot.zig index 157552c6d26..88e3d53e2b7 100644 --- a/zig/src/families/deepseek_v41/metal/boot.zig +++ b/zig/src/families/deepseek_v41/metal/boot.zig @@ -106,6 +106,21 @@ pub const Boot = struct { b.eng = try eh.Host.load(b.map_bytes, c.engram_heads, c.engram_pad, c.engram_vocab, &b.lbuf); const rows_dir = try b.rowsDir(dir, o); defer b.a.free(rows_dir); + // a boot that serves Engram layers proves the tables match the release: the manifest, headers, sampled rows + const full = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_DSV41_ENGRAM_VERIFY") orelse ""), "full"); + const em = host.engram_manifest; + const ck = em.check(b.a, io, dir, rows_dir, full) catch |e| { + switch (e) { + error.EngramManifestMissing => std.log.err("dsv41 metal: no {s} beside the Engram tables {s}: run tf-dsv41-engram-pack pack", .{ em.filename, rows_dir }), + error.EngramManifestBad => std.log.err("dsv41 metal: the Engram manifest {s}/{s} is missing or unreadable", .{ rows_dir, em.filename }), + error.EngramSourceMismatch => std.log.err("dsv41 metal: the release no longer matches the Engram tables in {s}: the shards moved or changed", .{rows_dir}), + error.EngramShardCorrupt => std.log.err("dsv41 metal: an Engram table in {s} is corrupt or does not match its manifest", .{rows_dir}), + error.EngramSampleMismatch => std.log.err("dsv41 metal: an Engram table in {s} differs from the release shards on a sampled row", .{rows_dir}), + else => std.log.err("dsv41 metal: the Engram tables in {s} cannot be verified ({s})", .{ rows_dir, @errorName(e) }), + } + return e; + }; + std.log.info("dsv41 metal: Engram tables checked against the release ({d} sampled rows{s})", .{ ck.samples, if (full) " and every byte" else "" }); const cap: u32 = if (std.c.getenv("TF_DSV41_ENGRAM_CACHE")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else 65536; b.rows = try er.Rows.open(b.a, rows_dir, 0, 1, b.eng.layer_ids, cap); std.log.info("dsv41 metal: Engram rows from {s} ({d} records a layer cached)", .{ rows_dir, cap }); diff --git a/zig/src/families/deepseek_v41/metal/forward.zig b/zig/src/families/deepseek_v41/metal/forward.zig index c6022aa34cb..622ec616b42 100644 --- a/zig/src/families/deepseek_v41/metal/forward.zig +++ b/zig/src/families/deepseek_v41/metal/forward.zig @@ -245,11 +245,11 @@ const max_buffers = 16; /// Layers lo .. hi of the forward, the embedding before layer 0 and the head after the last. fn encodeRange(x: *const Ctx, e: mtl.ComputeEncoder, m: Model, b: *const Bufs, lo: usize, hi: usize) !void { if (lo == 0) { - encodeEmbed(x, e, m, b); + try encodeEmbed(x, e, m, b); e.barrier(); } for (lo..hi) |i| try encodeLayer(x, e, m, b, i); - if (hi == m.layers.len) encodeHead(x, e, m, b); + if (hi == m.layers.len) try encodeHead(x, e, m, b); } test "a round's command buffers: one layer, then four even parts of the rest" { diff --git a/zig/tests/dsv41_metal/timing_real.zig b/zig/tests/dsv41_metal/timing_real.zig index 2243c25f6f7..5e41dbd0e09 100644 --- a/zig/tests/dsv41_metal/timing_real.zig +++ b/zig/tests/dsv41_metal/timing_real.zig @@ -340,7 +340,7 @@ fn rest(a: std.mem.Allocator, r: *fixture.Rig, junk: mtl.Buffer, n: usize) !void const Args = struct { r: *fixture.Rig, rows: u32, head: mtl.Buffer, logits: mtl.Buffer, ew: mtl.Buffer, es: mtl.Buffer, ex: mtl.Buffer, ekv: mtl.Buffer, k: u32, n: u32 }; const F = struct { fn headMv(x: Args, e: mtl.ComputeEncoder) !void { - pkg.pipelines.glue.mv(&x.r.p.glue, e, x.r.s.hn, .{ .buf = x.head }, .{ .buf = x.logits }, x.rows, x.r.c.hidden, x.r.c.vocab, true); + try pkg.pipelines.glue.mv(&x.r.p.glue, e, x.r.s.hn, .{ .buf = x.head }, .{ .buf = x.logits }, x.rows, x.r.c.hidden, x.r.c.vocab, true); } fn wkv(x: Args, e: mtl.ComputeEncoder) !void { const fp8 = pkg.pipelines.fp8; From c713807eb88821ff1a00b7d745053829e9fd5c5e Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:16:45 +0100 Subject: [PATCH 027/186] lanes: the plain-streams check in depth.zig on one line, so the file stays within 600 lines --- zig/src/core/lanes/depth.zig | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig index 7a7ecaa7c7d..258d86b170e 100644 --- a/zig/src/core/lanes/depth.zig +++ b/zig/src/core/lanes/depth.zig @@ -297,10 +297,8 @@ pub const Rule = struct { /// Head drafts for the streams of a shared round, sized by allocation over their chain chances. pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void { - if (r.plain_streams > 0 and whos.len >= r.plain_streams) { // drafts lose here: none asked, no head pass - @memset(out, 0); - return; - } + // drafts lose here: none asked, no head pass + if (r.plain_streams > 0 and whos.len >= r.plain_streams) return @memset(out, 0); if (r.confidence_floor != null) { for (whos, out) |*w, *o| o.* = try r.headDepth(w, null, probe); return; From dcd2e9f485fbb278d745dd9b5d5bc446261c1d3f Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:16:45 +0100 Subject: [PATCH 028/186] dsv41 metal: boot and entry read their settings through platform/environment.zig --- docs/environment.md | 5 +++++ zig/src/families/deepseek_v41/metal/boot.zig | 9 +++++---- zig/src/families/deepseek_v41/metal/entry.zig | 5 +++-- zig/src/platform/environment.zig | 5 +++++ zig/src/platform/environment_test.zig | 5 +++++ 5 files changed, 23 insertions(+), 6 deletions(-) diff --git a/docs/environment.md b/docs/environment.md index 81972702dd9..5b937ab5861 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -47,6 +47,11 @@ Diagnostics quote and escape the supplied text so it cannot create another termi | `TF_FLASHNEXT_SCALE_INDEX` | Flash Next on Metal before tensor units: exactly `0` or `1`; `1` has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs, the same values with fewer weight bytes; default 0 | | `TF_FLASHNEXT_GROUPED_ROWS` | Flash Next on Metal without tensor units: `0`, or the rows (`1` through `32`) from which a decode window runs its 4-bit experts grouped by expert, with the same output bits; default 5; ignored with `--ssd-experts` and in speed-up mode | | `GLM_EHPROJ_NAX` | GLM-5.3-Flash on Metal: exactly `0` or `1`; `0` serves the MTP head's prompt eh_proj on the scalar kernel; default 1, the tile kernel where it builds | +| `TF_DSV41_SMALL` | DeepSeek-V4.1 on Metal: exactly `0` or `1`; `1` reads the checkpoint's config at its own small shapes (the served smoke test's tiny model); default 0 | +| `TF_DSV41_ASSETS` | DeepSeek-V4.1 on Metal: a directory path without control characters that holds the Engram token map; unset or blank reads it beside the checkpoint | +| `TF_DSV41_ENGRAM_DIR` | DeepSeek-V4.1 on Metal: a directory path without control characters that holds the packed Engram tables; unset or blank takes the checkpoint's `engram_table_dir` | +| `TF_DSV41_ENGRAM_VERIFY` | DeepSeek-V4.1 on Metal: exactly `sample` or `full`; `full` also checks every byte of the Engram tables against their manifest at boot; default `sample` | +| `TF_DSV41_ENGRAM_CACHE` | DeepSeek-V4.1 on Metal: whole u32 count of Engram records a layer kept cached; default 65536 | | `XPU_A8_BITS` | Exactly `4` or `8`, for the XPU quantization study; unset keeps the study's existing setting | | `LFM3_FUSE` | Whole u32 fusion mask; default 7 | | `TF_SLIDE_POCKET_KEY` | Exactly `auto`, `reply`, `subject` or `both`; default `auto` | diff --git a/zig/src/families/deepseek_v41/metal/boot.zig b/zig/src/families/deepseek_v41/metal/boot.zig index 88e3d53e2b7..4c21cdf5c34 100644 --- a/zig/src/families/deepseek_v41/metal/boot.zig +++ b/zig/src/families/deepseek_v41/metal/boot.zig @@ -14,6 +14,7 @@ const target_ = @import("target.zig"); const Config = host.config.Config; const eh = host.engram_host; const er = host.engram_rows; +const environment = @import("sys").environment; const Allocator = std.mem.Allocator; /// The files a boot reads beside the shards: the Engram token map (here or under TF_DSV41_ASSETS). @@ -95,7 +96,7 @@ pub const Boot = struct { fn engram(b: *Boot, io: std.Io, dir: []const u8, o: Options) !?forward.Lookup { const c = &b.c; if (c.engram_layers.len == 0) return null; - const assets = if (std.c.getenv("TF_DSV41_ASSETS")) |v| std.mem.span(v) else dir; + const assets = try environment.path("TF_DSV41_ASSETS", environment.Process.get(.{}, "TF_DSV41_ASSETS")) orelse dir; const path = try std.fs.path.join(b.a, &.{ assets, token_map }); defer b.a.free(path); b.map_bytes = std.Io.Dir.cwd().readFileAllocOptions(io, path, b.a, .limited(1 << 26), .@"4", null) catch |e| { @@ -107,7 +108,7 @@ pub const Boot = struct { const rows_dir = try b.rowsDir(dir, o); defer b.a.free(rows_dir); // a boot that serves Engram layers proves the tables match the release: the manifest, headers, sampled rows - const full = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_DSV41_ENGRAM_VERIFY") orelse ""), "full"); + const full = try environment.read(u32, "TF_DSV41_ENGRAM_VERIFY", 0) == 1; const em = host.engram_manifest; const ck = em.check(b.a, io, dir, rows_dir, full) catch |e| { switch (e) { @@ -121,7 +122,7 @@ pub const Boot = struct { return e; }; std.log.info("dsv41 metal: Engram tables checked against the release ({d} sampled rows{s})", .{ ck.samples, if (full) " and every byte" else "" }); - const cap: u32 = if (std.c.getenv("TF_DSV41_ENGRAM_CACHE")) |v| try std.fmt.parseInt(u32, std.mem.span(v), 10) else 65536; + const cap = try environment.read(u32, "TF_DSV41_ENGRAM_CACHE", 65536); b.rows = try er.Rows.open(b.a, rows_dir, 0, 1, b.eng.layer_ids, cap); std.log.info("dsv41 metal: Engram rows from {s} ({d} records a layer cached)", .{ rows_dir, cap }); return .{ .a = b.a, .host = &b.eng, .source = b.rows.?.source() }; @@ -130,7 +131,7 @@ pub const Boot = struct { /// The row shards' folder: the option, TF_DSV41_ENGRAM_DIR, else engram_table_dir (relative to the checkpoint). fn rowsDir(b: *Boot, dir: []const u8, o: Options) ![]u8 { if (o.engram_dir) |d| return b.a.dupe(u8, d); - if (std.c.getenv("TF_DSV41_ENGRAM_DIR")) |v| return b.a.dupe(u8, std.mem.span(v)); + if (try environment.path("TF_DSV41_ENGRAM_DIR", environment.Process.get(.{}, "TF_DSV41_ENGRAM_DIR"))) |v| return b.a.dupe(u8, v); const t = b.c.engram_table_dir.items(); if (t.len == 0) return error.EngramRowsUnset; if (std.fs.path.isAbsolute(t)) return b.a.dupe(u8, t); diff --git a/zig/src/families/deepseek_v41/metal/entry.zig b/zig/src/families/deepseek_v41/metal/entry.zig index 7eb864254b5..9efb4f120e8 100644 --- a/zig/src/families/deepseek_v41/metal/entry.zig +++ b/zig/src/families/deepseek_v41/metal/entry.zig @@ -5,6 +5,7 @@ const pkg = @import("dsv41_metal"); const host = @import("dsv41"); const serve_lanes = host.serve_lanes; +const env = @import("sys").environment; const Allocator = std.mem.Allocator; pub const backends: []const []const u8 = &.{"metal"}; @@ -16,7 +17,7 @@ pub fn chip(_: Allocator) ?[]const u8 { /// The engine for `o.dir`, or null with `problem` set (the native_engines contract). pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { - if (@import("sys").experimental.refusal(o.model_type, @import("sys").getenv("TF_EXPERIMENTAL")) != null) { + if (@import("sys").experimental.refusal(o.model_type, env.Process.get(.{}, "TF_EXPERIMENTAL")) != null) { problem.* = "DeepSeek-V4.1 on Metal is experimental: its served lanes have no whole-model gate yet; set TF_EXPERIMENTAL=deepseek to opt in"; return null; } @@ -29,7 +30,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c .context = if (o.context) |c| @intCast(@max(c, 1)) else null, .drafts = o.drafts, .cache_bytes = if (o.prompt_cache_gib) |g| @intFromFloat(@max(g, 0) * (1 << 30)) else 4 << 30, - .small = if (std.c.getenv("TF_DSV41_SMALL")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false, + .small = try env.read(bool, "TF_DSV41_SMALL", false), }; const m = pkg.boot.Boot.open(gpa, io, o.dir, opts) catch |e| { problem.* = try std.fmt.allocPrint(a, "the DeepSeek-V4.1 Metal engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); diff --git a/zig/src/platform/environment.zig b/zig/src/platform/environment.zig index bf0c3efaad5..884df7d5e35 100644 --- a/zig/src/platform/environment.zig +++ b/zig/src/platform/environment.zig @@ -61,6 +61,11 @@ pub const rules = [_]Rule{ .{ .name = "TF_FLASHNEXT_SCALE_INDEX", .scope = .metal, .kind = .flag, .forms = "0 or 1 (1: Flash Next's 4-bit expert kernels read scales and biases through a 16-bit index)" }, .{ .name = "GLM_EHPROJ_NAX", .scope = .metal, .kind = .flag, .forms = "0 or 1 (0 serves GLM's MTP eh_proj on the scalar kernel)" }, .{ .name = "TF_FLASHNEXT_GROUPED_ROWS", .scope = .metal, .kind = .size, .forms = "0 (never) or the rows, 1 to 32, from which a window's experts run grouped", .max = 32 }, + .{ .name = "TF_DSV41_SMALL", .scope = .metal, .kind = .flag, .forms = "0 or 1 (1 reads a DeepSeek-V4.1 Metal checkpoint's config at its own small shapes)" }, + .{ .name = "TF_DSV41_ASSETS", .scope = .metal, .kind = .path, .forms = "a directory path without control characters holding DeepSeek-V4.1's Engram token map" }, + .{ .name = "TF_DSV41_ENGRAM_DIR", .scope = .metal, .kind = .path, .forms = "a directory path without control characters holding DeepSeek-V4.1's packed Engram tables" }, + .{ .name = "TF_DSV41_ENGRAM_VERIFY", .scope = .metal, .kind = .word, .words = &.{ "sample", "full" }, .forms = "sample or full (full also checks every byte of the Engram tables at boot)" }, + .{ .name = "TF_DSV41_ENGRAM_CACHE", .scope = .metal, .kind = .u32, .forms = "a whole Engram record count a layer from 0 to 4294967295 (default 65536)" }, .{ .name = "XPU_A8_BITS", .scope = .xpu, .kind = .bit_width, .forms = "4 or 8" }, .{ .name = "LFM3_FUSE", .scope = .xpu, .kind = .u32, .forms = "a whole fusion mask from 0 to 4294967295" }, .{ .name = "TF_SLIDE_POCKET_KEY", .scope = .slide, .kind = .word, .words = &.{ "auto", "reply", "subject", "both" }, .forms = "auto, reply, subject or both" }, diff --git a/zig/src/platform/environment_test.zig b/zig/src/platform/environment_test.zig index 9baa5497c6a..04093602fc8 100644 --- a/zig/src/platform/environment_test.zig +++ b/zig/src/platform/environment_test.zig @@ -100,6 +100,11 @@ const valid = [_]Case{ .{ .name = "TF_SLIDE_POCKET_DUMP", .raw = "/tmp/pocket rows", .n = 0 }, .{ .name = "TENSORFOLD_CUDA_KERNELS", .raw = "/tmp/captured set", .n = 0 }, .{ .name = "TF_SLIDE_FORGET_FAIL", .raw = " 0 ", .n = 0 }, + .{ .name = "TF_DSV41_SMALL", .raw = "1", .n = 1 }, + .{ .name = "TF_DSV41_ASSETS", .raw = "deepseek assets", .n = 0 }, + .{ .name = "TF_DSV41_ENGRAM_DIR", .raw = "engram-rows", .n = 0 }, + .{ .name = "TF_DSV41_ENGRAM_VERIFY", .raw = "full", .n = 1 }, + .{ .name = "TF_DSV41_ENGRAM_CACHE", .raw = "0", .n = 0 }, }; test "environment: valid forms remain raw numbers before existing consumer clamps and defaults" { From 3ee1eabf41019f77c3d73591fe1d7760c97c72ce Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:43:49 +0100 Subject: [PATCH 029/186] dsv41 metal: TF_DSV41_DRAFT_PLAIN_STREAMS takes 0 to 16 and refuses anything else by name --- zig/src/families/deepseek_v41/draft/calib_env.zig | 8 +++++++- .../families/deepseek_v41/draft/plain_streams_test.zig | 8 +++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/zig/src/families/deepseek_v41/draft/calib_env.zig b/zig/src/families/deepseek_v41/draft/calib_env.zig index 187987d3ea8..525d0a041b4 100644 --- a/zig/src/families/deepseek_v41/draft/calib_env.zig +++ b/zig/src/families/deepseek_v41/draft/calib_env.zig @@ -206,6 +206,9 @@ pub fn store(gpa: Allocator, io: std.Io, dir: []const u8, want: Shape, c: Costs, return path; } +/// TF_DSV41_DRAFT_PLAIN_STREAMS' largest value: the most slots a server runs (TF_DSV41_SLOTS). +pub const max_plain_streams = 16; + /// The depth's knobs (TF_DSV41_DEPTH, _DRAFT_DEPTH, _DRAFT_SKIP, _DEPTH_JOINT, _DEPTH_FAIR, _DEPTH_HOST_MS, rows cap). pub fn depthSettings(get: *const fn ([*:0]const u8) ?[]const u8) !depth.Settings { var s: depth.Settings = .{}; @@ -225,7 +228,10 @@ pub fn depthSettings(get: *const fn ([*:0]const u8) ?[]const u8) !depth.Settings s.fair = try float(get("TF_DSV41_DEPTH_FAIR"), joint.fair, 0.0, 4.0); s.host_ms = try float(get("TF_DSV41_DEPTH_HOST_MS"), joint.host_ms, 0.0, 100.0); if (trimmed(get("TF_DSV41_GRAPH_ROWS_MAX"))) |v| s.max_rows = @max(1, @min(try std.fmt.parseInt(u64, v, 10), costs_mod.max_rows)); - if (trimmed(get("TF_DSV41_DRAFT_PLAIN_STREAMS"))) |v| s.plain_streams = try std.fmt.parseInt(usize, v, 10); + if (trimmed(get("TF_DSV41_DRAFT_PLAIN_STREAMS"))) |v| { + s.plain_streams = std.fmt.parseInt(usize, v, 10) catch return error.DraftPlainStreams; + if (s.plain_streams > max_plain_streams) return error.DraftPlainStreams; + } return s; } diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 3d0fa7239e7..64350155d6d 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -17,13 +17,15 @@ const Env = struct { } }; -test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, bad values error" { +test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, 0 to 16 taken, anything else refused by name" { Env.pairs = &.{}; try std.testing.expectEqual(@as(usize, 0), (try calib_env.depthSettings(&Env.get)).plain_streams); Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "3" }}; try std.testing.expectEqual(@as(usize, 3), (try calib_env.depthSettings(&Env.get)).plain_streams); - Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "oops" }}; - try std.testing.expectError(error.InvalidCharacter, calib_env.depthSettings(&Env.get)); + for ([_][]const u8{ "16", "17", "4294967296", "18446744073709551616", "-1", "oops", "3x" }) |raw| { + Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", raw }}; + if (std.mem.eql(u8, raw, "16")) try std.testing.expectEqual(@as(usize, 16), (try calib_env.depthSettings(&Env.get)).plain_streams) else try std.testing.expectError(error.DraftPlainStreams, calib_env.depthSettings(&Env.get)); + } Env.pairs = &.{}; } From 0c65d2363b9db2b91e6e92f2bcc91b151ea40761 Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:44:36 +0100 Subject: [PATCH 030/186] dsv41 metal: the plain switch holds a stream's ask only while plain_streams streams decode --- zig/src/families/deepseek_v41/draft/lanes.zig | 28 +- .../deepseek_v41/draft/plain_streams_test.zig | 274 +++++++----------- 2 files changed, 119 insertions(+), 183 deletions(-) diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 831f937e71b..93ce5ce3f68 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -43,8 +43,6 @@ pub const St = struct { n_sib: u32 = 0, shown: std.ArrayList(u32) = .empty, // the drafts this round's window verifies (node indices, window order) shown_set: bool = false, - burst_seen: bool = false, // the plain switch: an ask batch reached plain_streams, the normal rules apply from then on - wide_asks: u8 = 0, // the plain switch: asks made while never wide, a lone stream drafts from its second /// Node i of the held tree: the main chain's drafts, then each sibling's rows. pub fn nodes(st: *const St) u32 { @@ -435,16 +433,23 @@ pub const Lanes = struct { } } - /// The hold before a stream's first draft (the plain switch): until its burst has formed or its second ask. - const lone_asks_before_draft = 2; // a stream admitted alone drafts from its second ask + /// The plain switch: plain_streams or more streams decoding (past their prompts), so the round's windows are plain. + fn wideRound(x: *const Lanes) bool { + const at = x.opts.depth.plain_streams; + if (at == 0) return false; + var n: usize = 0; + for (x.by_slot) |o| if (o) |st| { + n += @intFromBool(st.len > 0); + }; + return n >= at; + } fn draft(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { const x = self(ptr); try x.settle(); _ = x.scratch.reset(.retain_capacity); const a = x.scratch.allocator(); - const switch_on = x.opts.depth.plain_streams > 0; - const burst = switch_on and requests.len >= x.opts.depth.plain_streams; // the round's burst: the normal rules apply + const wide = x.wideRound(); const asks = try a.alloc(iface.Ask, requests.len); const whose = try a.alloc(*St, requests.len); const depths = try a.alloc(u32, requests.len); @@ -454,11 +459,12 @@ pub const Lanes = struct { var n: usize = 0; for (requests) |r| { const st = try x.get(r.stream); + const asks_now = r.depth > 0 and !wide; var at = st.len; if (r.rows) |path| { if (r.start != st.base) return error.PositionMismatch; const done = if (x.spec) |*sp| sp.ingested(st.slot, r.start, path) else false; - if (!done) try x.ingest(st, r.start, path); // the drafter's ring holds every committed row, drafted or not // a shared round drafts before its keep: the path, not st.len + if (!done) try x.ingest(st, r.start, path); // a shared round drafts before its keep: the path, not st.len at = r.start + path.len; } if (r.position != at + 1) return error.PositionMismatch; @@ -472,11 +478,7 @@ pub const Lanes = struct { st.n = 0; st.n_sib = 0; st.start = at; - if (r.depth == 0) continue; - if (switch_on) { - if (burst) st.burst_seen = true else st.wide_asks += 1; - if (!st.burst_seen and st.wide_asks < lone_asks_before_draft) continue; // the plain switch: the hold until its burst has formed or its second ask - } + if (!asks_now) continue; // no drafts asked, or the plain switch: a round this wide drafts none if (n == x.opts.slots) return error.TooManySlots; asks[n] = .{ .slot = st.slot, .anchor = anchor, .start = at, .params = dspark.Params.of(r.stream.sampling, at) }; whose[n] = st; @@ -536,7 +538,7 @@ pub const Lanes = struct { fn heldTree(ptr: *anyopaque, s: *Stream, gpa: Allocator) anyerror!?lanes.stream.Held { const x = self(ptr); const st = try x.get(s); - if (st.n == 0) return .{ .count = 0 }; // the ask was skipped (the plain switch's first round): the engine plans a plain window + if (st.n == 0) return .{ .count = 0 }; // no ask this round (the plain switch): the engine plans a plain window const total = st.nodes(); const tokens = try gpa.alloc(u32, total); errdefer gpa.free(tokens); diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 64350155d6d..613d6166219 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -1,4 +1,4 @@ -//! TF_DSV41_DRAFT_PLAIN_STREAMS: the env knob, and a round with enough active windows decodes plain. +//! TF_DSV41_DRAFT_PLAIN_STREAMS: the env knob, and a round with enough decoding streams decodes plain. const std = @import("std"); const lanes = @import("lanes"); const twin = @import("twin.zig"); @@ -8,6 +8,7 @@ const depth = @import("depth.zig"); const calib_env = @import("calib_env.zig"); const gpa = std.testing.allocator; +const t_ = std.testing; const Env = struct { var pairs: []const [2][]const u8 = &.{}; @@ -17,19 +18,39 @@ const Env = struct { } }; -test "TF_DSV41_DRAFT_PLAIN_STREAMS: parsed, 0 when unset, 0 to 16 taken, anything else refused by name" { - Env.pairs = &.{}; - try std.testing.expectEqual(@as(usize, 0), (try calib_env.depthSettings(&Env.get)).plain_streams); - Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", "3" }}; - try std.testing.expectEqual(@as(usize, 3), (try calib_env.depthSettings(&Env.get)).plain_streams); - for ([_][]const u8{ "16", "17", "4294967296", "18446744073709551616", "-1", "oops", "3x" }) |raw| { - Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", raw }}; - if (std.mem.eql(u8, raw, "16")) try std.testing.expectEqual(@as(usize, 16), (try calib_env.depthSettings(&Env.get)).plain_streams) else try std.testing.expectError(error.DraftPlainStreams, calib_env.depthSettings(&Env.get)); - } +fn plainStreams(raw: []const u8) !usize { + Env.pairs = &.{.{ "TF_DSV41_DRAFT_PLAIN_STREAMS", raw }}; + defer Env.pairs = &.{}; + return (try calib_env.depthSettings(&Env.get)).plain_streams; +} + +test "TF_DSV41_DRAFT_PLAIN_STREAMS: 0 when unset, 0 to 16 taken, anything else refused by name" { Env.pairs = &.{}; + try t_.expectEqual(@as(usize, 0), (try calib_env.depthSettings(&Env.get)).plain_streams); + for ([_][]const u8{ "0", "3", " 16 " }, [_]usize{ 0, 3, 16 }) |raw, want| try t_.expectEqual(want, try plainStreams(raw)); + for ([_][]const u8{ "17", "4294967296", "18446744073709551616", "-1", "oops", "3x" }) |raw| + try t_.expectError(error.DraftPlainStreams, plainStreams(raw)); } -fn runPlain(cases: []const []const u32, setup: Setup, max_news: []const u32) !Run { +/// A run's settings: the switch, staggered arrivals, sampling. +pub const Opts = struct { set: depth.Settings = .{}, stagger: bool = false, sampling: ?lanes.Sampling = null }; + +pub const Run = struct { + drafted: u64, + accepted: u64, + ok: bool, + passes: u64, + ingests: u64, + out: [][]u32, + + pub fn deinit(r: Run) void { + for (r.out) |o| gpa.free(o); + gpa.free(r.out); + } +}; + +/// Streams of these prompts and lengths through the lanes over the twin; staggered: one added each step. +pub fn run(cases: []const []const u32, max_news: []const u32, o: Opts) !Run { const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; // a confident drafter (lanes_test's): drafts are chosen const t = try twin.Twin.init(gpa, dims, cases.len + 1); defer t.deinit(); @@ -37,7 +58,7 @@ fn runPlain(cases: []const []const u32, setup: Setup, max_news: []const u32) !Ru defer d.deinit(); var c = try costs_mod.defaults(gpa, 64); defer c.deinit(gpa); - const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = setup, .spec = .{ .on = true } }); + const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = o.set, .spec = .{ .on = true } }); defer x.deinit(); var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -49,176 +70,89 @@ fn runPlain(cases: []const []const u32, setup: Setup, max_news: []const u32) !Ru x.attach(&e); const streams = try gpa.alloc(lanes.Stream, cases.len); defer gpa.free(streams); - for (cases, streams, max_news) |p, *s, *n| s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = n.* }); + for (cases, streams, max_news) |p, *s, n| s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = p, .max_new = n, .sampling = o.sampling }); defer for (streams) |*s| s.deinit(gpa); - for (streams) |*s| try e.addStream(s); - while (e.activeCount() > 0) try e.step(); + var added: usize = 0; + if (!o.stagger) { + for (streams) |*s| try e.addStream(s); + added = streams.len; + } + var steps: usize = 0; + while (e.activeCount() > 0 or added < streams.len) : (steps += 1) { + if (added < streams.len) { + try e.addStream(&streams[added]); + added += 1; + } + try e.step(); + if (steps > 1000) return error.TooManySteps; + } var ok = true; for (streams) |*s| ok = ok and s.emitted().len > 0; - return .{ .drafted = e.drafted, .ok = ok, .passes = d.passes, .out = blk: { - const out = try gpa.alloc([]u32, cases.len); - for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); - break :blk out; - } }; + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*q, *s| q.* = try gpa.dupe(u32, s.emitted()); + return .{ .drafted = e.drafted, .accepted = e.accepted, .ok = ok, .passes = d.passes, .ingests = d.ingests, .out = out }; } -fn freeRun(r: *Run) void { - for (r.out) |o| gpa.free(o); - gpa.free(r.out); +/// Each reply is its prompt's serial decode at its own length. +pub fn expectSerial(r: Run, cases: []const []const u32, max_news: []const u32, sampling: ?lanes.Sampling) !void { + const t = try twin.Twin.init(gpa, .{ .bigram = 8.0, .head = 0.08 }, 1); + defer t.deinit(); + try t_.expect(r.ok); + for (r.out, cases, max_news) |got, p, n| { + const want = try t.serial(0, p, n, sampling, null); + defer gpa.free(want); + try t_.expectEqualSlices(u32, want, got); + } } -const Setup = depth.Settings; - -const Run = struct { drafted: u64, ok: bool, passes: u64, out: [][]u32 }; - -const a1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; +pub const a1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; +const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; +const fours = [_]u32{ 40, 40, 40, 40 }; -test "a round with plain_streams or more active windows decodes plain; below it drafts; 0 keeps drafts" { - // identical prompts and greedy sampling keep every round a full 4-window round until the streams end together - const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; - const fours = [_]u32{ 40, 40, 40, 40 }; - // 4 active, threshold 3: every drafted window is forced to count 0 - const plain = try runPlain(&four, .{ .plain_streams = 3 }, &fours); - defer freeRun(@constCast(&plain)); - try std.testing.expect(plain.ok); - try std.testing.expectEqual(@as(u64, 0), plain.drafted); - // 3 active, threshold 3: at it, plain too - const at = try runPlain(four[0..3], .{ .plain_streams = 3 }, fours[0..3]); - defer freeRun(@constCast(&at)); - try std.testing.expect(at.ok); - try std.testing.expectEqual(@as(u64, 0), at.drafted); - // 2 active, threshold 3: below it, drafts are kept - const below = try runPlain(four[0..2], .{ .plain_streams = 3 }, fours[0..2]); - defer freeRun(@constCast(&below)); - try std.testing.expect(below.ok); - try std.testing.expect(below.drafted > 0); - // 4 active, threshold 0 (off): drafts are kept - const off = try runPlain(&four, .{}, &fours); - defer freeRun(@constCast(&off)); - try std.testing.expect(off.ok); - try std.testing.expect(off.drafted > 0); +test "a round with plain_streams or more decoding streams decodes plain; below it drafts; 0 keeps drafts" { + // identical prompts and greedy sampling keep every round a full round until the streams end together + const plain = try run(&four, &fours, .{ .set = .{ .plain_streams = 3 } }); + defer plain.deinit(); + try t_.expectEqual(@as(u64, 0), plain.drafted); + const at = try run(four[0..3], fours[0..3], .{ .set = .{ .plain_streams = 3 } }); + defer at.deinit(); + try t_.expectEqual(@as(u64, 0), at.drafted); + const below = try run(four[0..2], fours[0..2], .{ .set = .{ .plain_streams = 3 } }); + defer below.deinit(); + try t_.expect(below.drafted > 0); + const off = try run(&four, &fours, .{}); + defer off.deinit(); + try t_.expect(off.drafted > 0); + for ([_]Run{ plain, at, below, off }, [_]usize{ 4, 3, 2, 4 }) |r, n| try expectSerial(r, four[0..n], fours[0..n], null); } -test "a stream's first draft ask waits for its first round; the replies stay the serial ones" { - const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; - const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; - const fours = [_]u32{ 40, 40, 40, 40 }; - // off, 4 streams: every round drafts, every round runs the pass - const off = try runPlain(&four, .{}, &fours); - defer freeRun(@constCast(&off)); - try std.testing.expect(off.passes > 0); - // at the threshold: shared rounds run no pass and the opening asks are skipped, so no pass runs at all - const on = try runPlain(&four, .{ .plain_streams = 3 }, &fours); - defer freeRun(@constCast(&on)); - try std.testing.expectEqual(@as(u64, 0), on.drafted); - try std.testing.expectEqual(@as(u64, 0), on.passes); - // tokens identical drafted or plain - for (off.out, on.out) |a, b| try std.testing.expectEqualSlices(u32, a, b); - // serial reference - { - const t = try twin.Twin.init(gpa, dims, 1); - defer t.deinit(); - for (on.out, four) |got, p| { - const want = try t.serial(0, p, 40, null, null); - defer gpa.free(want); - try std.testing.expectEqualSlices(u32, want, got); - } - } +test "below the threshold the switch changes nothing: 1 and 2 streams draft from their first ask as with it off" { + for ([_]usize{ 1, 2 }) |n| for ([_]bool{ false, true }) |stagger| { + const on = try run(four[0..n], fours[0..n], .{ .set = .{ .plain_streams = 3 }, .stagger = stagger }); + defer on.deinit(); + const off = try run(four[0..n], fours[0..n], .{ .stagger = stagger }); + defer off.deinit(); + try t_.expect(on.drafted > 0); + try t_.expectEqual(off.passes, on.passes); + try t_.expectEqual(off.drafted, on.drafted); + try t_.expectEqual(off.accepted, on.accepted); + try t_.expectEqual(off.ingests, on.ingests); + for (off.out, on.out) |a, b| try t_.expectEqualSlices(u32, a, b); + }; } -test "a stream drafts no round until its burst has formed or its second ask while never wide" { - // staggered arrival, plain_streams 3: the burst forms at the third stream, from which every stream drafts - const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; - { - const t = try twin.Twin.init(gpa, dims, 5); - defer t.deinit(); - const d = try twin.Drafter.init(gpa, t, 5); - defer d.deinit(); - var c = try costs_mod.defaults(gpa, 64); - defer c.deinit(gpa); - const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = .{ .plain_streams = 3 }, .spec = .{ .on = true } }); - defer x.deinit(); - var arena = std.heap.ArenaAllocator.init(gpa); - defer arena.deinit(); - var cfg = try lanes.Config.init(gpa, try x.model(arena.allocator(), 64), 16, 15); - defer cfg.deinit(gpa); - var clock: lanes.fake.FixedClock = .{}; - var e = lanes.Engine.init(gpa, &cfg, x.backend(), clock.clock()); - defer e.deinit(); - x.attach(&e); - const streams = try gpa.alloc(lanes.Stream, 4); - defer gpa.free(streams); - var i: usize = 0; - var steps: u64 = 0; - while (e.activeCount() > 0 or i < 4) { - try e.step(); - steps += 1; - if (i < 4) { - const s = &streams[i]; - s.* = try lanes.Stream.init(gpa, .{ .id = "s", .prompt = &a1, .max_new = 40 }); - try e.addStream(s); - i += 1; - } - if (steps > 400) return error.TooManySteps; - } - defer for (streams) |*s| s.deinit(gpa); - var ok = true; - for (streams) |*s| ok = ok and s.emitted().len > 0; - try std.testing.expect(ok); - try std.testing.expect(d.passes > 0); // the burst formed at three streams: the later rounds draft - try std.testing.expect(e.drafted > 0); - // tokens equal to the serial reference - for (streams) |*s| { - const want = try t.serial(0, &a1, 40, null, null); - defer gpa.free(want); - try std.testing.expectEqualSlices(u32, want, s.emitted()); - } - } - // the truly-lone control: one stream admitted alone drafts from its second ask - const lone = try runPlain(&.{&a1}, .{ .plain_streams = 3 }, &.{40}); - defer freeRun(@constCast(&lone)); - try std.testing.expect(lone.ok); - try std.testing.expect(lone.drafted > 0); - try std.testing.expect(lone.passes > 0); - { - const t = try twin.Twin.init(gpa, dims, 1); - defer t.deinit(); - const want = try t.serial(0, &a1, 40, null, null); - defer gpa.free(want); - try std.testing.expectEqualSlices(u32, want, lone.out[0]); - } +test "at the threshold only the asks made before plain_streams streams decode run a pass; replies unchanged" { + const on = try run(&four, &fours, .{ .set = .{ .plain_streams = 3 } }); + defer on.deinit(); + try t_.expectEqual(@as(u64, 0), on.drafted); + try t_.expect(on.passes <= 2); // the first two streams' opening asks, made while fewer than three decoded + try expectSerial(on, &four, &fours, null); } -test "below the threshold a stream drafts on its first ask; drafting resumes once streams finish" { - const four = [_][]const u32{ &a1, &a1, &a1, &a1 }; - const dims: twin.Dims = .{ .bigram = 8.0, .head = 0.08 }; - const fours = [_]u32{ 40, 40, 40, 40 }; - // the threshold not reached (1-2 active): the first ask is held, the second drafts - const one = try runPlain(four[0..1], .{ .plain_streams = 3 }, fours[0..1]); - defer freeRun(@constCast(&one)); - try std.testing.expect(one.ok); - try std.testing.expect(one.drafted > 0); - try std.testing.expect(one.passes > 0); - const two = try runPlain(four[0..2], .{ .plain_streams = 3 }, fours[0..2]); - defer freeRun(@constCast(&two)); - try std.testing.expect(two.ok); - try std.testing.expect(two.drafted > 0); - try std.testing.expect(two.passes > 0); - // 4 streams, two finish early: admitted four-deep, the burst has formed from round one, so every ask drafts - const twos = [_]u32{ 6, 6, 40, 40 }; - const mixed = try runPlain(&four, .{ .plain_streams = 3 }, &twos); - defer freeRun(@constCast(&mixed)); - try std.testing.expect(mixed.ok); - try std.testing.expect(mixed.drafted > 0); // after two streams finished, 2 active: drafts asked again - try std.testing.expect(mixed.passes > 0); // every pass is from the resumed rounds - // serial reference, each stream's own length - { - const t = try twin.Twin.init(gpa, dims, 1); - defer t.deinit(); - for (mixed.out, four, &twos) |got, p, *n| { - const want = try t.serial(0, p, n.*, null, null); - defer gpa.free(want); - try std.testing.expectEqualSlices(u32, want, got); - } - } +test "staggered arrivals: streams draft until plain_streams decode, plain after, drafting again as they end" { + const lens = [_]u32{ 40, 46, 52, 58 }; + const on = try run(&four, &lens, .{ .set = .{ .plain_streams = 3 }, .stagger = true }); + defer on.deinit(); + try t_.expect(on.passes > 0 and on.drafted > 0); + try expectSerial(on, &four, &lens, null); } From 3f9cd87438047ba7183aecf060184abec48d2a6d Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:46:10 +0100 Subject: [PATCH 031/186] dsv41 metal: the taps ring takes a window's kept rows only, so it never holds a rejected draft's taps --- zig/src/families/deepseek_v41/metal/target.zig | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/zig/src/families/deepseek_v41/metal/target.zig b/zig/src/families/deepseek_v41/metal/target.zig index 2e83b64b236..7958b7924b4 100644 --- a/zig/src/families/deepseek_v41/metal/target.zig +++ b/zig/src/families/deepseek_v41/metal/target.zig @@ -50,7 +50,7 @@ const Slot = struct { base: u64 = 0, rows: u32 = 0, open: bool = false, - /// the taps of the last `ring rows` positions, position p at row p % ring rows + /// the taps of the last `ring rows` committed positions, position p at row p % ring rows ring: []u16, /// the taps `taps` names after a prompt: its last rows in order ptaps: mtl.Buffer, @@ -307,19 +307,18 @@ pub const Target = struct { } for (group) |g| { const s = &t.slots[g.slot]; + if (s.open) t.commitTaps(s, s.rows); // a window never kept keeps all its rows s.whole = false; try s.hist.appendSlice(t.a, g.tokens); s.high = @max(s.high, s.hist.items.len); } t.launch(group) catch |e| return t.fail(e); const lg = t.b.logitsOf(c, total); - const w = tapWidth(c); var r0: usize = 0; for (group, choices) |g, out| { const s = &t.slots[g.slot]; const n = g.tokens.len; for (0..n) |r| out[r] = try t.choose(lg[(r0 + r) * V ..][0..V], g.draws[r], g.sampling); - if (s.kv.taps) |tb| for (0..n) |r| t.ringPut(s, g.start + r, tb.slice(u16, n * w)[r * w ..][0..w]); s.base = g.start; s.rows = @intCast(n); s.open = true; @@ -370,6 +369,7 @@ pub const Target = struct { for (path, 0..) |r, i| if (r != i) return error.NotAPath; s.open = false; const m: u32 = @intCast(path.len); + t.commitTaps(s, m); if (m == s.rows) return; const len = s.base + m; s.hist.shrinkRetainingCapacity(@intCast(len)); @@ -379,6 +379,13 @@ pub const Target = struct { t.stats.rewinds += 1; } + /// The ring takes the first `m` rows of the slot's last window (its kept rows); rows not kept never land in it. + fn commitTaps(t: *const Target, s: *Slot, m: u32) void { + const tb = s.kv.taps orelse return; + const w = tapWidth(t.env.c); + for (0..m) |r| t.ringPut(s, s.base + r, tb.slice(u16, @as(usize, s.rows) * w)[r * w ..][0..w]); + } + fn taps(ptr: *anyopaque, slot: u32) iface.Taps { const t = self(ptr); const s = &t.slots[slot]; @@ -395,6 +402,7 @@ pub const Target = struct { const s = &t.slots[slot]; if (!s.live) return; s.live = false; + if (s.open) t.commitTaps(s, s.rows); s.open = false; if (s.whole and t.cache.on()) t.cache.save(t.env.c, &s.kv, s.hist.items, s.ring, null) catch {}; } From dd60b6ee51cbad61d5bdd556d9cc1b4ddc53aabf Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:46:17 +0100 Subject: [PATCH 032/186] dsv41 metal: plain rounds leave their rows in the target's history, and the next ask catches the drafter up --- CHANGELOG.md | 1 + zig/src/families/deepseek_v41/draft/draft.zig | 1 + zig/src/families/deepseek_v41/draft/iface.zig | 11 ++++ .../deepseek_v41/draft/lane_ingest.zig | 39 +++++++++++++++ zig/src/families/deepseek_v41/draft/lanes.zig | 20 ++------ .../deepseek_v41/draft/plain_catchup_test.zig | 50 +++++++++++++++++++ .../deepseek_v41/draft/plain_streams_test.zig | 10 ++-- .../families/deepseek_v41/draft/spec_test.zig | 17 +++++-- zig/src/families/deepseek_v41/draft/twin.zig | 30 +++++++++++ .../families/deepseek_v41/metal/target.zig | 19 ++++++- zig/tests/dsv41_metal/serve_test.zig | 22 ++++++++ 11 files changed, 193 insertions(+), 27 deletions(-) create mode 100644 zig/src/families/deepseek_v41/draft/lane_ingest.zig create mode 100644 zig/src/families/deepseek_v41/draft/plain_catchup_test.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index f7dbd790269..a363d7927ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -84,6 +84,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. - DeepSeek-V4.1 on Metal: rounds with four or more active streams decode plain instead of drafting, since drafts lose there (#1720, part of #1676). - DeepSeek-V4.1 on Metal: rounds with three or more active streams decode plain and skip the draft pass, since drafts lose there (#1720, part of #1676). +- DeepSeek-V4.1 on Metal: rounds with three or more decoding streams decode plain, with no draft pass and no drafter ingest, since drafts lose there; when fewer streams decode again the drafter takes the rows it missed from the target, so its drafts are the ones it would have made (#1720, part of #1676). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. diff --git a/zig/src/families/deepseek_v41/draft/draft.zig b/zig/src/families/deepseek_v41/draft/draft.zig index 69559911c81..78f48ed0ae2 100644 --- a/zig/src/families/deepseek_v41/draft/draft.zig +++ b/zig/src/families/deepseek_v41/draft/draft.zig @@ -39,6 +39,7 @@ test { _ = calib_env; _ = @import("calib_env_test.zig"); _ = @import("plain_streams_test.zig"); + _ = @import("plain_catchup_test.zig"); _ = calib_measure; _ = @import("calib_measure_test.zig"); _ = branches_oracle; diff --git a/zig/src/families/deepseek_v41/draft/iface.zig b/zig/src/families/deepseek_v41/draft/iface.zig index f8bfae834e6..d63cae1b376 100644 --- a/zig/src/families/deepseek_v41/draft/iface.zig +++ b/zig/src/families/deepseek_v41/draft/iface.zig @@ -48,8 +48,19 @@ pub const Target = struct { warm: ?*const fn (ptr: *anyopaque, asks: []const Warm) void = null, /// After `prefill`: the prompt rows the slot took from kept state (its first row computed is this one). resumed: ?*const fn (ptr: *anyopaque, slot: u32) u64 = null, + /// Committed rows' taps for positions [start, start + rows), in order; null: not kept. null fn: none kept. + history: ?*const fn (ptr: *anyopaque, slot: u32, start: u64, rows: u32) ?Taps = null, }; + /// The target keeps committed rows' taps, so the drafter may take a plain round's rows later (`history`). + pub fn keepsHistory(t: Target) bool { + return t.vtable.history != null; + } + + pub fn history(t: Target, slot: u32, start: u64, rows: u32) ?Taps { + return if (t.vtable.history) |f| f(t.ptr, slot, start, rows) else null; + } + /// The prompt rows the slot's last `prefill` took from kept state (0: none, or a target that keeps none). pub fn resumedRows(t: Target, slot: u32) u64 { return if (t.vtable.resumed) |f| f(t.ptr, slot) else 0; diff --git a/zig/src/families/deepseek_v41/draft/lane_ingest.zig b/zig/src/families/deepseek_v41/draft/lane_ingest.zig new file mode 100644 index 00000000000..3adae056a8e --- /dev/null +++ b/zig/src/families/deepseek_v41/draft/lane_ingest.zig @@ -0,0 +1,39 @@ +//! The drafter's committed rows for the lanes: a window's kept rows, after any rows that were left for later. +const iface = @import("iface.zig"); +const branches = @import("branches.zig"); +const ln = @import("lanes.zig"); + +/// A round's kept rows into the drafter: first the rows a skipped absorb left out (`catchUp`), then these. +pub fn ingest(x: *ln.Lanes, st: *ln.St, start: u64, rows: []const u32) !void { + try catchUp(x, st, start, start + rows.len); + st.ring_end = start + rows.len; + const span = iface.ingestSpan(rows.len, x.opts.shape.window); + if (span.count == 0) return; + const taps = x.target.taps(st.slot); + const m = taps.map orelse return x.pass.ingest(st.slot, start + span.skip, taps, rows[span.skip..]); + // a tree resolved as chains: the path's rows in the taps of the chain that holds it + var mapped: [branches.max_rows]u32 = undefined; + if (span.count > mapped.len) return error.TreeRows; + for (rows[span.skip..], mapped[0..span.count]) |r, *o| { + if (r >= m.len or m[r] == branches.unrun) return error.TreeTaps; + o.* = m[r]; + } + try x.pass.ingest(st.slot, start + span.skip, taps, mapped[0..span.count]); +} + +/// Absorb skipped: the rows stay in the target's history, and the next ingest takes them from there. +pub fn deferred(x: *const ln.Lanes) bool { + return x.target.keepsHistory(); +} + +/// The committed rows between the drafter's end and `start`, the last `window` before `end` only (all it reads). +fn catchUp(x: *ln.Lanes, st: *ln.St, start: u64, end: u64) !void { + const from = st.ring_end orelse return; + const lo = @max(from, end -| x.opts.shape.window); + if (lo >= start) return; + const n: u32 = @intCast(start - lo); + const taps = x.target.history(st.slot, lo, n) orelse return error.HistoryLost; + const rows = try x.scratch.allocator().alloc(u32, n); + for (rows, 0..) |*r, i| r.* = @intCast(i); + try x.pass.ingest(st.slot, lo, taps, rows); +} diff --git a/zig/src/families/deepseek_v41/draft/lanes.zig b/zig/src/families/deepseek_v41/draft/lanes.zig index 93ce5ce3f68..70d176a0656 100644 --- a/zig/src/families/deepseek_v41/draft/lanes.zig +++ b/zig/src/families/deepseek_v41/draft/lanes.zig @@ -11,7 +11,6 @@ const depth = @import("depth.zig"); const calib = @import("calib.zig"); const costs_mod = @import("costs.zig"); const Policy = @import("policy.zig").Policy; -const branches = @import("branches.zig"); const spec_mod = @import("spec.zig"); pub const max_block = 16; @@ -43,6 +42,7 @@ pub const St = struct { n_sib: u32 = 0, shown: std.ArrayList(u32) = .empty, // the drafts this round's window verifies (node indices, window order) shown_set: bool = false, + ring_end: ?u64 = null, // the drafter holds the committed rows before this position (null: none yet) /// Node i of the held tree: the main chain's drafts, then each sibling's rows. pub fn nodes(st: *const St) u32 { @@ -151,20 +151,7 @@ pub const Lanes = struct { }; } - fn ingest(x: *Lanes, st: *St, start: u64, rows: []const u32) !void { - const span = iface.ingestSpan(rows.len, x.opts.shape.window); - if (span.count == 0) return; - const taps = x.target.taps(st.slot); - const m = taps.map orelse return x.pass.ingest(st.slot, start + span.skip, taps, rows[span.skip..]); - // a tree resolved as chains: the path's rows in the taps of the chain that holds it - var mapped: [branches.max_rows]u32 = undefined; - if (span.count > mapped.len) return error.TreeRows; - for (rows[span.skip..], mapped[0..span.count]) |r, *o| { - if (r >= m.len or m[r] == branches.unrun) return error.TreeTaps; - o.* = m[r]; - } - try x.pass.ingest(st.slot, start + span.skip, taps, mapped[0..span.count]); - } + const ingest = @import("lane_ingest.zig").ingest; fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { const x = self(ptr); @@ -464,7 +451,8 @@ pub const Lanes = struct { if (r.rows) |path| { if (r.start != st.base) return error.PositionMismatch; const done = if (x.spec) |*sp| sp.ingested(st.slot, r.start, path) else false; - if (!done) try x.ingest(st, r.start, path); // a shared round drafts before its keep: the path, not st.len + // a shared round drafts before its keep: the path, not st.len; no ask: the next ingest catches it up + if (!done and (asks_now or !@import("lane_ingest.zig").deferred(x))) try x.ingest(st, r.start, path); at = r.start + path.len; } if (r.position != at + 1) return error.PositionMismatch; diff --git a/zig/src/families/deepseek_v41/draft/plain_catchup_test.zig b/zig/src/families/deepseek_v41/draft/plain_catchup_test.zig new file mode 100644 index 00000000000..9ca0ff3cd52 --- /dev/null +++ b/zig/src/families/deepseek_v41/draft/plain_catchup_test.zig @@ -0,0 +1,50 @@ +//! Plain rounds absorb nothing into the drafter; the next ask catches it up from the target's history, exactly. +const std = @import("std"); +const lanes = @import("lanes"); +const ps = @import("plain_streams_test.zig"); + +const t_ = std.testing; +const warm: lanes.Sampling = .{ .seed = 5, .temperature = 1.0, .top_k = 20, .top_p = 0.95, .min_p = 0.02 }; +const four = [_][]const u32{ &ps.a1, &ps.a1, &ps.a1, &ps.a1 }; + +test "plain rounds run no ingest when the target keeps history; without it each round absorbs its rows" { + const lens = [_]u32{ 40, 40, 40, 40 }; + const kept = try ps.run(&four, &lens, .{ .set = .{ .plain_streams = 3 } }); + defer kept.deinit(); + const each = try ps.run(&four, &lens, .{ .set = .{ .plain_streams = 3 }, .history = false }); + defer each.deinit(); + try t_.expectEqual(@as(u64, 4), kept.ingests); // the four prompts' rows only + try t_.expect(each.ingests > 4 * 30); + for (kept.out, each.out) |a, b| try t_.expectEqualSlices(u32, a, b); + try ps.expectSerial(kept, &four, &lens, null); +} + +test "dropping below the threshold catches the drafter up: the same drafts as a drafter fed every round" { + // two streams end early: rounds turn plain to drafting; the long pair's gap is past the drafter's window (16) + const cases = [_][4]u32{ .{ 6, 6, 40, 40 }, .{ 30, 30, 70, 70 } }; + for (&cases) |*lens| for ([_]?lanes.Sampling{ null, warm }) |smp| { + const kept = try ps.run(&four, lens, .{ .set = .{ .plain_streams = 3 }, .check_ring = true, .sampling = smp }); + defer kept.deinit(); + const each = try ps.run(&four, lens, .{ .set = .{ .plain_streams = 3 }, .history = false, .check_ring = true, .sampling = smp }); + defer each.deinit(); + try t_.expect(kept.drafted > 0); + try t_.expectEqual(each.drafted, kept.drafted); + try t_.expectEqual(each.accepted, kept.accepted); + try t_.expectEqual(each.passes, kept.passes); + try t_.expect(kept.ingests < each.ingests); + for (kept.out, each.out) |a, b| try t_.expectEqualSlices(u32, a, b); + try ps.expectSerial(kept, &four, lens, smp); + }; +} + +test "a never-plain run and a run that went plain give the same replies, greedy and sampled" { + const lens = [_]u32{ 8, 12, 50, 50 }; + for ([_]?lanes.Sampling{ null, warm }) |smp| { + const never = try ps.run(&four, &lens, .{ .sampling = smp, .check_ring = true }); + defer never.deinit(); + const went = try ps.run(&four, &lens, .{ .set = .{ .plain_streams = 3 }, .sampling = smp, .check_ring = true }); + defer went.deinit(); + try t_.expect(never.drafted > 0 and went.drafted > 0); + for (never.out, went.out) |a, b| try t_.expectEqualSlices(u32, a, b); + } +} diff --git a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig index 613d6166219..bc8151c1b30 100644 --- a/zig/src/families/deepseek_v41/draft/plain_streams_test.zig +++ b/zig/src/families/deepseek_v41/draft/plain_streams_test.zig @@ -32,8 +32,8 @@ test "TF_DSV41_DRAFT_PLAIN_STREAMS: 0 when unset, 0 to 16 taken, anything else r try t_.expectError(error.DraftPlainStreams, plainStreams(raw)); } -/// A run's settings: the switch, staggered arrivals, sampling. -pub const Opts = struct { set: depth.Settings = .{}, stagger: bool = false, sampling: ?lanes.Sampling = null }; +/// A run's settings: the switch, the target's history, the drafter's ring check, staggered arrivals, sampling. +pub const Opts = struct { set: depth.Settings = .{}, history: bool = true, check_ring: bool = false, stagger: bool = false, sampling: ?lanes.Sampling = null }; pub const Run = struct { drafted: u64, @@ -56,9 +56,11 @@ pub fn run(cases: []const []const u32, max_news: []const u32, o: Opts) !Run { defer t.deinit(); const d = try twin.Drafter.init(gpa, t, cases.len + 1); defer d.deinit(); + d.check_ring = o.check_ring; var c = try costs_mod.defaults(gpa, 64); defer c.deinit(gpa); - const x = try ln.Lanes.init(gpa, t.target(), d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = o.set, .spec = .{ .on = true } }); + const tgt = if (o.history) t.target() else t.targetNoHistory(); + const x = try ln.Lanes.init(gpa, tgt, d.pass(), c, .{ .shape = d.shape, .slots = 4, .depth = o.set, .spec = .{ .on = true } }); defer x.deinit(); var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -151,7 +153,7 @@ test "at the threshold only the asks made before plain_streams streams decode ru test "staggered arrivals: streams draft until plain_streams decode, plain after, drafting again as they end" { const lens = [_]u32{ 40, 46, 52, 58 }; - const on = try run(&four, &lens, .{ .set = .{ .plain_streams = 3 }, .stagger = true }); + const on = try run(&four, &lens, .{ .set = .{ .plain_streams = 3 }, .stagger = true, .check_ring = true }); defer on.deinit(); try t_.expect(on.passes > 0 and on.drafted > 0); try expectSerial(on, &four, &lens, null); diff --git a/zig/src/families/deepseek_v41/draft/spec_test.zig b/zig/src/families/deepseek_v41/draft/spec_test.zig index fbeb227aee1..b42716cd9da 100644 --- a/zig/src/families/deepseek_v41/draft/spec_test.zig +++ b/zig/src/families/deepseek_v41/draft/spec_test.zig @@ -121,7 +121,7 @@ const Run = struct { const Mode = enum { off, sync, device }; -fn run(cases: []const Case, mode: Mode, settings: spec.Settings, siblings: u32) !Run { +fn run(cases: []const Case, mode: Mode, settings: spec.Settings, siblings: u32, history: bool) !Run { const dims: twin.Dims = .{}; const slots: u32 = 4; const t = try twin.Twin.init(gpa, dims, slots + 1); @@ -133,7 +133,7 @@ fn run(cases: []const Case, mode: Mode, settings: spec.Settings, siblings: u32) const pass = if (mode == .device) d.passAsync() else d.pass(); var set = settings; set.on = mode != .off; - const x = try ln.Lanes.init(gpa, t.target(), pass, c, .{ .shape = d.shape, .slots = slots, .siblings = siblings, .spec = set }); + const x = try ln.Lanes.init(gpa, if (history) t.target() else t.targetNoHistory(), pass, c, .{ .shape = d.shape, .slots = slots, .siblings = siblings, .spec = set }); defer x.deinit(); var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -155,7 +155,14 @@ fn run(cases: []const Case, mode: Mode, settings: spec.Settings, siblings: u32) } fn expectSame(cases: []const Case, settings: spec.Settings, siblings: u32) ![3]spec.Stats { - const off = try run(cases, .off, settings, siblings); + var stats: [3]spec.Stats = undefined; + // with history an absorb without an ask waits for the next ingest, so ingest calls are counted without it + for ([_]bool{ true, false }) |history| stats = try expectSameAs(cases, settings, siblings, history); + return stats; +} + +fn expectSameAs(cases: []const Case, settings: spec.Settings, siblings: u32, history: bool) ![3]spec.Stats { + const off = try run(cases, .off, settings, siblings, history); defer off.deinit(); for (cases, off.out) |cs, got| { // the reference itself is the serial decode const t = try twin.Twin.init(gpa, .{}, 1); @@ -166,7 +173,7 @@ fn expectSame(cases: []const Case, settings: spec.Settings, siblings: u32) ![3]s } var stats: [3]spec.Stats = undefined; for ([_]Mode{ .sync, .device }, 1..) |m, i| { - const r = try run(cases, m, settings, siblings); + const r = try run(cases, m, settings, siblings, history); defer r.deinit(); for (off.out, r.out) |a, b| try std.testing.expectEqualSlices(u32, a, b); // the same drafts: the same windows, acceptance and kept counts as the run without speculation @@ -175,7 +182,7 @@ fn expectSame(cases: []const Case, settings: spec.Settings, siblings: u32) ![3]s try std.testing.expectEqualSlices(u64, &off.kept, &r.kept); // every used speculation replaced a pass, every launch ran one; skipped commits saved their ingests try std.testing.expectEqual(off.passes + r.stats.launched - r.stats.hit, r.passes); - try std.testing.expectEqual(off.ingests + r.stats.launched_slots, r.ingests + r.stats.ingests_skipped); + if (!history) try std.testing.expectEqual(off.ingests + r.stats.launched_slots, r.ingests + r.stats.ingests_skipped); stats[i] = r.stats; } stats[0] = .{}; diff --git a/zig/src/families/deepseek_v41/draft/twin.zig b/zig/src/families/deepseek_v41/draft/twin.zig index c665150483d..5380f891658 100644 --- a/zig/src/families/deepseek_v41/draft/twin.zig +++ b/zig/src/families/deepseek_v41/draft/twin.zig @@ -16,6 +16,7 @@ const Slot = struct { rows: std.ArrayList(f32) = .empty, // the last window's keys, values and taps [rows][3 dim] parents: std.ArrayList(i32) = .empty, taps: std.ArrayList(f32) = .empty, + hist: std.ArrayList(f32) = .empty, // committed rows' taps [len][dim] (Target.history) live: bool = false, fn deinit(s: *Slot, gpa: Allocator) void { @@ -24,6 +25,7 @@ const Slot = struct { s.rows.deinit(gpa); s.parents.deinit(gpa); s.taps.deinit(gpa); + s.hist.deinit(gpa); } fn len(s: Slot, d: usize) usize { @@ -174,6 +176,8 @@ pub const Twin = struct { choices[r] = try t.choose(a, row, draws[r], smp); } s.base = committed; + s.hist.shrinkRetainingCapacity(committed * d); + try s.hist.appendSlice(a, s.taps.items); for (0..n) |r| { // every row stays until keep (a chain that is never kept keeps them all) try s.k.appendSlice(a, s.rows.items[r * 2 * d ..][0..d]); try s.v.appendSlice(a, s.rows.items[r * 2 * d + d ..][0..d]); @@ -182,6 +186,11 @@ pub const Twin = struct { } pub fn target(t: *Twin) iface.Target { + return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .history = history } }; + } + + /// The target without `history`: a plain round's rows are then absorbed in their own round. + pub fn targetNoHistory(t: *Twin) iface.Target { return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release } }; } @@ -225,9 +234,11 @@ pub const Twin = struct { for (path[1..], path[0 .. path.len - 1]) |r, parent| if (s.parents.items[r] != @as(i32, @intCast(parent))) return error.NotAPath; s.k.shrinkRetainingCapacity(s.base * d); s.v.shrinkRetainingCapacity(s.base * d); + s.hist.shrinkRetainingCapacity(s.base * d); for (path) |r| { try s.k.appendSlice(t.gpa, s.rows.items[r * 2 * d ..][0..d]); try s.v.appendSlice(t.gpa, s.rows.items[r * 2 * d + d ..][0..d]); + try s.hist.appendSlice(t.gpa, s.taps.items[r * d ..][0..d]); } } @@ -236,6 +247,7 @@ pub const Twin = struct { const s = &t.slots[slot]; s.k.shrinkRetainingCapacity(s.base * t.dims.dim); s.v.shrinkRetainingCapacity(s.base * t.dims.dim); + s.hist.shrinkRetainingCapacity(s.base * t.dims.dim); } fn taps(ptr: *anyopaque, slot: u32) iface.Taps { @@ -244,6 +256,22 @@ pub const Twin = struct { return .{ .address = @intFromPtr(s.taps.items.ptr), .rows = @intCast(s.taps.items.len / t.dims.dim), .stride = t.dims.dim }; } + fn history(ptr: *anyopaque, slot: u32, start: u64, rows: u32) ?iface.Taps { + const t = self(ptr); + const s = &t.slots[slot]; + const d = t.dims.dim; + if (start + rows > s.hist.items.len / d) return null; + return .{ .address = @intFromPtr(s.hist.items[@intCast(start * d)..].ptr), .rows = rows, .stride = d }; + } + + /// A committed row's taps (what the drafter's ring should hold at `pos`), for the tests' ring check. + pub fn committedTaps(t: *const Twin, slot: u32, pos: u64) ?[]const f32 { + const s = &t.slots[slot]; + const d = t.dims.dim; + if (pos >= s.hist.items.len / d) return null; + return s.hist.items[@intCast(pos * d)..][0..d]; + } + fn release(ptr: *anyopaque, slot: u32) void { const t = self(ptr); t.slots[slot].live = false; @@ -286,6 +314,7 @@ pub const Drafter = struct { cw: []f32, passes: u64 = 0, ingests: u64 = 0, // ingest calls (tests) + check_ring: bool = false, // each context row must equal the target's committed taps at its position (tests) /// `begin` / `collect` (passAsync): the pass reads rings at `begin`, the host takes results at `collect`; all slots stash_drafts: []u32, stash_conf: []f32, @@ -435,6 +464,7 @@ pub const Drafter = struct { const pos = ctx.first + i; if (x.tags[ask.slot][pos % ring] != @as(i64, @intCast(pos))) return error.RingStale; const row = x.rings[ask.slot][(pos % ring) * d ..][0..d]; + if (x.check_ring and !std.mem.eql(f32, t.committedTaps(ask.slot, pos) orelse &.{}, row)) return error.RingWrong; t.project(t.wk, row, kbuf[i * 2 * d ..][0..d]); t.project(t.wv, row, kbuf[i * 2 * d + d ..][0..d]); keys[i] = kbuf[i * 2 * d ..][0..d]; diff --git a/zig/src/families/deepseek_v41/metal/target.zig b/zig/src/families/deepseek_v41/metal/target.zig index 7958b7924b4..ef832b6798d 100644 --- a/zig/src/families/deepseek_v41/metal/target.zig +++ b/zig/src/families/deepseek_v41/metal/target.zig @@ -39,7 +39,7 @@ pub const Options = struct { }; /// What the target did, for the tests and the logs. -pub const Stats = struct { prompts: u64 = 0, hits: u64 = 0, cached_rows: u64 = 0, windows: u64 = 0, rows: u64 = 0, rewinds: u64 = 0, launches: u64 = 0 }; +pub const Stats = struct { prompts: u64 = 0, hits: u64 = 0, cached_rows: u64 = 0, windows: u64 = 0, rows: u64 = 0, rewinds: u64 = 0, launches: u64 = 0, history_rows: u64 = 0 }; /// A slot: its device state, prompt tail, ids (rows run, the last window's included until its keep) and taps. const Slot = struct { @@ -138,7 +138,7 @@ pub const Target = struct { } pub fn target(t: *Target) iface.Target { - return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .resumed = resumedFn } }; + return .{ .ptr = t, .vtable = &.{ .prefill = prefill, .window = window, .keep = keep, .taps = taps, .release = release, .resumed = resumedFn, .history = history } }; } fn self(ptr: *anyopaque) *Target { @@ -386,6 +386,21 @@ pub const Target = struct { for (0..m) |r| t.ringPut(s, s.base + r, tb.slice(u16, @as(usize, s.rows) * w)[r * w ..][0..w]); } + /// Committed rows' taps [start, start + rows) from the ring, in order, in the prompt-taps buffer. + fn history(ptr: *anyopaque, slot: u32, start: u64, rows: u32) ?iface.Taps { + const t = self(ptr); + if (slot >= t.slots.len) return null; + const s = &t.slots[slot]; + const c = t.env.c; + const w = tapWidth(c); + const end: u64 = if (s.open) s.base else s.hist.items.len; + if (w == 0 or s.after_prompt or rows == 0 or start + rows > end or start + ringRows(c) < end) return null; + const out = s.ptaps.slice(u16, @as(usize, rows) * w); + for (0..rows) |i| @memcpy(out[i * w ..][0..w], s.ring[((start + i) % ringRows(c)) * w ..][0..w]); + t.stats.history_rows += rows; + return .{ .address = @intFromPtr(s.ptaps.contents()), .rows = rows, .stride = @intCast(w) }; + } + fn taps(ptr: *anyopaque, slot: u32) iface.Taps { const t = self(ptr); const s = &t.slots[slot]; diff --git a/zig/tests/dsv41_metal/serve_test.zig b/zig/tests/dsv41_metal/serve_test.zig index 9bf82fd51a5..8a9a0b9de56 100644 --- a/zig/tests/dsv41_metal/serve_test.zig +++ b/zig/tests/dsv41_metal/serve_test.zig @@ -265,3 +265,25 @@ test "dsv41 metal serve: (e) a request cancelled during its prompt leaves the sl try testing.expect(gone.reason.? == .cancelled); try testing.expect(ok0 and ok1); } + +test "dsv41 metal serve: (f) plain rounds at four streams, drafts again at two from the target's history == serial" { + var ps: [4][]u32 = undefined; + for (&ps, 0..) |*p, i| p.* = try ids(90 + 40 * i, 20 + i); + defer for (ps) |p| gpa.free(p); + const asks = [_]Ask{ .{ .prompt = ps[0], .max = 12 }, .{ .prompt = ps[1], .max = 12 }, .{ .prompt = ps[2], .max = 60 }, .{ .prompt = ps[3], .max = 60 } }; + var want: [4][]u32 = undefined; + try reference(&asks, &want); + defer freeAll(&want); + const m = try fs.open(gpa, io, seed, .{}); + const s = try Served.open(gpa, io, m.model(), ".", null); + defer Served.close(s); + var drafted: u64 = 0; + var got: [4][]u32 = undefined; + try serve(s.engine(), &asks, true, &got, &drafted); + defer freeAll(&got); + const bad = same(&want, &got); + const caught = m.t.stats.history_rows; + say(bad == 0 and caught > 0 and drafted > 0, "(f) 4 together vs serial: {d} differ; {d} rows caught up from history, {d} drafted rows, {d} ingests", .{ bad, caught, drafted, m.ad.ingests }); + try testing.expectEqual(@as(usize, 0), bad); + try testing.expect(caught > 0 and drafted > 0); +} From 1385658e890e71fb24718f153f44af260761d25f Mon Sep 17 00:00:00 2001 From: Kevin Gilbertson <880013+gilby@users.noreply.github.com> Date: Thu, 12 Nov 2026 00:32:44 +0100 Subject: [PATCH 033/186] dsv41 metal: a round of plain_streams streams verifies no copy windows either --- zig/src/core/lanes/drafted_copy_test.zig | 61 +++++++++++++++++++++--- zig/src/core/lanes/windows.zig | 9 ++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/zig/src/core/lanes/drafted_copy_test.zig b/zig/src/core/lanes/drafted_copy_test.zig index 9f4833503d5..46a4b64408c 100644 --- a/zig/src/core/lanes/drafted_copy_test.zig +++ b/zig/src/core/lanes/drafted_copy_test.zig @@ -47,6 +47,7 @@ const Watch = struct { var streams: [4]?*sm.Stream = @splat(null); var wasted: u32 = 0; var copies: u32 = 0; + var wide: u32 = 0; // copy windows verified in rounds of three streams fn slot(s: *sm.Stream) usize { for (&streams, 0..) |*x, i| if (x.* == null or x.* == s) { @@ -65,19 +66,21 @@ const Watch = struct { fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { for (windows) |w| if (w.held == 0 and w.tokens.len > 0) { copies += 1; + wide += @intFromBool(windows.len >= 3); if (at[slot(w.stream)] == w.positions[0]) wasted += 1; }; return inner_verify(ptr, windows, out); } }; -const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), wasted: u32 = 0, copies: u32 = 0 }; +const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), wasted: u32 = 0, copies: u32 = 0, wide: u32 = 0 }; -fn run(n: usize, drafts: bool, copy_first: bool) !Run { +fn run(n: usize, drafts: bool, copy_first: bool, plain_streams: u32) !Run { Watch.at = @splat(null); Watch.streams = @splat(null); Watch.wasted = 0; Watch.copies = 0; + Watch.wide = 0; var target: fake.Fake = .{ .gpa = gpa }; defer target.deinit(); var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .copy_first = copy_first } }; @@ -95,7 +98,7 @@ fn run(n: usize, drafts: bool, copy_first: bool) !Run { Watch.inner_verify = verified.verify; verified.verify = Watch.verify; seen.vtable = &verified; - const m: @import("config.zig").Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 }; + const m: @import("config.zig").Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4, .plain_streams = plain_streams }; var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15); defer cfg.deinit(gpa); var clock: fake.FixedClock = .{}; @@ -107,7 +110,7 @@ fn run(n: usize, drafts: bool, copy_first: bool) !Run { defer for (streams[0..n]) |*s| s.deinit(gpa); for (streams[0..n]) |*s| try engine.addStream(s); while (engine.activeCount() > 0) try engine.step(); - var out: Run = .{ .wasted = Watch.wasted, .copies = Watch.copies }; + var out: Run = .{ .wasted = Watch.wasted, .copies = Watch.copies, .wide = Watch.wide }; for (streams[0..n], 0..) |*s, i| { out.lens[i] = s.emitted().len; @memcpy(out.tokens[i][0..out.lens[i]], s.emitted()); @@ -117,9 +120,9 @@ fn run(n: usize, drafts: bool, copy_first: bool) !Run { test "a copy ahead takes the round: with copy_first no hold is asked for it, and replies equal one-token decode" { for ([_]usize{ 1, 3 }) |n| { - const plain = try run(n, false, false); - const before = try run(n, true, false); - const after = try run(n, true, true); + const plain = try run(n, false, false, 0); + const before = try run(n, true, false, 0); + const after = try run(n, true, true, 0); for (0..n) |i| { try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], before.tokens[i][0..before.lens[i]]); try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], after.tokens[i][0..after.lens[i]]); @@ -128,3 +131,47 @@ test "a copy ahead takes the round: with copy_first no hold is asked for it, and try std.testing.expectEqual(@as(u32, 0), after.wasted); } } + +test "plain_streams: rounds that wide verify no copy, narrower ones still copy; replies equal one-token decode" { + const plain = try run(3, false, false, 0); + const wide = try run(3, true, false, 3); + const narrow = try run(3, true, false, 4); + for (0..3) |i| { + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], wide.tokens[i][0..wide.lens[i]]); + try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], narrow.tokens[i][0..narrow.lens[i]]); + } + try std.testing.expectEqual(@as(u32, 0), wide.wide); + try std.testing.expect(narrow.wide > 0); +} + +/// A row policy that keeps every window as planned (a family's trim in place of the built-in allocation). +const KeepAll = struct { + const trim = @import("trim.zig"); + fn choose(_: *anyopaque, _: []const trim.Window, _: bool, _: std.mem.Allocator, _: []trim.Choice) anyerror!void {} + fn commit(_: *anyopaque, _: trim.Commit) anyerror!void {} + const vtable: trim.Trim.VTable = .{ .choose = choose, .commit = commit }; +}; + +test "plain_streams: a family's row policy keeps copies below the switch, and none from it on" { + const win = @import("windows.zig"); + const copied = [_]u32{ 5, 6, 7, 8, 9, 10, 11 }; + var cfg = try Config.init(gpa, .{ .exact_width = 16, .hidden_rows = true, .batch_rows = 32, .plain_streams = 3 }, 16, 15); + defer cfg.deinit(gpa); + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var e = Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer e.deinit(); + var keep: u8 = 0; + e.trim = .{ .ptr = &keep, .vtable = &KeepAll.vtable }; + var streams: [3]sm.Stream = undefined; + for (&streams) |*s| s.* = try sm.Stream.init(gpa, .{ .id = "copied", .prompt = &.{ 1, 2 }, .max_new = 64 }); + defer for (&streams) |*s| s.deinit(gpa); + for ([_]usize{ 2, 3 }) |n| { + _ = e.arena.reset(.retain_capacity); + var plans: [3]win.Plan = undefined; + for (plans[0..n], streams[0..n]) |*p, *s| p.* = .{ .stream = s, .position = 2, .kind = .copy, .tokens = &copied }; + try win.allocate(&e, plans[0..n]); + for (plans[0..n]) |p| try std.testing.expectEqual(@as(usize, if (n < 3) copied.len else 0), p.count()); + } +} diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index b3758b7c384..0d93234bfee 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -233,6 +233,14 @@ pub fn who(s: *Stream) depth.Who { return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .gain = s.grant_gain, .stream = s, .sampled = s.sampling != null }; } +/// cfg.plain_streams: a shared round this wide drafts none, its copies included (each window its pending row). +fn plainCopies(e: *const Engine, plans: []Plan) void { + if (e.cfg.plain_streams == 0 or plans.len < e.cfg.plain_streams) return; + for (plans) |*p| if (p.kind == .copy) { + p.* = .{ .stream = p.stream, .position = p.position, .kind = .none }; + }; +} + /// Whether any of a round's streams draws keyed samples (its rows carry their host work). pub fn anySampled(plans: []const Plan) bool { for (plans) |p| if (p.stream.sampling != null) return true; @@ -241,6 +249,7 @@ pub fn anySampled(plans: []const Plan) bool { /// Trim draft prefixes by landing probability to the shared width with the most expected tokens a ms. pub fn allocate(e: *Engine, plans: []Plan) !void { + plainCopies(e, plans); if (e.trim) |t| return trimmed(e, plans, t); const a = e.arena.allocator(); const fixed = try a.alloc(u32, plans.len); From 7a96db004e9c65d549b765e4ed64a1b09e962742 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:16:03 +0000 Subject: [PATCH 034/186] changelog: #1720's line from lines/1720.txt in place of its own --- CHANGELOG.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a363d7927ac..db707fab730 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- DeepSeek-V4.1 on Metal (experimental, opt in with `TF_EXPERIMENTAL=deepseek`): Macs now serve the release checkpoint natively with `tensorfold-dsv41` (`zig build dsv41-metal-serve`) for text prompts; image prompts come with #1722. Prompts and replies share rounds through the lane core with DSpark drafts, and the prompt cache reports `cached_tokens`. Rounds with three or more decoding streams decode plain (no draft pass, drafter ingest or copy windows), and when fewer streams decode again the drafter catches up from the target, so replies are unchanged. On an M3 Ultra and an M5 Ultra with 512 GB, drafts make one stream 1.21x to 1.50x faster than `--no-drafts` and two streams up to 1.13x faster, three and four streams are level, and `token_sha256` is equal in every cell (#1720, part of #1676). Thanks to @gilby, and to @3utsu for an independent M5 Ultra receipt. + - Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, making room for the execution witness's hook; nothing a user sees changes (#2896, refs #2741). Thanks to @kotevcode. - Qwen3.8-27B on Metal: prompt passes no longer write the GDN tails buffer, which nothing read. Replies are unchanged; on an M5 Ultra, time to first token is about 0.5% shorter at 2K, 8K and 32K prompts, one and four streams (#2915, for #2905). @@ -82,9 +84,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Qwen3.8-27B on Metal: three glue kernels nothing launched (`q27_norm`, `q27_mlp`, `q27_gdn_post`) and their host mirrors are removed; saved prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#2816, part of #2251). Thanks to @jschmied. - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. -- DeepSeek-V4.1 on Metal: rounds with four or more active streams decode plain instead of drafting, since drafts lose there (#1720, part of #1676). -- DeepSeek-V4.1 on Metal: rounds with three or more active streams decode plain and skip the draft pass, since drafts lose there (#1720, part of #1676). -- DeepSeek-V4.1 on Metal: rounds with three or more decoding streams decode plain, with no draft pass and no drafter ingest, since drafts lose there; when fewer streams decode again the drafter takes the rows it missed from the target, so its drafts are the ones it would have made (#1720, part of #1676). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. From 3c9fc8e9eb9de2ced1f80481e9c36c486fb36e1b Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 16:41:14 +0000 Subject: [PATCH 035/186] template: a chat template is held to a nesting depth, and its render to steps, stack and memory A checkpoint's chat template had no limit at all. 100,000 nested parentheses ran the parser off its stack, a sum of 100,000 terms built a tree that tall for the scope pass and the renderer, two nested loops over range(100000) never finished, `'' * 10**18` counted to its end, and a macro that calls itself inside a filter chain overflowed the stack before the macro limit. The parser now refuses more than 256 levels of its own recursion and any expression taller than that. A render stops by name past 50 million steps, 4 MiB of stack or 1 GiB of scratch, and an empty sequence repeated is empty at once. The stack is measured, not counted in levels, because frames differ too much in size between expression kinds. --- zig/src/core/template/eval.zig | 20 +++++- zig/src/core/template/limits_test.zig | 90 +++++++++++++++++++++++++++ zig/src/core/template/parser.zig | 24 ++++++- zig/src/core/template/template.zig | 78 +++++++++++++++++++---- 4 files changed, 196 insertions(+), 16 deletions(-) create mode 100644 zig/src/core/template/limits_test.zig diff --git a/zig/src/core/template/eval.zig b/zig/src/core/template/eval.zig index 1286e19e22e..18a2a2aaeeb 100644 --- a/zig/src/core/template/eval.zig +++ b/zig/src/core/template/eval.zig @@ -18,12 +18,26 @@ pub const Activation = struct { /// Macro nesting limit; Python's own frame budget runs out earlier, so this only guards the stack. const max_depth = 400; +/// A render's steps, stack and scratch bytes; only a template written to exhaust the server reaches one of them. +pub const Limits = struct { steps: u64 = 50_000_000, stack: usize = 4 << 20, bytes: usize = 1 << 30 }; pub const Ctx = struct { rt: v.Rt, vars: *const v.Map, out: *v.Buffer, depth: usize = 0, + limits: Limits = .{}, + steps: u64 = 0, + /// The stack where the render began; frames differ too much in size for a count of levels to bound it. + base: usize = 0, + + /// One more statement list or expression: a step of the render, taken only while its recursion has stack left. + fn enter(c: *Ctx) Error!void { + c.steps += 1; + if (c.steps > c.limits.steps) return c.rt.fail("the chat template ran more than {d} steps", .{c.limits.steps}); + if (c.base == 0) c.base = @frameAddress(); + if (c.base -| @frameAddress() > c.limits.stack) return c.rt.fail("the chat template recursed more than {d} KiB of stack deep", .{c.limits.stack >> 10}); + } fn resolve(c: *Ctx, name: []const u8) Error!Value { if (c.vars.get(name)) |x| return x; @@ -113,6 +127,7 @@ fn assignByName(c: *Ctx, act: *Activation, t: ast.Target, x: Value) Error!void { } fn exec(c: *Ctx, act: *Activation, body: []const ast.Stmt) Error!Flow { + try c.enter(); for (body) |st| switch (st) { .output => |pieces| for (pieces) |pc| switch (pc) { .text => |t| try c.out.literal(c.rt.a, t), @@ -383,7 +398,9 @@ fn binop(c: *Ctx, op: ast.BinOp, l: Value, r: Value) Error!Value { return .{ .list = try l.list.concat(rt.a, r.list) }; if (op == .mul and (l == .str or l == .list or r == .str or r == .list) and (isNum(l) or isNum(r)) and l != .float and r != .float) { const seq = if (isNum(l)) r else l; - const n = @max(0, intOf(if (isNum(l)) l else r)); + // an empty sequence stays empty at any count, and the loops below would still count it out + const empty = if (seq == .str) seq.str.s.len == 0 else seq.list.items.len == 0; + const n = if (empty) 0 else @max(0, intOf(if (isNum(l)) l else r)); if (seq == .str) { var out: v.Buffer = .{}; for (0..@intCast(n)) |_| try out.append(rt.a, seq.str); @@ -429,6 +446,7 @@ pub fn compare(c: *Ctx, op: ast.CmpOp, l: Value, r: Value) Error!bool { } pub fn eval(c: *Ctx, act: *Activation, e: *const ast.Expr) Error!Value { + try c.enter(); const rt = &c.rt; return switch (e.*) { .constant => |k| switch (k) { diff --git a/zig/src/core/template/limits_test.zig b/zig/src/core/template/limits_test.zig new file mode 100644 index 00000000000..e203a90208d --- /dev/null +++ b/zig/src/core/template/limits_test.zig @@ -0,0 +1,90 @@ +//! The limits a chat template is held to: nesting when it is parsed, then steps, recursion and memory when it renders. +const std = @import("std"); +const template = @import("template.zig"); + +/// `head`, then `open` and `close` each `n` times around `middle`, then `tail`. +fn nested(a: std.mem.Allocator, head: []const u8, open: []const u8, n: usize, middle: []const u8, close: []const u8, tail: []const u8) ![]u8 { + var out: std.ArrayList(u8) = .empty; + try out.appendSlice(a, head); + for (0..n) |_| try out.appendSlice(a, open); + try out.appendSlice(a, middle); + for (0..n) |_| try out.appendSlice(a, close); + try out.appendSlice(a, tail); + return out.items; +} + +/// Compiling `source` fails with a message that holds `part`. +fn refused(source: []const u8, part: []const u8) !void { + const a = std.testing.allocator; + var diag = template.Diag{}; + defer if (diag.msg.len > 0) a.free(diag.msg); + var t = template.compile(a, source, &diag) catch |err| { + try std.testing.expectEqual(error.TemplateFailed, err); + return std.testing.expect(std.mem.indexOf(u8, diag.msg, part) != null); + }; + t.deinit(); + return error.TestExpectedError; +} + +/// `source` rendered with no variables under `limits`: its text, or a failure whose message holds one of `part`. +fn rendered(source: []const u8, limits: ?template.Limits, want: []const u8, part: ?[]const u8) !void { + const a = std.testing.allocator; + var diag = template.Diag{}; + defer if (diag.msg.len > 0) a.free(diag.msg); + var t = try template.compile(a, source, &diag); + defer t.deinit(); + if (limits) |l| t.limits = l; + const out = template.render(a, &t, .null, .null, .null, true, &diag) catch |err| { + try std.testing.expectEqual(error.TemplateFailed, err); + var parts = std.mem.splitScalar(u8, part orelse return error.TestUnexpectedError, '|'); + while (parts.next()) |one| if (std.mem.indexOf(u8, diag.msg, one) != null) return; + std.debug.print("failed with: {s}\n", .{diag.msg}); + return error.TestUnexpectedError; + }; + defer a.free(out); + try std.testing.expect(part == null); + try std.testing.expectEqualStrings(want, out); +} + +test "security: a template nested past the parser's depth is refused before anything recurses through it" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const n = 100_000; + // each shape recursed once a level in the parser, or built a tree that tall for the scope pass and the renderer + const shapes = [_][5][]const u8{ + .{ "{{ ", "(", "1", ")", " }}" }, .{ "{{ ", "[", "1", "]", " }}" }, + .{ "{{ ", "not ", "x", "", " }}" }, .{ "{{ ", "- ", "1", "", " }}" }, + .{ "{{ 1", "", "", " if x else 1", " }}" }, .{ "", "{% if x %}", "", "{% endif %}", "" }, + .{ "{{ 1", "", "", " + 1", " }}" }, .{ "{{ x", "", "", ".a", " }}" }, + .{ "{{ x", "", "", "|string", " }}" }, .{ "{{ x", "", "", " or x", " }}" }, + .{ "{% set ", "(", "y", ")", " = 1 %}" }, .{ "{{ f", "", "", "(1)", " }}" }, + }; + for (shapes) |s| try refused(try nested(a, s[0], s[1], n, s[2], s[3], s[4]), "nested more than"); + // what a real template nests is far inside the limit and renders as before + try rendered(try nested(a, "{{ ", "(", 100, "1", ")", " }}"), null, "1", null); + try rendered(try nested(a, "{{ 1", "", 199, "", " + 1", " }}"), null, "200", null); + try rendered(try nested(a, "", "{% if true %}", 40, "deep", "{% endif %}", ""), null, "deep", null); +} + +test "security: a render stops at its steps, its recursion and its memory, and says which" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const loops = "{% for a in range(100000) %}{% for b in range(100000) %}{% endfor %}{% endfor %}done"; + try rendered(loops, .{ .steps = 100_000 }, "", "ran more than 100000 steps"); + try rendered("{% for a in range(1000) %}{% endfor %}done", .{ .steps = 100_000 }, "done", null); + // a macro that calls itself stops at the macro limit, or sooner at the stack's when its frames are large + try rendered("{% macro f(n) %}{{ f(n + 1) }}{% endmacro %}{{ f(0) }}", null, "", "recursion depth exceeded|of stack deep"); + // whatever the call is nested in, the stack limit stops it: each shape's frames are a different size + for ([_][2][]const u8{ .{ "not ", "" }, .{ "", " | string" }, .{ "[", "]" }, .{ "namespace(a=", ")" }, .{ "- ", "" } }) |shape| { + const source = try nested(a, "{% macro f(n) %}{{ ", shape[0], 40, "f(n + 1)", shape[1], " }}{% endmacro %}{{ f(0) }}"); + try rendered(source, null, "", "recursion depth exceeded|of stack deep"); + try rendered(source, .{ .stack = 256 << 10 }, "", "recursed more than 256 KiB of stack deep"); + } + const doubling = "{% set ns = namespace(s='0123456789abcdef') %}{% for i in range(40) %}{% set ns.s = ns.s ~ ns.s %}{% endfor %}{{ ns.s | length }}"; + try rendered(doubling, .{ .bytes = 8 << 20 }, "", "needs more than 8 MiB"); + try rendered("{% set s = '0123456789abcdef' * 1000 %}{{ s | length }}", .{ .bytes = 8 << 20 }, "16000", null); + // an empty sequence repeated is empty at once, whatever the count + try rendered("{{ ('' * 1000000000000000000) | length }}|{{ ([] * 1000000000000000000) | length }}", null, "0|0", null); +} diff --git a/zig/src/core/template/parser.zig b/zig/src/core/template/parser.zig index 2b66c788753..12714c59a13 100644 --- a/zig/src/core/template/parser.zig +++ b/zig/src/core/template/parser.zig @@ -6,11 +6,21 @@ const Allocator = std.mem.Allocator; pub const Error = error{ SyntaxError, OutOfMemory }; +/// How deep the parser recurses and how tall an expression stands: the scope pass and the renderer recurse as deep. +pub const max_depth = 256; + pub const Parser = struct { a: Allocator, toks: []const lex.Token, i: usize = 0, msg: []const u8 = "", + depth: u32 = 0, + + /// One more production open; the caller closes it with `defer p.depth -= 1`. + fn deeper(p: *Parser) Error!void { + if (p.depth >= max_depth) return p.fail("nested more than {d} levels deep", .{max_depth}); + p.depth += 1; + } fn fail(p: *Parser, comptime fmt: []const u8, args: anytype) Error { p.msg = std.fmt.allocPrint(p.a, "line {d}: " ++ fmt, .{p.cur().line} ++ args) catch return error.OutOfMemory; @@ -88,6 +98,8 @@ pub const Parser = struct { } fn subparse(p: *Parser, end: []const []const u8) Error![]ast.Stmt { + try p.deeper(); + defer p.depth -= 1; var body: std.ArrayList(ast.Stmt) = .empty; var pieces: std.ArrayList(ast.Piece) = .empty; while (true) { @@ -250,6 +262,8 @@ pub const Parser = struct { fn targetItem(p: *Parser, namespace: bool) Error!ast.Target { if (p.skipOp(.lparen)) { + try p.deeper(); + defer p.depth -= 1; const inner = try p.assignTarget(&.{}, false); try p.expectOp(.rparen); return inner; @@ -283,6 +297,8 @@ pub const Parser = struct { } fn condExpr(p: *Parser) Error!*ast.Expr { + try p.deeper(); + defer p.depth -= 1; var yes = try p.orExpr(); while (p.skipName("if")) { const t = try p.orExpr(); @@ -305,8 +321,10 @@ pub const Parser = struct { } fn notExpr(p: *Parser) Error!*ast.Expr { - if (p.skipName("not")) return p.expr(.{ .not_ = try p.notExpr() }); - return p.compare(); + if (!p.skipName("not")) return p.compare(); + try p.deeper(); + defer p.depth -= 1; + return p.expr(.{ .not_ = try p.notExpr() }); } fn compare(p: *Parser) Error!*ast.Expr { @@ -372,6 +390,8 @@ pub const Parser = struct { } fn unary(p: *Parser, with_filter: bool) Error!*ast.Expr { + try p.deeper(); + defer p.depth -= 1; var node = if (p.skipOp(.sub)) try p.expr(.{ .neg = try p.unary(false) }) else if (p.skipOp(.add)) try p.expr(.{ .pos = try p.unary(false) }) else try p.primary(); node = try p.postfix(node); if (with_filter) node = try p.filterExpr(node); diff --git a/zig/src/core/template/template.zig b/zig/src/core/template/template.zig index 58e558dc36c..a5c83bf5cc1 100644 --- a/zig/src/core/template/template.zig +++ b/zig/src/core/template/template.zig @@ -14,6 +14,7 @@ const Allocator = std.mem.Allocator; pub const Error = error{ TemplateFailed, OutOfMemory }; pub const Str = v.Str; +pub const Limits = eval.Limits; pub const Span = v.Span; /// Why rendering failed; `raised` marks the template's own raise_exception message. @@ -24,30 +25,41 @@ pub const Compiled = struct { /// Its child allocator also backs each render's scratch, so concurrent renders need it thread-safe. arena: std.heap.ArenaAllocator, tree: ast.Template, + limits: eval.Limits = .{}, pub fn deinit(t: *Compiled) void { t.arena.deinit(); } }; -fn checkNames(stmts: []const ast.Stmt, bad: *?[]const u8) void { +/// What a walk of the tree found: a filter or test the engine lacks, and an expression taller than the parser's depth. +const Check = struct { bad: ?[]const u8 = null, tall: bool = false, room: u32 = parser.max_depth }; + +fn checkNames(stmts: []const ast.Stmt, found: *Check) void { const W = struct { - fn expr(e: *const ast.Expr, out: *?[]const u8) void { + fn expr(e: *const ast.Expr, out: *Check) void { + // a loop builds `a + a + a ...` without recursing: its height is checked here, before any whole walk + if (out.room == 0) { + out.tall = true; + return; + } + out.room -= 1; + defer out.room += 1; switch (e.*) { .filter => |f| if (!access.among(f.name, &filters.implemented)) { - out.* = f.name; + out.bad = f.name; }, .is => |t| if (!access.among(t.name, &filters.tests)) { - out.* = t.name; + out.bad = t.name; }, else => {}, } ast.eachChild(error{}, e, out, child) catch |err| switch (err) {}; } - fn child(out: *?[]const u8, e: *ast.Expr) error{}!void { + fn child(out: *Check, e: *ast.Expr) error{}!void { expr(e, out); } - fn body(b: []const ast.Stmt, out: *?[]const u8) void { + fn body(b: []const ast.Stmt, out: *Check) void { for (b) |st| switch (st) { .output => |pieces| for (pieces) |pc| if (pc == .expr) expr(pc.expr, out), .for_ => |f| { @@ -70,16 +82,52 @@ fn checkNames(stmts: []const ast.Stmt, bad: *?[]const u8) void { .break_, .continue_ => {}, }; } - fn ifs(n: *const ast.If, out: *?[]const u8) void { + fn ifs(n: *const ast.If, out: *Check) void { expr(n.test_expr, out); body(n.body, out); for (n.elifs) |e| ifs(e, out); body(n.else_body, out); } }; - W.body(stmts, bad); + W.body(stmts, found); } +/// The render arena's source of memory, which gives out `left` bytes and then refuses. +const Capped = struct { + child: Allocator, + left: usize, + spent: bool = false, + + fn allocator(c: *Capped) Allocator { + return .{ .ptr = c, .vtable = &.{ .alloc = alloc, .resize = resize, .remap = remap, .free = free } }; + } + + fn take(c: *Capped, n: usize) bool { + if (n > c.left) c.spent = true else c.left -= n; + return !c.spent; + } + + fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 { + const c: *Capped = @ptrCast(@alignCast(ctx)); + return if (c.take(len)) c.child.rawAlloc(len, alignment, ra) else null; + } + + fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) bool { + const c: *Capped = @ptrCast(@alignCast(ctx)); + return c.take(new_len -| memory.len) and c.child.rawResize(memory, alignment, new_len, ra); + } + + fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) ?[*]u8 { + const c: *Capped = @ptrCast(@alignCast(ctx)); + return if (c.take(new_len -| memory.len)) c.child.rawRemap(memory, alignment, new_len, ra) else null; + } + + fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ra: usize) void { + const c: *Capped = @ptrCast(@alignCast(ctx)); + c.child.rawFree(memory, alignment, ra); + } +}; + fn report(gpa: Allocator, diag: ?*Diag, comptime fmt: []const u8, args: anytype) Error { if (diag) |d| d.msg = std.fmt.allocPrint(gpa, fmt, args) catch return error.OutOfMemory; return error.TemplateFailed; @@ -94,11 +142,12 @@ pub fn compile(gpa: Allocator, source: []const u8, diag: ?*Diag) Error!Compiled const tokens = lexer.tokenize(&lx) catch |err| return if (err == error.OutOfMemory) error.OutOfMemory else report(gpa, diag, "template syntax error: {s}", .{lx.msg}); var p = parser.Parser{ .a = a, .toks = tokens }; const body = p.parse() catch |err| return if (err == error.OutOfMemory) error.OutOfMemory else report(gpa, diag, "template syntax error: {s}", .{p.msg}); + var found: Check = .{}; + checkNames(body, &found); + if (found.tall) return report(gpa, diag, "template syntax error: an expression is nested more than {d} levels deep", .{parser.max_depth}); var z = scope.Analyzer{ .a = a }; const frame = z.run(body) catch |err| return if (err == error.OutOfMemory) error.OutOfMemory else report(gpa, diag, "template syntax error: {s}", .{z.msg}); - var bad: ?[]const u8 = null; - checkNames(body, &bad); - if (bad) |name| return report(gpa, diag, "template uses unsupported filter or test '{s}'", .{name}); + if (found.bad) |name| return report(gpa, diag, "template uses unsupported filter or test '{s}'", .{name}); return .{ .arena = arena, .tree = .{ .body = body, .frame = frame } }; } @@ -117,13 +166,15 @@ pub fn renderProvenance(gpa: Allocator, t: *const Compiled, messages: std.json.V /// Checkpoint-owned globals are template data; request context overrides them and retains its provenance. pub fn renderProvenanceGlobals(gpa: Allocator, t: *const Compiled, messages: std.json.Value, tools: std.json.Value, context: std.json.Value, template_globals: std.json.Value, add_generation_prompt: bool, diag: ?*Diag) Error!Str { // not on `gpa`: an arena there, such as a server request's, would keep the freed scratch until it is reset - var arena = std.heap.ArenaAllocator.init(t.arena.child_allocator); + var capped: Capped = .{ .child = t.arena.child_allocator, .left = t.limits.bytes }; + var arena = std.heap.ArenaAllocator.init(capped.allocator()); defer arena.deinit(); const a = arena.allocator(); var out: v.Buffer = .{}; var vars: v.Map = .empty; - var c = eval.Ctx{ .rt = .{ .a = a }, .vars = &vars, .out = &out }; + var c = eval.Ctx{ .rt = .{ .a = a }, .vars = &vars, .out = &out, .limits = t.limits }; run(&c, &vars, t, messages, tools, context, template_globals, add_generation_prompt) catch |err| { + if (err == error.OutOfMemory and capped.spent) return report(gpa, diag, "the chat template needs more than {d} MiB to render this request", .{t.limits.bytes >> 20}); if (err == error.OutOfMemory) return error.OutOfMemory; if (diag) |d| d.raised = c.rt.raised; return report(gpa, diag, "{s}", .{c.rt.msg}); @@ -163,6 +214,7 @@ pub fn renderChat(gpa: Allocator, source: []const u8, messages: std.json.Value, test { _ = @import("tests.zig"); + _ = @import("limits_test.zig"); _ = @import("unicode.zig"); _ = @import("value.zig"); _ = @import("provenance_test.zig"); From def39d21248b966cf2a03b0cdf418d76a8eb6bde Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:00:06 +0000 Subject: [PATCH 036/186] template: a step is a statement list or a loop filter's test, and the byte cap counts only what was given Every expression took a step and a stack check, which slowed renders of large requests by a few percent. The check now runs once a statement list, which every macro call and loop body passes through; an expression between two checks is at most the parser's 256 levels tall. A loop's filter runs no statement list, so each item it tests is now a step of its own, or a filtered loop inside another loop ran uncounted. The render arena asks to grow a block in place before it takes a new one, and the cap kept the bytes of every refused request; it gives them back. test-template-limits-host runs the limits tests on Linux and macOS. --- build.zig | 3 +++ zig/build/host_ci.zig | 2 ++ zig/src/core/template/eval.zig | 4 ++-- zig/src/core/template/limits_test.zig | 6 ++++++ zig/src/core/template/template.zig | 31 ++++++++++++++++++++++++--- 5 files changed, 41 insertions(+), 5 deletions(-) diff --git a/build.zig b/build.zig index 74f659e4d3b..c871eac8561 100644 --- a/build.zig +++ b/build.zig @@ -253,6 +253,9 @@ pub fn build(b: *std.Build) void { @import("zig/build/server.zig").cpuArtefacts(b, b.graph.host, optimize, build_options); @import("zig/build/mcdma_pair.zig").add(b, target); @import("zig/build/cache_index_timing.zig").step(b); + // the chat template's nesting, step, stack and memory limits (zig/src/core/template/limits_test.zig), on every host + const template_limits = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = b.graph.host, .optimize = .Debug, .link_libc = true }); + b.step("test-template-limits-host", "Chat template nesting, render steps, stack and memory limits without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = template_limits, .filters = &.{"security:"} })).step); sys_build.wire(b); @import("zig/build/checkpoint_config.zig").tests(b); @import("zig/build/checkpoint_config.zig").wire(b); diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index fac4a2aa517..95928ec7c31 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -85,6 +85,7 @@ const linux_fast = [_][]const u8{ "test-grammar-error-host", "test-info-config", "test-native-options-host", + "test-template-limits-host", }; const linux_hip = [_][]const u8{ @@ -305,6 +306,7 @@ const macos_fast = [_][]const u8{ "test-qwen27-prompt-floor-host", "test-qwen27-reply-cache-host", "test-qwen27-startup-host", + "test-template-limits-host", }; const macos_lanes = [_][]const u8{ diff --git a/zig/src/core/template/eval.zig b/zig/src/core/template/eval.zig index 18a2a2aaeeb..30c3badbb7c 100644 --- a/zig/src/core/template/eval.zig +++ b/zig/src/core/template/eval.zig @@ -31,7 +31,7 @@ pub const Ctx = struct { /// The stack where the render began; frames differ too much in size for a count of levels to bound it. base: usize = 0, - /// One more statement list or expression: a step of the render, taken only while its recursion has stack left. + /// One more statement list or loop-filter test: a step of the render, taken only while the stack has room. fn enter(c: *Ctx) Error!void { c.steps += 1; if (c.steps > c.limits.steps) return c.rt.fail("the chat template ran more than {d} steps", .{c.limits.steps}); @@ -178,6 +178,7 @@ fn execFor(c: *Ctx, act: *Activation, f: *const ast.For) Error!void { const tf = try enter(c, act, f.test_frame); var kept: std.ArrayList(Value) = .empty; for (items) |item| { + try c.enter(); // the filter runs no statement list, so each item it tests is a step of its own try assignByName(c, tf, f.target, item); if (v.truthy(try eval(c, tf, t))) try kept.append(c.rt.a, item); } @@ -446,7 +447,6 @@ pub fn compare(c: *Ctx, op: ast.CmpOp, l: Value, r: Value) Error!bool { } pub fn eval(c: *Ctx, act: *Activation, e: *const ast.Expr) Error!Value { - try c.enter(); const rt = &c.rt; return switch (e.*) { .constant => |k| switch (k) { diff --git a/zig/src/core/template/limits_test.zig b/zig/src/core/template/limits_test.zig index e203a90208d..930050470f0 100644 --- a/zig/src/core/template/limits_test.zig +++ b/zig/src/core/template/limits_test.zig @@ -88,3 +88,9 @@ test "security: a render stops at its steps, its recursion and its memory, and s // an empty sequence repeated is empty at once, whatever the count try rendered("{{ ('' * 1000000000000000000) | length }}|{{ ([] * 1000000000000000000) | length }}", null, "0|0", null); } + +test "security: a loop's filter takes a step for each item it tests" { + // the filter runs no statement list, so each item it tests is counted where it is tested + try rendered("{% for b in range(100000) if b < 0 %}{% endfor %}done", .{ .steps = 50_000 }, "", "ran more than 50000 steps"); + try rendered("{% for b in range(10) if b is even %}{{ b }}{% endfor %}", .{ .steps = 50_000 }, "02468", null); +} diff --git a/zig/src/core/template/template.zig b/zig/src/core/template/template.zig index a5c83bf5cc1..5a0d7734ede 100644 --- a/zig/src/core/template/template.zig +++ b/zig/src/core/template/template.zig @@ -109,17 +109,28 @@ const Capped = struct { fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 { const c: *Capped = @ptrCast(@alignCast(ctx)); - return if (c.take(len)) c.child.rawAlloc(len, alignment, ra) else null; + if (!c.take(len)) return null; + if (c.child.rawAlloc(len, alignment, ra)) |p| return p; + c.left += len; // the child refused, so nothing was spent + return null; } fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) bool { const c: *Capped = @ptrCast(@alignCast(ctx)); - return c.take(new_len -| memory.len) and c.child.rawResize(memory, alignment, new_len, ra); + const more = new_len -| memory.len; + if (!c.take(more)) return false; + if (c.child.rawResize(memory, alignment, new_len, ra)) return true; + c.left += more; // the arena asks to grow a block in place before it takes a new one, and is often refused + return false; } fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) ?[*]u8 { const c: *Capped = @ptrCast(@alignCast(ctx)); - return if (c.take(new_len -| memory.len)) c.child.rawRemap(memory, alignment, new_len, ra) else null; + const more = new_len -| memory.len; + if (!c.take(more)) return null; + if (c.child.rawRemap(memory, alignment, new_len, ra)) |p| return p; + c.left += more; + return null; } fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ra: usize) void { @@ -128,6 +139,20 @@ const Capped = struct { } }; +test "security: a render's byte cap counts what its allocator gave, not what it refused" { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 1, .resize_fail_index = 0 }); + var capped: Capped = .{ .child = failing.allocator(), .left = 4096 }; + const a = capped.allocator(); + const block = try a.alloc(u8, 1024); + defer a.free(block); + // the child grows nothing in place and gives no second block, so the cap has given nothing more + try std.testing.expect(!a.resize(block, 3072)); + try std.testing.expect(a.remap(block, 3072) == null); + try std.testing.expectError(error.OutOfMemory, a.alloc(u8, 2048)); + try std.testing.expectEqual(@as(usize, 3072), capped.left); + try std.testing.expect(!capped.spent); +} + fn report(gpa: Allocator, diag: ?*Diag, comptime fmt: []const u8, args: anytype) Error { if (diag) |d| d.msg = std.fmt.allocPrint(gpa, fmt, args) catch return error.OutOfMemory; return error.TemplateFailed; From 0591c4e85dca440749063d171d4305d65a6edb50 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:30:07 +0000 Subject: [PATCH 037/186] changelog: #3072 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index db707fab730..74b421459c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Chat templates: a template is held to a nesting depth at load, and its render to a step count, stack depth and scratch memory, so a crafted template can't hang or exhaust the server; real templates use a small fraction of each limit, render byte-identically, and render as fast as before (#3072). + - DeepSeek-V4.1 on Metal (experimental, opt in with `TF_EXPERIMENTAL=deepseek`): Macs now serve the release checkpoint natively with `tensorfold-dsv41` (`zig build dsv41-metal-serve`) for text prompts; image prompts come with #1722. Prompts and replies share rounds through the lane core with DSpark drafts, and the prompt cache reports `cached_tokens`. Rounds with three or more decoding streams decode plain (no draft pass, drafter ingest or copy windows), and when fewer streams decode again the drafter catches up from the target, so replies are unchanged. On an M3 Ultra and an M5 Ultra with 512 GB, drafts make one stream 1.21x to 1.50x faster than `--no-drafts` and two streams up to 1.13x faster, three and four streams are level, and `token_sha256` is equal in every cell (#1720, part of #1676). Thanks to @gilby, and to @3utsu for an independent M5 Ultra receipt. - Internal: `Admit` and `Ledger` move unchanged from `core/engine_api.zig` to `core/engine_admission.zig` and keep their names under `engine_api`, making room for the execution witness's hook; nothing a user sees changes (#2896, refs #2741). Thanks to @kotevcode. From c835f68dad483357b7f08e5f319df948107cb7ec Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 16:47:31 +0000 Subject: [PATCH 038/186] slide rollback: a fact graph is removed only when its own folders file names the folder A rollback unlinked every path its snapshot's graphs list named that read as a fact graph. The list sits in the model folder, so a folder copied or downloaded with a snapshot inside could name any fact graph of the user's. A --slide server now also writes GRAPH.folders beside the fact graph, listing the folders that learned from it, and a rollback removes a listed graph only when that file names the folder. The ones it leaves are printed. A graph taught by an earlier version has no such file until a server next teaches the folder. --- docs/living-weights.md | 6 ++- zig/src/core/originals.zig | 72 +++++++++++++++++++++++++-------- zig/src/core/originals_test.zig | 37 +++++++++++++++++ zig/src/server/slide_cli.zig | 2 + 4 files changed, 100 insertions(+), 17 deletions(-) diff --git a/docs/living-weights.md b/docs/living-weights.md index b42e2bda92c..de48a07078a 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -378,7 +378,11 @@ The rollback: - only then removes what the lessons added: lesson files (`living-weights-00001.safetensors` and on, or an older `living-weights.safetensors`), writes a lesson left cut short (`NAME.tmp` beside a kept or lesson file), and the fact graph files that `--slide-graph` named while the folder learned. A fact graph can live outside the folder; the - rollback removes it all the same, since it describes weights that are gone, and prints every path it removes. + rollback removes it all the same, since it describes weights that are gone, and prints every path it removes. It + removes one only when the `GRAPH.folders` file a `--slide` server writes beside the graph lists this folder, because + the snapshot's own list sits inside the model folder, where a downloaded repository could plant one. A graph it + leaves for that reason is printed too: one taught by an earlier version has no such file until a `--slide` server + next teaches the folder, so delete it by hand if you roll back first. - leaves everything else: files you added after the snapshot, files a lesson never writes (the tokenizer, templates, a download's `.cache`), and the snapshot itself, so you can teach the folder again and roll back again. diff --git a/zig/src/core/originals.zig b/zig/src/core/originals.zig index fe17d2d9b15..e3fd3a5a72c 100644 --- a/zig/src/core/originals.zig +++ b/zig/src/core/originals.zig @@ -11,6 +11,8 @@ pub const folder_name = ".tensorfold-original"; const partial_name = folder_name ++ ".partial"; // a snapshot being made; the next attempt starts it again const manifest_name = "SHA256SUMS"; // `shasum -a 256 -c SHA256SUMS` inside the snapshot checks it by hand const graphs_name = "graphs"; // the fact graph files --slide servers kept for the folder, one absolute path a line +/// Beside a fact graph: the folders that learned from it, one absolute path a line; a snapshot's own list can lie. +pub const taught_suffix = ".folders"; const put_back_suffix = ".tensorfold-rollback"; // a file being put back, renamed over its name once whole /// A file a lesson may write: a safetensors file, the shard index or config.json. @@ -46,6 +48,8 @@ pub const Restored = struct { cloned: usize = 0, removed: []const []const u8 = &.{}, graphs: []const []const u8 = &.{}, + /// Fact graphs the snapshot lists that stayed: nothing beside them says this folder learned from them. + left: []const []const u8 = &.{}, seconds: f64 = 0, }; @@ -117,7 +121,7 @@ fn keepFiles(io: Io, dir: []const u8, graph: ?[]const u8, o: Options, a: Allocat why.* = try std.fmt.allocPrint(a, text, .{ snap, name, manifest_name, dir }); return error.Damaged; } - if (graph) |g| try noteGraph(a, io, snap, g); + if (graph) |g| try noteGraph(a, io, snap, g, dir); var out: Kept = .{ .checked = !o.checked, .files = entries.len }; for (entries) |e| out.bytes += (try Io.Dir.cwd().statFile(io, try join(a, snap, e.name), .{})).size; out.seconds = since(io, t0); @@ -168,7 +172,7 @@ fn keepFiles(io: Io, dir: []const u8, graph: ?[]const u8, o: Options, a: Allocat try text.print(a, "{s} {s}\n", .{ &hex, name }); } try files.writeWhole(a, try join(a, part, manifest_name), text.items); - if (graph) |g| try noteGraph(a, io, part, g); + if (graph) |g| try noteGraph(a, io, part, g, dir); files.rename(try a.dupeSentinel(u8, part, 0), try a.dupeSentinel(u8, snap, 0)) catch return error.SnapshotFolder; files.syncDir(a, dir); out.seconds = since(io, t0); @@ -241,7 +245,7 @@ fn rollbackFiles(io: Io, dir: []const u8, o: Options, a: Allocator, why: *[]cons return error.NotRestored; } out.removed = try removeAdded(a, io, dir, entries); - out.graphs = try removeGraphs(a, io, snap); + try removeGraphs(a, io, snap, dir, &out); files.syncDir(a, dir); out.seconds = since(io, t0); return out; @@ -413,40 +417,76 @@ fn lessonFile(name: []const u8) bool { return true; } -/// The fact graph files the snapshot lists, removed where they still read as one (a JSON object with facts and next). -fn removeGraphs(a: Allocator, io: Io, snap: []const u8) ![]const []const u8 { +/// The fact graphs the snapshot lists, removed where they still read as one and their own folders file names `dir`. +fn removeGraphs(a: Allocator, io: Io, snap: []const u8, dir: []const u8, out: *Restored) !void { const at = try join(a, snap, graphs_name); const list = Io.Dir.cwd().readFileAlloc(io, at, a, .limited(1 << 20)) catch |e| switch (e) { - error.FileNotFound => return &.{}, + error.FileNotFound => return, else => return e, }; + const home = try Io.Dir.cwd().realPathFileAlloc(io, dir, a); var gone: std.ArrayList([]const u8) = .empty; + var left: std.ArrayList([]const u8) = .empty; var lines = std.mem.splitScalar(u8, list, '\n'); while (lines.next()) |path| { if (!std.fs.path.isAbsolute(path)) continue; const text = Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 26)) catch continue; const v = std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}) catch continue; if (v != .object or v.object.get("facts") == null or v.object.get("next") == null) continue; + // the snapshot's list sits in the folder, which a copied repository fills: the graph's side decides + const taught = try std.mem.concat(a, u8, &.{ path, taught_suffix }); + const folders = try lineList(a, io, taught); + if (!inLines(home, folders)) { + try left.append(a, path); + continue; + } if (platform.c.unlink(try a.dupeSentinel(u8, path, 0)) != 0) return error.RemoveFailed; try gone.append(a, path); + try writeLines(a, taught, folders, home); } - return gone.items; + out.graphs = gone.items; + out.left = left.items; } -/// `graph`, made absolute, listed once in the snapshot's graphs file. -fn noteGraph(a: Allocator, io: Io, snap: []const u8, graph: []const u8) !void { +/// A list file's text, empty when there is none. +fn lineList(a: Allocator, io: Io, file: []const u8) ![]const u8 { + return Io.Dir.cwd().readFileAlloc(io, file, a, .limited(1 << 20)) catch |e| switch (e) { + error.FileNotFound => "", + else => return e, + }; +} + +/// Whether `line` is one of the lines of `list`. +fn inLines(line: []const u8, list: []const u8) bool { + var lines = std.mem.splitScalar(u8, list, '\n'); + while (lines.next()) |l| if (std.mem.eql(u8, l, line)) return true; + return false; +} + +/// `list` without `drop`, written back; a list left empty takes its file with it. +fn writeLines(a: Allocator, file: []const u8, list: []const u8, drop: []const u8) !void { + var rest: std.ArrayList(u8) = .empty; + var lines = std.mem.splitScalar(u8, list, '\n'); + while (lines.next()) |l| if (l.len > 0 and !std.mem.eql(u8, l, drop)) try rest.print(a, "{s}\n", .{l}); + if (rest.items.len > 0) return files.writeWhole(a, file, rest.items); + _ = platform.c.unlink(try a.dupeSentinel(u8, file, 0)); +} + +/// `graph`, made absolute, listed once in the snapshot's graphs file, and `dir` once in the graph's folders file. +fn noteGraph(a: Allocator, io: Io, snap: []const u8, graph: []const u8, dir: []const u8) !void { const path = if (std.fs.path.isAbsolute(graph)) graph else blk: { const here = try Io.Dir.cwd().realPathFileAlloc(io, ".", a); break :blk try std.fs.path.join(a, &.{ here, graph }); }; const file = try join(a, snap, graphs_name); - const old = Io.Dir.cwd().readFileAlloc(io, file, a, .limited(1 << 20)) catch |e| switch (e) { - error.FileNotFound => "", - else => return e, - }; - var lines = std.mem.splitScalar(u8, old, '\n'); - while (lines.next()) |line| if (std.mem.eql(u8, line, path)) return; - try files.writeWhole(a, file, try std.mem.concat(a, u8, &.{ old, path, "\n" })); + const old = try lineList(a, io, file); + if (!inLines(path, old)) try files.writeWhole(a, file, try std.mem.concat(a, u8, &.{ old, path, "\n" })); + const home = try Io.Dir.cwd().realPathFileAlloc(io, dir, a); + const taught = try std.mem.concat(a, u8, &.{ path, taught_suffix }); + const folders = try lineList(a, io, taught); + if (inLines(home, folders)) return; + if (std.fs.path.dirname(path)) |parent| try Io.Dir.cwd().createDirPath(io, parent); + try files.writeWhole(a, taught, try std.mem.concat(a, u8, &.{ folders, home, "\n" })); } const teach_fix = "free some space, or teach a copy on APFS, Btrfs or XFS"; diff --git a/zig/src/core/originals_test.zig b/zig/src/core/originals_test.zig index 4cd2ebd822a..122b6665e27 100644 --- a/zig/src/core/originals_test.zig +++ b/zig/src/core/originals_test.zig @@ -265,11 +265,48 @@ test "a rollback puts back every kept file byte for byte, removes what lessons a try std.testing.expect(!exists(io, try path(a, dir, shard2 ++ ".tmp"))); try std.testing.expectEqual(@as(usize, 1), r.graphs.len); try std.testing.expect(std.fs.path.isAbsolute(r.graphs[0]) and !exists(io, graph)); + try std.testing.expect(r.left.len == 0 and !exists(io, try std.mem.concat(a, u8, &.{ graph, originals.taught_suffix }))); try std.testing.expect(exists(io, try path(a, dir, ".cache/huggingface/x.metadata"))); try std.testing.expect(!(try originals.keep(io, dir, null, o, a, &why)).made); // the next lesson reuses it } } +test "security: a rollback leaves a fact graph its snapshot lists unless the graph's own side names the folder" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try model(a, io, tmp); + var why: []const u8 = ""; + const o: originals.Options = .{ .clones = false, .free = plenty }; + _ = try originals.keep(io, dir, null, o, a, &why); + // two fact graphs of the user's: one no folder is tied to, one another folder learned from + const here = try Io.Dir.cwd().realPathFileAlloc(io, ".", a); + const mine = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/mine.json", .{ here, tmp.sub_path }); + const other = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/other.json", .{ here, tmp.sub_path }); + const others = try std.mem.concat(a, u8, &.{ other, originals.taught_suffix }); + const facts = "{\"facts\": [{\"id\": 1, \"text\": \"I like blue.\"}], \"next\": 2}"; + for ([_][]const u8{ mine, other }) |graph| try Io.Dir.cwd().writeFile(io, .{ .sub_path = graph, .data = facts }); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = others, .data = "/another/model/folder\n" }); + // the folder's side of the story, as a repository could ship it: its snapshot lists both + const list = try std.fmt.allocPrint(a, "{s}\n{s}\n", .{ mine, other }); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = try path(a, dir, originals.folder_name ++ "/graphs"), .data = list }); + try teach(a, io, dir); + const r = try originals.rollback(io, dir, o, a, &why); + try std.testing.expectEqual(@as(usize, 0), r.graphs.len); + try std.testing.expectEqual(@as(usize, 2), r.left.len); + try std.testing.expect(exists(io, mine) and exists(io, other)); + // once a server teaches the folder with `mine`, the graph's folders file names it and a rollback takes that one + _ = try originals.keep(io, dir, mine, o, a, &why); + const again = try originals.rollback(io, dir, o, a, &why); + try std.testing.expectEqual(@as(usize, 1), again.graphs.len); + try std.testing.expectEqualStrings(mine, again.graphs[0]); + try std.testing.expect(!exists(io, mine) and exists(io, other)); + try std.testing.expectEqualStrings("/another/model/folder\n", try read(a, io, others)); +} + test "a rollback removes only what TensorFold writes there: files the user added after the snapshot stay" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/server/slide_cli.zig b/zig/src/server/slide_cli.zig index 9029fb99f3a..26b1c9d9972 100644 --- a/zig/src/server/slide_cli.zig +++ b/zig/src/server/slide_cli.zig @@ -36,6 +36,8 @@ pub fn main(init: std.process.Init, argv: []const [:0]const u8) !u8 { try w.print(done, .{ dir, r.files, gb, r.cloned, api.originals.folder_name, r.seconds }); for (r.removed) |name| try w.print("tensorfold: removed {s}, which a lesson added\n", .{name}); for (r.graphs) |path| try w.print("tensorfold: removed the fact graph {s}\n", .{path}); + const left = "tensorfold: left the fact graph {s}: no {s} file beside it says this folder learned from it\n"; + for (r.left) |path| try w.print(left, .{ path, api.originals.taught_suffix }); try w.flush(); return 0; } From 8d588d68ba9cebabec10411c9597974a2804decd Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:00:15 +0000 Subject: [PATCH 039/186] slide rollback: a removed fact graph takes its folders file with it, and a folders file it cannot read leaves the graph After a rollback removed a graph, its folders file kept the other folders' names, so a graph made later at the same path took them on and an old folder's rollback could remove it; the file now goes with the graph. A folders file that cannot be read stopped a rollback that had already put the files back; it now counts as one that does not name the folder, the graph stays, and the rollback prints why. A folders file left behind after its graph is removed is printed too, not fatal. test-originals-host runs the snapshot and rollback tests on Linux and macOS. --- docs/living-weights.md | 7 ++-- zig/build/cuda_server_host.zig | 1 + zig/build/host_ci.zig | 2 ++ zig/src/core/originals.zig | 45 +++++++++++++----------- zig/src/core/originals_test.zig | 62 +++++++++++++++++++++++++++++++-- zig/src/server/slide_cli.zig | 9 +++-- 6 files changed, 99 insertions(+), 27 deletions(-) diff --git a/docs/living-weights.md b/docs/living-weights.md index de48a07078a..a679bd4cfe5 100644 --- a/docs/living-weights.md +++ b/docs/living-weights.md @@ -380,9 +380,10 @@ The rollback: fact graph files that `--slide-graph` named while the folder learned. A fact graph can live outside the folder; the rollback removes it all the same, since it describes weights that are gone, and prints every path it removes. It removes one only when the `GRAPH.folders` file a `--slide` server writes beside the graph lists this folder, because - the snapshot's own list sits inside the model folder, where a downloaded repository could plant one. A graph it - leaves for that reason is printed too: one taught by an earlier version has no such file until a `--slide` server - next teaches the folder, so delete it by hand if you roll back first. + the snapshot's own list sits inside the model folder, where a downloaded repository could plant one, and removes + that file with the graph. A graph it leaves for that reason, or because its `GRAPH.folders` file cannot be read, is + printed too: one taught by an earlier version has no such file until a `--slide` server next teaches the folder, so + delete it by hand if you roll back first. - leaves everything else: files you added after the snapshot, files a lesson never writes (the tokenizer, templates, a download's `.cache`), and the snapshot itself, so you can teach the folder again and roll back again. diff --git a/zig/build/cuda_server_host.zig b/zig/build/cuda_server_host.zig index f85c933597e..c77531a9b00 100644 --- a/zig/build/cuda_server_host.zig +++ b/zig/build/cuda_server_host.zig @@ -58,6 +58,7 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod b.step("test-lesson-gate-host", "Production server rounds and earlier-fact retention").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Learning persistence:"} })).step); b.step("test-slide-delete-host", "Single and multiple fact deletion without devices").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{ "delete", "a graph file", "allocation that fails anywhere in recording" } })).step); b.step("test-frame-cache-identity-host", "Image frame identities through the actual shared cache").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"Frame cache:"} })).step); + b.step("test-originals-host", "Model folder snapshots, holds and slide rollback without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"originals_test."} })).step); const media_gate = b.step("test-media-request-ownership-host", "Shared request ownership, image frame caches and media expansion on the CPU"); media_gate.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api })).step); media_gate.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{ "image", "video" } })).step); diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index 95928ec7c31..85206a8dd76 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -86,6 +86,7 @@ const linux_fast = [_][]const u8{ "test-info-config", "test-native-options-host", "test-template-limits-host", + "test-originals-host", }; const linux_hip = [_][]const u8{ @@ -301,6 +302,7 @@ const macos_fast = [_][]const u8{ "test-glm-disk-charge-host", "test-living-metal-host", "test-metal-formats-host", + "test-originals-host", "test-persistence-identity-host", "test-qwen27-bank-boundary-host", "test-qwen27-prompt-floor-host", diff --git a/zig/src/core/originals.zig b/zig/src/core/originals.zig index e3fd3a5a72c..b92f2cab1e1 100644 --- a/zig/src/core/originals.zig +++ b/zig/src/core/originals.zig @@ -28,6 +28,7 @@ pub const Options = struct { tamper: ?*const fn ([:0]const u8) void = null, // tests: a kept or put-back file changed before it is checked uncloned: ?*const fn ([:0]const u8) bool = null, // tests: the files this file system will not clone checked: bool = false, // this process made or checked the snapshot already: a later lesson reuses it unread + graph_gone: ?*const fn ([]const u8) void = null, // tests: a removed fact graph's folders file, before it goes }; /// A snapshot made now (`made`) or found (`checked`: hashed now), its files, bytes, clones and seconds; `taught` too. @@ -47,12 +48,16 @@ pub const Restored = struct { bytes: u64 = 0, cloned: usize = 0, removed: []const []const u8 = &.{}, - graphs: []const []const u8 = &.{}, - /// Fact graphs the snapshot lists that stayed: nothing beside them says this folder learned from them. - left: []const []const u8 = &.{}, + /// Fact graphs removed; `err` is why the folders file beside one stayed. + graphs: []const Graph = &.{}, + /// Fact graphs the snapshot lists that stayed: nothing beside them names this folder, or (`err`) it was unreadable. + left: []const Graph = &.{}, seconds: f64 = 0, }; +/// A fact graph a rollback removed or left, and the error from the folders file beside it. +pub const Graph = struct { path: []const u8, err: ?anyerror = null }; + /// A process's hold on a model folder, gone when it is released or the process ends. pub const Hold = struct { fd: platform.c.fd_t, @@ -245,7 +250,7 @@ fn rollbackFiles(io: Io, dir: []const u8, o: Options, a: Allocator, why: *[]cons return error.NotRestored; } out.removed = try removeAdded(a, io, dir, entries); - try removeGraphs(a, io, snap, dir, &out); + try removeGraphs(a, io, snap, dir, o, &out); files.syncDir(a, dir); out.seconds = since(io, t0); return out; @@ -418,15 +423,15 @@ fn lessonFile(name: []const u8) bool { } /// The fact graphs the snapshot lists, removed where they still read as one and their own folders file names `dir`. -fn removeGraphs(a: Allocator, io: Io, snap: []const u8, dir: []const u8, out: *Restored) !void { +fn removeGraphs(a: Allocator, io: Io, snap: []const u8, dir: []const u8, o: Options, out: *Restored) !void { const at = try join(a, snap, graphs_name); const list = Io.Dir.cwd().readFileAlloc(io, at, a, .limited(1 << 20)) catch |e| switch (e) { error.FileNotFound => return, else => return e, }; const home = try Io.Dir.cwd().realPathFileAlloc(io, dir, a); - var gone: std.ArrayList([]const u8) = .empty; - var left: std.ArrayList([]const u8) = .empty; + var gone: std.ArrayList(Graph) = .empty; + var left: std.ArrayList(Graph) = .empty; var lines = std.mem.splitScalar(u8, list, '\n'); while (lines.next()) |path| { if (!std.fs.path.isAbsolute(path)) continue; @@ -435,14 +440,23 @@ fn removeGraphs(a: Allocator, io: Io, snap: []const u8, dir: []const u8, out: *R if (v != .object or v.object.get("facts") == null or v.object.get("next") == null) continue; // the snapshot's list sits in the folder, which a copied repository fills: the graph's side decides const taught = try std.mem.concat(a, u8, &.{ path, taught_suffix }); - const folders = try lineList(a, io, taught); + // the files are back already: a folders file that cannot be read names no folder, and the graph stays + const folders = lineList(a, io, taught) catch |e| { + try left.append(a, .{ .path = path, .err = e }); + continue; + }; if (!inLines(home, folders)) { - try left.append(a, path); + try left.append(a, .{ .path = path }); continue; } if (platform.c.unlink(try a.dupeSentinel(u8, path, 0)) != 0) return error.RemoveFailed; - try gone.append(a, path); - try writeLines(a, taught, folders, home); + if (o.graph_gone) |hook| hook(taught); + // its folders go with it, or a graph made later at this path would take them on; one left is only reported + const stayed: ?anyerror = if (Io.Dir.cwd().deleteFile(io, taught)) null else |e| switch (e) { + error.FileNotFound => null, + else => e, + }; + try gone.append(a, .{ .path = path, .err = stayed }); } out.graphs = gone.items; out.left = left.items; @@ -463,15 +477,6 @@ fn inLines(line: []const u8, list: []const u8) bool { return false; } -/// `list` without `drop`, written back; a list left empty takes its file with it. -fn writeLines(a: Allocator, file: []const u8, list: []const u8, drop: []const u8) !void { - var rest: std.ArrayList(u8) = .empty; - var lines = std.mem.splitScalar(u8, list, '\n'); - while (lines.next()) |l| if (l.len > 0 and !std.mem.eql(u8, l, drop)) try rest.print(a, "{s}\n", .{l}); - if (rest.items.len > 0) return files.writeWhole(a, file, rest.items); - _ = platform.c.unlink(try a.dupeSentinel(u8, file, 0)); -} - /// `graph`, made absolute, listed once in the snapshot's graphs file, and `dir` once in the graph's folders file. fn noteGraph(a: Allocator, io: Io, snap: []const u8, graph: []const u8, dir: []const u8) !void { const path = if (std.fs.path.isAbsolute(graph)) graph else blk: { diff --git a/zig/src/core/originals_test.zig b/zig/src/core/originals_test.zig index 122b6665e27..fb5147e28d8 100644 --- a/zig/src/core/originals_test.zig +++ b/zig/src/core/originals_test.zig @@ -264,7 +264,7 @@ test "a rollback puts back every kept file byte for byte, removes what lessons a try std.testing.expect(!exists(io, try path(a, dir, "living-weights-00001.safetensors"))); try std.testing.expect(!exists(io, try path(a, dir, shard2 ++ ".tmp"))); try std.testing.expectEqual(@as(usize, 1), r.graphs.len); - try std.testing.expect(std.fs.path.isAbsolute(r.graphs[0]) and !exists(io, graph)); + try std.testing.expect(std.fs.path.isAbsolute(r.graphs[0].path) and r.graphs[0].err == null and !exists(io, graph)); try std.testing.expect(r.left.len == 0 and !exists(io, try std.mem.concat(a, u8, &.{ graph, originals.taught_suffix }))); try std.testing.expect(exists(io, try path(a, dir, ".cache/huggingface/x.metadata"))); try std.testing.expect(!(try originals.keep(io, dir, null, o, a, &why)).made); // the next lesson reuses it @@ -302,11 +302,69 @@ test "security: a rollback leaves a fact graph its snapshot lists unless the gra _ = try originals.keep(io, dir, mine, o, a, &why); const again = try originals.rollback(io, dir, o, a, &why); try std.testing.expectEqual(@as(usize, 1), again.graphs.len); - try std.testing.expectEqualStrings(mine, again.graphs[0]); + try std.testing.expectEqualStrings(mine, again.graphs[0].path); try std.testing.expect(!exists(io, mine) and exists(io, other)); try std.testing.expectEqualStrings("/another/model/folder\n", try read(a, io, others)); } +/// A removed graph's folders file turned into a folder, which no unlink of a file removes. +fn stuck(taught: []const u8) void { + Io.Dir.cwd().deleteFile(std.testing.io, taught) catch {}; + Io.Dir.cwd().createDir(std.testing.io, taught, .default_dir) catch {}; +} + +test "security: a removed fact graph takes its folders file along, and one that stays is reported, not fatal" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try model(a, io, tmp); + var why: []const u8 = ""; + const here = try Io.Dir.cwd().realPathFileAlloc(io, ".", a); + const graph = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/shared.json", .{ here, tmp.sub_path }); + const taught = try std.mem.concat(a, u8, &.{ graph, originals.taught_suffix }); + const facts = "{\"facts\": [{\"id\": 1, \"text\": \"I like blue.\"}], \"next\": 2}"; + // another folder learned from this graph too: a graph made here later must not inherit its name + try Io.Dir.cwd().writeFile(io, .{ .sub_path = taught, .data = "/another/model/folder\n" }); + for ([_]?*const fn ([]const u8) void{ null, stuck }) |hook| { + const o: originals.Options = .{ .clones = false, .free = plenty, .graph_gone = hook }; + _ = try originals.keep(io, dir, graph, o, a, &why); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = graph, .data = facts }); + try teach(a, io, dir); + const r = try originals.rollback(io, dir, o, a, &why); + try std.testing.expectEqual(@as(usize, kept_files), r.files); + try std.testing.expect(r.graphs.len == 1 and r.left.len == 0 and !exists(io, graph)); + // a folders file that stays is the rollback's to report, after the files are back + try std.testing.expectEqual(hook != null, r.graphs[0].err != null); + try std.testing.expectEqual(hook != null, exists(io, taught)); + } +} + +test "security: a folders file that cannot be read leaves its fact graph, and the rollback still finishes" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try model(a, io, tmp); + var why: []const u8 = ""; + const o: originals.Options = .{ .clones = false, .free = plenty }; + _ = try originals.keep(io, dir, null, o, a, &why); + const here = try Io.Dir.cwd().realPathFileAlloc(io, ".", a); + const graph = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}/mine.json", .{ here, tmp.sub_path }); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = graph, .data = "{\"facts\": [], \"next\": 1}" }); + try Io.Dir.cwd().createDir(io, try std.mem.concat(a, u8, &.{ graph, originals.taught_suffix }), .default_dir); + const list = try path(a, dir, originals.folder_name ++ "/graphs"); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = list, .data = try std.mem.concat(a, u8, &.{ graph, "\n" }) }); + try teach(a, io, dir); + const r = try originals.rollback(io, dir, o, a, &why); + try std.testing.expectEqual(@as(usize, kept_files), r.files); + try std.testing.expect(r.graphs.len == 0 and r.left.len == 1 and r.left[0].err != null and exists(io, graph)); +} + test "a rollback removes only what TensorFold writes there: files the user added after the snapshot stay" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/server/slide_cli.zig b/zig/src/server/slide_cli.zig index 26b1c9d9972..ef7d789f983 100644 --- a/zig/src/server/slide_cli.zig +++ b/zig/src/server/slide_cli.zig @@ -35,9 +35,14 @@ pub fn main(init: std.process.Init, argv: []const [:0]const u8) !u8 { const gb = @as(f64, @floatFromInt(r.bytes)) / 1e9; try w.print(done, .{ dir, r.files, gb, r.cloned, api.originals.folder_name, r.seconds }); for (r.removed) |name| try w.print("tensorfold: removed {s}, which a lesson added\n", .{name}); - for (r.graphs) |path| try w.print("tensorfold: removed the fact graph {s}\n", .{path}); + const suffix = api.originals.taught_suffix; + for (r.graphs) |g| { + try w.print("tensorfold: removed the fact graph {s}\n", .{g.path}); + if (g.err) |e| try w.print("tensorfold: could not remove {s}{s} ({s}): delete it by hand\n", .{ g.path, suffix, @errorName(e) }); + } const left = "tensorfold: left the fact graph {s}: no {s} file beside it says this folder learned from it\n"; - for (r.left) |path| try w.print(left, .{ path, api.originals.taught_suffix }); + const unread = "tensorfold: left the fact graph {s}: its {s} file could not be read ({s})\n"; + for (r.left) |g| if (g.err) |e| try w.print(unread, .{ g.path, suffix, @errorName(e) }) else try w.print(left, .{ g.path, suffix }); try w.flush(); return 0; } From 13d7ce05229f6a31d7f8e51c2e3aac2693b8f815 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 09:44:21 +0000 Subject: [PATCH 040/186] changelog: #3073 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 74b421459c0..e65663831a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). + - Chat templates: a template is held to a nesting depth at load, and its render to a step count, stack depth and scratch memory, so a crafted template can't hang or exhaust the server; real templates use a small fraction of each limit, render byte-identically, and render as fast as before (#3072). - DeepSeek-V4.1 on Metal (experimental, opt in with `TF_EXPERIMENTAL=deepseek`): Macs now serve the release checkpoint natively with `tensorfold-dsv41` (`zig build dsv41-metal-serve`) for text prompts; image prompts come with #1722. Prompts and replies share rounds through the lane core with DSpark drafts, and the prompt cache reports `cached_tokens`. Rounds with three or more decoding streams decode plain (no draft pass, drafter ingest or copy windows), and when fewer streams decode again the drafter catches up from the target, so replies are unchanged. On an M3 Ultra and an M5 Ultra with 512 GB, drafts make one stream 1.21x to 1.50x faster than `--no-drafts` and two streams up to 1.13x faster, three and four streams are level, and `token_sha256` is equal in every cell (#1720, part of #1676). Thanks to @gilby, and to @3utsu for an independent M5 Ultra receipt. From 07d5fe726deddb1dafd449d6fb52b495e8d17f83 Mon Sep 17 00:00:00 2001 From: gilby Date: Tue, 10 Nov 2026 03:36:26 +0100 Subject: [PATCH 041/186] flashnext hip: a window round's plan buffers A stream's words, a slot's descriptor and the rows its keep replays, one Plan a cache slot. No kernel launches yet. --- zig/src/families/flashnext/hip/gdn_replay.zig | 271 ++++++++++++++++++ 1 file changed, 271 insertions(+) create mode 100644 zig/src/families/flashnext/hip/gdn_replay.zig diff --git a/zig/src/families/flashnext/hip/gdn_replay.zig b/zig/src/families/flashnext/hip/gdn_replay.zig new file mode 100644 index 00000000000..35ca1335215 --- /dev/null +++ b/zig/src/families/flashnext/hip/gdn_replay.zig @@ -0,0 +1,271 @@ +//! A verify window's DeltaNet without a snapshot a row: the round runs from a plan, the keep replays it. +const std = @import("std"); +const hip = @import("hip"); + +const pops = hip.plan_ops; +const cache = @import("cache.zig"); +const forward = @import("forward.zig"); +const model = @import("model.zig"); + +const Ops = hip.ops.Ops; +const Tensor = hip.ops.Tensor; +const take = forward.take; + +/// Rows one window holds at most: the pending token and its drafts. +pub const window_rows = cache.max_window; + +/// Slots a window's plan runs: the window, then an empty slot that stops its wave writing the state back. +const slots = 2; + +/// Whether this model's shapes are the plan kernels': key dim 128, eight value rows a wave, gated heads. +pub fn supported(d: model.Dims, layers: usize) bool { + if (d.key_dim != 128) return false; + if (d.value_dim % 8 != 0 or d.key_heads < 1 or d.value_heads % d.key_heads != 0) return false; + if (d.conv < 1 or d.conv > 8) return false; + return layers > 0; +} + +/// The plan's words in four-byte units: a slot's first row, its rows, its kept rows, its descriptor pointer. +pub const Layout = struct { + first: usize, + count: usize, + keep: usize, + ptr: usize, + snaps: usize, + words: usize, + + pub fn of(layers: usize) Layout { + var l: Layout = undefined; + l.first = 0; + l.count = slots; + l.keep = 2 * slots; + l.ptr = std.mem.alignForward(usize, 3 * slots, 4); // eight-byte pointers start on an even word + l.snaps = l.ptr + 2 * slots; + l.words = l.snaps + 4 * layers; + return l; + } + + /// Words a slot's descriptor takes: cached positions, the last kept final row, two addresses a layer. + pub fn desc(layers: usize) usize { + return 2 * (2 + 2 * layers); + } + + /// Where the descriptors start (eight-byte aligned). + fn start(layers: usize) usize { + return std.mem.alignForward(usize, of(layers).words, 2); + } + + /// Every word of the plan: its control words, then a descriptor a slot. + pub fn total(layers: usize) usize { + return start(layers) + slots * desc(layers); + } +}; + +/// One stream's window plan: the round's words, its descriptors, and the rows its keep reads again. +pub const Plan = struct { + r: *const hip.Runtime, + layers: usize, + /// Rows the buffers hold at most: a window. + room: usize, + key_heads: usize, + value_heads: usize, + key_dim: usize, + value_dim: usize, + conv: usize, + channels: usize, + kept_stride: usize, + samples_stride: usize, + /// The words and descriptors, and the device copy a round and a keep read. + host: hip.HostBuffer, + dev: hip.DeviceBuffer, + /// A layer's k, v, gate and beta over the window's rows, section-major: what a keep replays from. + kept: hip.DeviceBuffer, + /// A layer's conv input rows over the window, section-major: where a keep takes its conv tail. + samples: hip.DeviceBuffer, + /// Rows the window last ran. + rows: usize = 0, + + pub fn init(d: *const hip.Runtime, dims: model.Dims) !Plan { + if (!supported(dims, dims.layers)) return error.UnsupportedPlan; + var host = try hip.HostBuffer.alloc(d, 4 * Layout.total(dims.layers)); + errdefer host.free(); + var dev = try hip.DeviceBuffer.alloc(d, host.bytes.len); + errdefer dev.free(); + const width = dims.key_heads * dims.key_dim + dims.value_heads * dims.value_dim + 2 * dims.value_heads; + var kept = try hip.DeviceBuffer.alloc(d, 4 * window_rows * width * dims.layers); + errdefer kept.free(); + var samples = try hip.DeviceBuffer.alloc(d, 4 * window_rows * dims.convWidth() * dims.layers); + errdefer samples.free(); + @memset(host.bytes, 0); + return .{ + .r = d, + .layers = dims.layers, + .room = window_rows, + .key_heads = dims.key_heads, + .value_heads = dims.value_heads, + .key_dim = dims.key_dim, + .value_dim = dims.value_dim, + .conv = dims.conv, + .channels = dims.convWidth(), + .kept_stride = 4 * window_rows * width, + .samples_stride = 4 * window_rows * dims.convWidth(), + .host = host, + .dev = dev, + .kept = kept, + .samples = samples, + }; + } + + pub fn deinit(p: *Plan) void { + p.samples.free(); + p.kept.free(); + p.dev.free(); + p.host.free(); + } + + /// Device bytes a stream's plan takes, whether or not one is made: the admission counts it the same either way. + pub fn bytes(dims: model.Dims) usize { + const width = dims.key_heads * dims.key_dim + dims.value_heads * dims.value_dim + 2 * dims.value_heads; + const words = 4 * Layout.total(dims.layers); + return words + 4 * window_rows * width * dims.layers + 4 * window_rows * dims.convWidth() * dims.layers; + } + + pub fn layout(p: *const Plan) Layout { + return Layout.of(p.layers); + } + + /// A slot's descriptor words: [0] cached positions, [1] the last kept final row (0: none), then a layer's pair. + fn descriptor(p: *Plan, slot: usize) []u64 { + const at = Layout.start(p.layers) + slot * Layout.desc(p.layers); + return p.host.slice(u64)[at / 2 ..][0 .. Layout.desc(p.layers) / 2]; + } + + fn descAt(p: *const Plan, slot: usize) u64 { + return p.dev.base() + 4 * (Layout.start(p.layers) + slot * Layout.desc(p.layers)); + } + + /// Linear layer `layer`'s section-major k, v, gate and beta over the window's row `at`. + pub fn keptOf(p: *const Plan, layer: usize, at: usize) struct { k: u64, v: u64, gate: u64, beta: u64 } { + const kw = p.key_heads * p.key_dim; + const vw = p.value_heads * p.value_dim; + const rows = p.rows; + const base = p.kept.base() + layer * p.kept_stride; + const k = base + 4 * at * kw; + const v = base + 4 * rows * kw + 4 * at * vw; + const gate = base + 4 * rows * (kw + vw) + 4 * at * p.value_heads; + return .{ .k = k, .v = v, .gate = gate, .beta = gate + 4 * rows * p.value_heads }; + } + + /// The conv samples of layer `layer` (the window's `channels` a row). + pub fn samplesOf(p: *const Plan, layer: usize) u64 { + return p.samples.base() + layer * p.samples_stride; + } + + /// The kept rows a round and a keep read. + pub fn keepWords(p: *const Plan) u64 { + return p.dev.base() + 4 * p.layout().keep; + } + + /// The window's words and descriptor for caches `c` over `n` rows, on the host: the copy goes up with `send`. + pub fn fill(p: *Plan, c: *const cache.Caches, n: usize) !void { + if (n == 0 or n > p.room) return error.TooManyRows; + const l = p.layout(); + const w = p.host.slice(u32); + const ptr = p.host.slice(u64); + const desc = p.descriptor(0); + for (0..p.layers) |i| { + // a sparse layer keeps its descriptor pair, conv samples and replay inputs zero: the kernels skip it + desc[2 + 2 * i ..][0..2].* = if (c.linear[i]) .{ c.linearAt(i).conv, c.linearAt(i).state } else .{ 0, 0 }; + ptr[l.snaps / 2 + 2 * i] = if (c.linear[i]) p.samplesOf(i) else 0; + ptr[l.snaps / 2 + 2 * i + 1] = if (c.linear[i]) p.kept.base() + i * p.kept_stride else 0; + } + // a round keeps no final row: a keep takes the conv tail and the state alone + desc[0] = 0; + desc[1] = 0; + @memset(p.descriptor(1), 0); + w[l.first] = 0; + w[l.count] = @intCast(n); + w[l.keep] = std.math.maxInt(u32); + ptr[l.ptr / 2] = p.descAt(0); + // the empty slot runs nothing and keeps nothing: the kernels stop at its zero row count + w[l.first + 1] = 0; + w[l.count + 1] = 0; + w[l.keep + 1] = std.math.maxInt(u32); + ptr[l.ptr / 2 + 1] = p.descAt(1); + p.rows = n; + } + + /// The words up to the device: a graph holds the copy, so its every replay sends the window's own words. + pub fn send(p: *Plan, s: hip.abi.Stream) !void { + try hip.raw.upload(p.dev, 0, p.host.bytes, s); + } + + /// The window's first `n` rows kept: the conv tail from their samples, the state replayed through them. + pub fn keep(p: *Plan, o: Ops, n: usize) !void { + // a window of one row left the state and its conv tail as the round wrote them + if (p.rows <= 1 or n == 0) return; + p.host.slice(u32)[p.layout().keep] = @bitCast(@as(i32, @intCast(n))); + try p.send(o.stream); + try pops.keep(o, p.ref(), .{ + .keep = p.keepWords(), + .hidden = 0, + .hidden_words = 0, + .channels = p.channels, + .taps = p.conv - 1, + .layers = p.layers, + }); + try pops.gdnReplay(o, p.keepWords(), p.ref(), p.key_heads, p.value_heads, p.key_dim, p.value_dim, p.layers); + } + + /// The words as launches take them: the window's rows from its first. + pub fn ref(p: *const Plan) pops.Plan { + const l = p.layout(); + const at = p.dev.base(); + return .{ + .args = .{ + .first = at + 4 * l.first, + .count = at + 4 * l.count, + .desc = at + 4 * l.ptr, + .snaps = at + 4 * l.snaps, + }, + .rows = p.rows, + .slots = slots, + }; + } +}; + +/// Every cache slot's window plan: a slot's keep reads its own rows, which the next round over it has not yet written. +pub const Slots = struct { + gpa: std.mem.Allocator, + plans: []Plan, + + pub fn init(gpa: std.mem.Allocator, d: *const hip.Runtime, dims: model.Dims, streams: usize) !Slots { + const plans = try gpa.alloc(Plan, @max(streams, 1)); + var made: usize = 0; + errdefer { + for (plans[0..made]) |*p| p.deinit(); + gpa.free(plans); + } + for (plans) |*p| { + p.* = try Plan.init(d, dims); + made += 1; + } + return .{ .gpa = gpa, .plans = plans }; + } + + pub fn deinit(s: *Slots) void { + for (s.plans) |*p| p.deinit(); + s.gpa.free(s.plans); + } + + /// Cache slot `id`'s plan (null: a slot this rank never leased, so nothing can round or keep it). + pub fn planFor(s: *Slots, id: u32) ?*Plan { + if (id >= s.plans.len) return null; + return &s.plans[id]; + } + + /// Device bytes the slots take. + pub fn bytes(dims: model.Dims, streams: usize) usize { + return Plan.bytes(dims) * @max(streams, 1); + } +}; From 18ec707370b21f79df1e7fab57737b721840c446 Mon Sep 17 00:00:00 2001 From: gilby Date: Tue, 10 Nov 2026 03:36:39 +0100 Subject: [PATCH 042/186] flashnext hip: a verify window's DeltaNet from its plan, the keep replaying the kept rows A window's conv split, gates and delta rule read its rows from the plan, and the keep rebuilds the conv tail from the kept inputs and re-runs the scan over the kept rows, so the state is serial decode's while the round writes no snapshot a row. --- zig/src/families/flashnext/hip/gdn_replay.zig | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/zig/src/families/flashnext/hip/gdn_replay.zig b/zig/src/families/flashnext/hip/gdn_replay.zig index 35ca1335215..8e8c6cd21f7 100644 --- a/zig/src/families/flashnext/hip/gdn_replay.zig +++ b/zig/src/families/flashnext/hip/gdn_replay.zig @@ -269,3 +269,75 @@ pub const Slots = struct { return Plan.bytes(dims) * @max(streams, 1); } }; + +/// A window's rows through the linear attention from its plan: a stream one conv split and delta rule, the norm after. +pub fn linear( + o: Ops, + m: *const model.Model, + l: model.Linear, + segs: []const forward.Seg, + index: usize, + qkv: Tensor, + z: Tensor, + a: Tensor, + b: Tensor, + rows: usize, + s: *Slots, +) !Tensor { + const d = m.dims; + const ch = d.convWidth(); + const kw = d.key_heads * d.key_dim; + const vw = d.value_heads * d.value_dim; + const y = try o.arena.of(f32, rows * vw); + var off: usize = 0; + for (segs) |g| { + const p = s.planFor(g.c.id) orelse return error.CachesNotPlanned; + const n = g.tokens.len; + const qn = try o.arena.of(f32, n * kw); + const kept = p.keptOf(index, 0); + try pops.convSplit( + o, + forward.at(qkv, off * ch), + l.conv, + qn, + kept.k, + kept.v, + p.ref(), + index, + ch, + d.conv, + kw, + vw, + .{ .q = m.qk.q, .k = m.qk.k, .eps = m.qk.eps }, + .{ + .a = forward.at(a, off * d.value_heads), + .b = forward.at(b, off * d.value_heads), + .a_log = l.a_log, + .dt_bias = l.dt_bias, + .gate = kept.gate, + .beta = kept.beta, + .count = n * d.value_heads, + .heads = d.value_heads, + }, + ); + try pops.gatedDelta( + o, + qn, + kept.k, + kept.v, + kept.gate, + kept.beta, + y + off * vw * 4, + p.ref(), + index, + d.key_heads, + d.value_heads, + d.key_dim, + d.value_dim, + ); + off += n; + } + const out = try take(o, m.act, rows * vw); + try o.gnormOut(y, l.norm, z, out, rows * d.value_heads, d.value_dim, d.eps, d.gate); + return o.project(out, l.out, rows, l.out.partial); +} From 43c82e4ccce271d33d0b4021d4d04c2152921691 Mon Sep 17 00:00:00 2001 From: gilby Date: Tue, 10 Nov 2026 03:37:17 +0100 Subject: [PATCH 043/186] flashnext hip: window rounds and their keeps on the plan path The engine makes a plan a cache slot where the shapes are the plan kernels' and sends its words inside a round's graph, a window's linear layer goes through the plan, and a keep replays the kept rows where the plan took it and copies the snapshots where it did not. The plan's bytes count in a stream's admission. --- zig/src/families/flashnext/hip/engine.zig | 64 ++++++++++++++++------ zig/src/families/flashnext/hip/forward.zig | 5 ++ zig/src/families/flashnext/hip/mixers.zig | 3 + zig/src/families/flashnext/hip/model.zig | 2 + 4 files changed, 56 insertions(+), 18 deletions(-) diff --git a/zig/src/families/flashnext/hip/engine.zig b/zig/src/families/flashnext/hip/engine.zig index 6f597278648..9f5029b3dfe 100644 --- a/zig/src/families/flashnext/hip/engine.zig +++ b/zig/src/families/flashnext/hip/engine.zig @@ -12,6 +12,7 @@ const forward = @import("forward.zig"); const head_mod = @import("head.zig"); const mixers = @import("mixers.zig"); const index = @import("../index.zig"); +const gdn = @import("gdn_replay.zig"); pub const Options = struct { device: c_int = 0, @@ -88,6 +89,8 @@ pub const Engine = struct { stopped: bool, /// The chunk the options asked for when the engine took their fallback (null: it took what they asked). narrowed: ?usize = null, + /// The cache slots' DeltaNet plans (null: a window keeps its state a row, from the snapshots). + gdn: ?*gdn.Slots = null, /// `load`, this rank's failure heard by the other ranks at the step they wait at. pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !*Engine { @@ -181,6 +184,14 @@ pub const Engine = struct { e.graphs = Graphs.init(gpa); e.chosen = null; e.pool = .{}; + // a window runs from the plan where its shapes are the plan kernels': else a snapshot a row stays + if (gdn.supported(e.weights.model.dims, e.cfg.layers)) { + const slots = try gpa.create(gdn.Slots); + errdefer gpa.destroy(slots); + slots.* = try gdn.Slots.init(gpa, &e.driver, e.weights.model.dims, @max(o.want, 1)); + e.gdn = slots; + e.weights.model.gdn = slots; + } _ = try steps.done(.{ 0, 0 }, .first); return e; } @@ -197,6 +208,11 @@ pub const Engine = struct { }, ); e.graphs.deinit(); + if (e.gdn) |p| { + p.deinit(); + e.gpa.destroy(p); + e.weights.model.gdn = null; + } e.pool.deinit(e.gpa); if (e.votes) |*v| v.deinit(); e.drawer.deinit(); @@ -247,7 +263,8 @@ pub const Engine = struct { head = hl.bytes() + hl.cinBytes() + 4; } return l.bytes() + l.windowBytes() + l.cinBytes() + head + - cache.max_window * e.weights.model.dims.wide() * e.act.size(); + cache.max_window * e.weights.model.dims.wide() * e.act.size() + + if (gdn.supported(e.weights.model.dims, e.cfg.layers)) gdn.Plan.bytes(e.weights.model.dims) else 0; } /// The scratch for `o.chunk` rows, or the fallback's when the larger chunk admits fewer of the streams asked for. @@ -335,23 +352,29 @@ pub const Engine = struct { if (kept == 0 or kept > tokens.len) return error.BadKeep; const l = c.layout; const s = e.stream.handle; - if (kept < tokens.len) for (c.linear, 0..) |linear, i| { - if (!linear) continue; - const from = c.snapshot(i, kept - 1); - const to = c.linearAt(i); - try hip.runtime.check(e.driver.api.hipMemcpyDtoDAsync( - @ptrFromInt(to.conv), - @ptrFromInt(from.conv), - l.conv_bytes, - s, - )); - try hip.runtime.check(e.driver.api.hipMemcpyDtoDAsync( - @ptrFromInt(to.state), - @ptrFromInt(from.state), - l.state_bytes, - s, - )); - }; + const plan = if (e.weights.model.gdn) |p| p.planFor(c.id) else null; + // a plan leaves the state where its window began, so a keep always runs; a one-row window wrote it back + if (plan) |p| { + try p.keep(e.ops(&e.arena), kept); + } else if (kept < tokens.len) { + for (c.linear, 0..) |linear, i| { + if (!linear) continue; + const from = c.snapshot(i, kept - 1); + const to = c.linearAt(i); + try hip.runtime.check(e.driver.api.hipMemcpyDtoDAsync( + @ptrFromInt(to.conv), + @ptrFromInt(from.conv), + l.conv_bytes, + s, + )); + try hip.runtime.check(e.driver.api.hipMemcpyDtoDAsync( + @ptrFromInt(to.state), + @ptrFromInt(from.state), + l.state_bytes, + s, + )); + } + } const row = l.tail_bytes / l.tail_rows; try hip.runtime.check(e.driver.api.hipMemcpyDtoDAsync( @ptrFromInt(c.tail()), @@ -404,6 +427,11 @@ pub const Engine = struct { e.host.used = 0; e.arena.reset(); const staged = try forward.stage(&e.weights.model, &e.host, segs, true); + // a window's plan words are written here, outside the graph its round may replay, which copies them up + if (e.weights.model.gdn) |s| for (segs) |g| { + const p = s.planFor(g.c.id) orelse return error.CachesNotPlanned; + try p.fill(g.c, g.tokens.len); + }; const began = nowNs(e.io); defer e.graphs.submit_ns += @intCast(nowNs(e.io) - began); e.graphs.rounds += 1; diff --git a/zig/src/families/flashnext/hip/forward.zig b/zig/src/families/flashnext/hip/forward.zig index f05995286d5..f7a7b9a682c 100644 --- a/zig/src/families/flashnext/hip/forward.zig +++ b/zig/src/families/flashnext/hip/forward.zig @@ -342,6 +342,11 @@ pub fn run(ops: Ops, m: *const model.Model, host: *Host, segs: []const Seg, wind g.c.cin.base() else try o.arena.take((g.c.layout.tail_rows + g.tokens.len) * d.wide() * m.act.size()); + // a window's words go up inside the round: a graph holds the copy, so every replay sends this round's own + if (window) if (m.gdn) |s| for (segs) |g| { + const p = s.planFor(g.c.id) orelse return error.CachesNotPlanned; + try p.send(o.stream); + }; for (m.layers, 0..) |layer, i| { const mark = o.arena.mark(); defer o.arena.release(mark); diff --git a/zig/src/families/flashnext/hip/mixers.zig b/zig/src/families/flashnext/hip/mixers.zig index 830a588f684..d1c9207551a 100644 --- a/zig/src/families/flashnext/hip/mixers.zig +++ b/zig/src/families/flashnext/hip/mixers.zig @@ -4,6 +4,7 @@ const hip = @import("hip"); const model = @import("model.zig"); const cache = @import("cache.zig"); const forward = @import("forward.zig"); +const gdn = @import("gdn_replay.zig"); const Ops = hip.ops.Ops; const Tensor = hip.ops.Tensor; @@ -35,6 +36,8 @@ pub fn linear( for (&outs, [_]model.Projection{ l.qkv, l.z, l.a, l.b }) |*t, p| t.* = try o.project(x, p, rows, false); } const qkv, const z, const a, const b = outs; + // a window on the plan path: every row of it through the stream's plan, no state snapshot a row + if (window) if (m.gdn) |s| return gdn.linear(o, m, l, segs, index, qkv, z, a, b, rows, s); const q = try o.arena.of(f32, rows * kw); const k = try o.arena.of(f32, rows * kw); const v = try o.arena.of(f32, rows * vw); diff --git a/zig/src/families/flashnext/hip/model.zig b/zig/src/families/flashnext/hip/model.zig index 56ecce8cc95..156c3ce0986 100644 --- a/zig/src/families/flashnext/hip/model.zig +++ b/zig/src/families/flashnext/hip/model.zig @@ -146,6 +146,8 @@ pub const Model = struct { head: Projection, ple: Ple, mtp: ?Head = null, + /// The cache slots' DeltaNet plans (null: the per-row snapshot path, which keeps the state a row). + gdn: ?*@import("gdn_replay.zig").Slots = null, }; test "a rank's dimensions: heads, value heads and the vocabulary split, the kv heads held two ranks a head" { From 1b2f54178782579a1ac2ba84f0176c2522961a35 Mon Sep 17 00:00:00 2001 From: gilby Date: Tue, 10 Nov 2026 03:39:07 +0100 Subject: [PATCH 044/186] hip kernel tests: Flash Next's plan path against its per-row path, q, k, v, state and conv tail At Flash Next's shape: the plan's conv split against today's launch and the kept conv tail for one to four rows, the plan's delta rule y against serial and the state a keep replays for each kept row (a one-row window is its own keep; a longer one leaves the state as the round found it), and two streams each replayed from its own rows. Host: the words' layout, the kept sections' offsets, and the shape guard, which is the only choice now. --- zig/src/families/flashnext/hip/gdn_replay.zig | 57 +++ zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/gdn_plan.zig | 410 ++++++++++++++++++ 3 files changed, 468 insertions(+) create mode 100644 zig/src/hip/kernel_tests/gdn_plan.zig diff --git a/zig/src/families/flashnext/hip/gdn_replay.zig b/zig/src/families/flashnext/hip/gdn_replay.zig index 8e8c6cd21f7..c06336a3325 100644 --- a/zig/src/families/flashnext/hip/gdn_replay.zig +++ b/zig/src/families/flashnext/hip/gdn_replay.zig @@ -4,6 +4,7 @@ const hip = @import("hip"); const pops = hip.plan_ops; const cache = @import("cache.zig"); +const config = @import("../config.zig"); const forward = @import("forward.zig"); const model = @import("model.zig"); @@ -341,3 +342,59 @@ pub fn linear( try o.gnormOut(y, l.norm, z, out, rows * d.value_heads, d.value_dim, d.eps, d.gate); return o.project(out, l.out, rows, l.out.partial); } + +test "the words keep every slot's arrays and put the pointers on even words" { + const l = Layout.of(48); + try std.testing.expect(l.ptr % 2 == 0); + try std.testing.expect(l.snaps == l.ptr + 2 * slots); + try std.testing.expect(l.words == l.snaps + 4 * 48); + try std.testing.expect(Layout.start(48) % 2 == 0); + try std.testing.expectEqual(@as(usize, 2 * (2 + 2 * 48)), Layout.desc(48)); + try std.testing.expectEqual(Layout.start(48) + slots * Layout.desc(48), Layout.total(48)); +} + +test "the plan path is the default: Flash Next's shapes take it, the shapes its kernels cannot take stay per-row" { + var c = try config.parse(std.testing.allocator, @embedFile("../fixtures/config.json")); + defer c.deinit(); + const d = try model.Dims.of(&c, 1); + try std.testing.expectEqual(@as(usize, 16), d.key_heads); + try std.testing.expectEqual(@as(usize, 48), d.value_heads); + try std.testing.expectEqual(@as(usize, 4), d.conv); + // the shapes the plan kernels are: no variable asks for it + try std.testing.expect(supported(d, 48)); + // every other shape keeps the per-row snapshots + var narrow = d; + narrow.key_dim = 64; + try std.testing.expect(!supported(narrow, 48)); + var odd = d; + odd.value_dim = 100; + try std.testing.expect(!supported(odd, 48)); + odd = d; + odd.value_heads = 50; + try std.testing.expect(!supported(odd, 48)); + var long = d; + long.conv = 9; + try std.testing.expect(!supported(long, 48)); + try std.testing.expect(!supported(d, 0)); +} + +test "one layer's kept sections sit one after another over the window's rows" { + const kw = 16 * 128; + const vw = 48 * 128; + const vh = 48; + var p: Plan = undefined; + p.rows = 4; + p.room = window_rows; + p.key_heads = 16; + p.key_dim = 128; + p.value_heads = vh; + p.value_dim = 128; + p.kept_stride = 4 * p.room * (kw + vw + 2 * vh); + p.kept = .{ .r = undefined, .ptr = @ptrFromInt(0x1000), .len = 0 }; + const first = p.keptOf(0, 0); + const second = p.keptOf(0, 1); + try std.testing.expectEqual(first.k + 4 * kw, second.k); + try std.testing.expectEqual(first.k, first.v - 4 * p.rows * kw); + try std.testing.expectEqual(first.v, first.gate - 4 * p.rows * vw); + try std.testing.expectEqual(first.beta, first.gate + 4 * p.rows * vh); +} diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index bb9c604d654..5e8109dd613 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -20,6 +20,7 @@ test "the code objects load on this device and every kernel the launchers use re test { _ = @import("kernel_tests/gdn.zig"); _ = @import("kernel_tests/gdn_bits.zig"); + _ = @import("kernel_tests/gdn_plan.zig"); _ = @import("kernel_tests/mamba.zig"); _ = @import("kernel_tests/decode.zig"); _ = @import("kernel_tests/tail_pin.zig"); diff --git a/zig/src/hip/kernel_tests/gdn_plan.zig b/zig/src/hip/kernel_tests/gdn_plan.zig new file mode 100644 index 00000000000..f96040147aa --- /dev/null +++ b/zig/src/hip/kernel_tests/gdn_plan.zig @@ -0,0 +1,410 @@ +//! A window's plan path against its per-row path at Flash Next's shape: the same rows, y, state and conv tail. +const std = @import("std"); +const rig = @import("rig.zig"); +const plan = @import("../launches/plan.zig"); +const plan_ops = @import("../ops/plan.zig"); +const DeviceBuffer = @import("../memory.zig").DeviceBuffer; +const Ops = @import("../ops/ops.zig").Ops; + +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; + +const key_heads = 16; +const value_heads = 48; +const dk = 128; +const dv = 128; +const conv = 4; +const taps = conv - 1; +const channels = 2 * key_heads * dk + value_heads * dv; +const kw = key_heads * dk; +const vw = value_heads * dv; +const sections = kw + vw + 2 * value_heads; + +/// A plan's words: a slot's first row, its rows, its kept rows and its descriptor, then an empty slot. +const Words = struct { + first: DeviceBuffer, + count: DeviceBuffer, + keep: DeviceBuffer, + desc: DeviceBuffer, + ptrs: DeviceBuffer, + snaps: DeviceBuffer, + slots: usize, + + /// `firsts` a slot's first row, `counts` its rows, `descs` its descriptor pair; the last slot must be empty. + fn make(t: *Rig, firsts: []const u32, counts: []const u32, descs: []const [4]u64, snaps: [2]u64) !Words { + var w: Words = .{ + .first = try t.upload(firsts), + .count = undefined, + .keep = undefined, + .desc = undefined, + .ptrs = undefined, + .snaps = undefined, + .slots = firsts.len, + }; + errdefer w.first.free(); + w.count = try t.upload(counts); + errdefer w.count.free(); + w.keep = try t.alloc(firsts.len * 4); + errdefer w.keep.free(); + w.desc = try t.alloc(descs.len * 4 * 8); + errdefer w.desc.free(); + w.ptrs = try t.alloc(firsts.len * 8); + errdefer w.ptrs.free(); + w.snaps = try t.alloc(2 * 8); + errdefer w.snaps.free(); + for (descs, 0..) |d, i| try w.desc.upload(i * 32, std.mem.sliceAsBytes(&d)); + const held = try gpa.alloc(u64, firsts.len); + defer gpa.free(held); + for (held, 0..) |*p, i| p.* = at(w.desc) + i * 32; + try w.ptrs.upload(0, std.mem.sliceAsBytes(held)); + try w.snaps.upload(0, std.mem.sliceAsBytes(&snaps)); + const marks = try gpa.alloc(i32, firsts.len); + defer gpa.free(marks); + @memset(marks, -1); + try w.keep.upload(0, std.mem.sliceAsBytes(marks)); + return w; + } + + /// The empty slot after the window's: no rows, no state, no kept inputs. + fn none() [4]u64 { + return .{ 0, 0, 0, 0 }; + } + + fn ref(w: *const Words, rows: usize) plan.PlanRef { + return .{ + .rows = rows, + .slots = w.slots, + .args = .{ + .first = at(w.first), + .count = at(w.count), + .desc = at(w.ptrs), + .snaps = at(w.snaps), + }, + }; + } + + /// The kept rows `n` of slot 0 (`-1` leaves it alone). + fn mark(w: *const Words, o: Ops, n: i32) !void { + try w.keep.upload(0, std.mem.sliceAsBytes(&[_]i32{ n, -1 })); + _ = o; + } + + fn free(w: *Words) void { + for ([_]*DeviceBuffer{ &w.snaps, &w.ptrs, &w.desc, &w.keep, &w.count, &w.first }) |b| b.free(); + } +}; + +test "a window's split from its plan gives the per-row launch's q, k and v bits, at one to four rows" { + const t = try Rig.open(1 << 20); + defer t.close(); + const o = t.ops(&t.on); + inline for (1..5) |rows| { + var rng: rig.Rng = .{ .state = 0xC04 + rows }; + const hx = try gpa.alloc(u16, rows * channels); + defer gpa.free(hx); + for (hx) |*e| e.* = t.bits(rng.unit() * 2); + const hw = try gpa.alloc(f32, channels * conv); + defer gpa.free(hw); + for (hw) |*e| e.* = rng.unit() / 4; + const hs = try gpa.alloc(f32, taps * channels); + defer gpa.free(hs); + for (hs) |*e| e.* = rng.unit(); + var x = try t.upload(hx); + defer x.free(); + var w = try t.upload(hw); + defer w.free(); + var state = try t.upload(hs); + defer state.free(); + var states = try t.alloc(rows * taps * channels * 4); + defer states.free(); + var a: [3]DeviceBuffer = .{ + try t.alloc(rows * kw * 4), + try t.alloc(rows * kw * 4), + try t.alloc(rows * vw * 4), + }; + defer for (&a) |*b| b.free(); + var b: [3]DeviceBuffer = .{ + try t.alloc(rows * kw * 4), + try t.alloc(rows * kw * 4), + try t.alloc(rows * vw * 4), + }; + defer for (&b) |*x2| x2.free(); + var samples = try t.alloc(rows * channels * 4); + defer samples.free(); + var tail = try t.upload(hs); + defer tail.free(); + var words = try Words.make(t, &.{ 0, 0 }, &.{ @intCast(rows), 0 }, &.{ .{ 0, 0, at(tail), 0 }, Words.none() }, .{ at(samples), 0 }); + defer words.free(); + // today's window: one launch over the rows, each row's conv tail kept aside + try o.l.tf_conv_split(.{ + .x = at(x), + .kind = if (o.bf16) 2 else 1, + .weight = at(w), + .state = at(state), + .states = at(states), + .qn = at(a[0]), + .kn = at(a[1]), + .v = at(a[2]), + .channels = channels, + .kernel = conv, + .rows = @intCast(rows), + .kw = kw, + .vw = vw, + }, o.stream); + // the plan's: the same launch from the window's words, which leaves the state where it found it + try plan_ops.convSplit( + o, + .{ .ptr = at(x), .kind = if (o.bf16) .bf16 else .f16 }, + at(w), + at(b[0]), + at(b[1]), + at(b[2]), + words.ref(rows), + 0, + channels, + conv, + kw, + vw, + null, + null, + ); + try t.stream.synchronize(); + for ([_][]const u8{ "q", "k", "v" }, 0..) |what, i| { + const want = try rig.download(u32, a[i]); + defer gpa.free(want); + const got = try rig.download(u32, b[i]); + defer gpa.free(got); + if (!std.mem.eql(u32, want, got)) std.debug.print("{s} differs at {d} rows\n", .{ what, rows }); + try std.testing.expectEqualSlices(u32, want, got); + } + // the kept rows' tails come from the conv inputs: a keep takes row n's tail from the samples it kept + for (1..rows + 1) |n| { + try words.mark(o, @intCast(n)); + try plan_ops.keep(o, words.ref(rows), .{ + .keep = at(words.keep), + .hidden = 0, + .hidden_words = 0, + .channels = channels, + .taps = taps, + .layers = 1, + }); + try t.stream.synchronize(); + const want = try rig.download(u32, .{ .r = states.r, .ptr = @ptrFromInt(at(states) + (n - 1) * taps * channels * 4), .len = taps * channels * 4 }); + defer gpa.free(want); + const got = try rig.download(u32, tail); + defer gpa.free(got); + if (!std.mem.eql(u32, want, got)) std.debug.print("conv tail differs after {d} of {d} rows\n", .{ n, rows }); + try std.testing.expectEqualSlices(u32, want, got); + try tail.upload(0, std.mem.sliceAsBytes(hs)); + } + } +} + +test "a window's delta rule from its plan, the first n rows kept, leaves the state today's snapshot of row n holds" { + const t = try Rig.open(1 << 20); + defer t.close(); + const o = t.ops(&t.on); + inline for (1..5) |rows| { + var rng: rig.Rng = .{ .state = 0x6D3 + rows }; + const hq = try gpa.alloc(f32, rows * kw); + defer gpa.free(hq); + const hk = try gpa.alloc(f32, rows * kw); + defer gpa.free(hk); + const hv = try gpa.alloc(f32, rows * vw); + defer gpa.free(hv); + const hg = try gpa.alloc(f32, rows * value_heads); + defer gpa.free(hg); + const hb = try gpa.alloc(f32, rows * value_heads); + defer gpa.free(hb); + const hs = try gpa.alloc(f32, value_heads * dk * dv); + defer gpa.free(hs); + for ([_][]f32{ hq, hk, hv }) |set| { + for (set) |*e| e.* = rng.unit(); + } + for (hg) |*e| e.* = 0.5 + 0.25 * @abs(rng.unit()); + for (hb) |*e| e.* = 0.5 + 0.25 * @abs(rng.unit()); + for (hs) |*e| e.* = 0.1 * rng.unit(); + var q = try t.upload(hq); + defer q.free(); + var k = try t.upload(hk); + defer k.free(); + var v = try t.upload(hv); + defer v.free(); + var gate = try t.upload(hg); + defer gate.free(); + var beta = try t.upload(hb); + defer beta.free(); + var serial_state = try t.upload(hs); + defer serial_state.free(); + const snap = value_heads * dk * dv; + var snaps_a = try t.alloc(rows * snap * 4); + defer snaps_a.free(); + var y_a = try t.alloc(rows * vw * 4); + defer y_a.free(); + var planned_state = try t.upload(hs); + defer planned_state.free(); + // the window's start state, kept for comparison: a round leaves it there + var start = try t.upload(hs); + defer start.free(); + var y_b = try t.alloc(rows * vw * 4); + defer y_b.free(); + // the rows the conv split keeps: k, v, gate and beta one section after another + const hkp = try gpa.alloc(f32, rows * sections); + defer gpa.free(hkp); + @memcpy(hkp[0 .. rows * kw], hk); + @memcpy(hkp[rows * kw .. rows * (kw + vw)], hv); + @memcpy(hkp[rows * (kw + vw) .. rows * (kw + vw) + rows * value_heads], hg); + @memcpy(hkp[rows * (kw + vw) + rows * value_heads ..], hb); + var kept = try t.upload(hkp); + defer kept.free(); + var words = try Words.make(t, &.{ 0, 0 }, &.{ @intCast(rows), 0 }, &.{ .{ 0, 0, 0, at(planned_state) }, Words.none() }, .{ 0, at(kept) }); + defer words.free(); + try o.gatedDelta( + at(q), + at(k), + at(v), + at(gate), + at(beta), + at(serial_state), + at(y_a), + rows, + key_heads, + value_heads, + dk, + dv, + at(snaps_a), + ); + try plan_ops.gatedDelta(o, at(q), at(k), at(v), at(gate), at(beta), at(y_b), words.ref(rows), 0, key_heads, value_heads, dk, dv); + try t.stream.synchronize(); + const want_y = try rig.download(u32, y_a); + defer gpa.free(want_y); + const got_y = try rig.download(u32, y_b); + defer gpa.free(got_y); + try std.testing.expectEqualSlices(u32, want_y, got_y); + // a window of one row writes its state back in place; a longer one leaves it as the round found it + const round_state = try rig.download(u32, planned_state); + defer gpa.free(round_state); + const row1 = try rig.download(u32, .{ .r = snaps_a.r, .ptr = @ptrFromInt(at(snaps_a)), .len = snap * 4 }); + defer gpa.free(row1); + if (rows == 1) { + // the round is its own keep: the state it holds is the serial run's after its one row + try std.testing.expectEqualSlices(u32, row1, round_state); + // and a keep replays nothing (the kernel stops at a one-row window): the state stays that one + try words.mark(o, 1); + try plan_ops.gdnReplay(o, at(words.keep), words.ref(rows), key_heads, value_heads, dk, dv, 1); + try t.stream.synchronize(); + const got = try rig.download(u32, planned_state); + defer gpa.free(got); + try std.testing.expectEqualSlices(u32, row1, got); + } else { + const as_found = try rig.download(u32, start); + defer gpa.free(as_found); + try std.testing.expectEqualSlices(u32, as_found, round_state); + for (1..rows + 1) |n| { + // a round leaves the state where it found it: a keep replays from the window's start + try planned_state.upload(0, std.mem.sliceAsBytes(hs)); + try words.mark(o, @intCast(n)); + try plan_ops.gdnReplay(o, at(words.keep), words.ref(rows), key_heads, value_heads, dk, dv, 1); + try t.stream.synchronize(); + const want = try rig.download(u32, .{ .r = snaps_a.r, .ptr = @ptrFromInt(at(snaps_a) + (n - 1) * snap * 4), .len = snap * 4 }); + defer gpa.free(want); + const got = try rig.download(u32, planned_state); + defer gpa.free(got); + if (!std.mem.eql(u32, want, got)) std.debug.print("state differs after {d} of {d} rows\n", .{ n, rows }); + try std.testing.expectEqualSlices(u32, want, got); + } + } + } +} + +test "two streams' plans keep their own rows: each state as its serial run's snapshot of the kept row holds it" { + const t = try Rig.open(1 << 20); + defer t.close(); + const o = t.ops(&t.on); + const counts = [2]usize{ 2, 3 }; + const snap = value_heads * dk * dv; + var rng: rig.Rng = .{ .state = 0x77E }; + const hs = try gpa.alloc(f32, snap); + defer gpa.free(hs); + for (hs) |*e| e.* = 0.1 * rng.unit(); + // a stream's rows: q, k, v, gate and beta, and the ones the conv split keeps (k, v, gate, beta one section each) + var q: [2]DeviceBuffer = undefined; + var k: [2]DeviceBuffer = undefined; + var v: [2]DeviceBuffer = undefined; + var gate: [2]DeviceBuffer = undefined; + var beta: [2]DeviceBuffer = undefined; + var kept: [2]DeviceBuffer = undefined; + var serial: [2]DeviceBuffer = undefined; + var planned: [2]DeviceBuffer = undefined; + var ys: [2]DeviceBuffer = undefined; + var snaps: [2]DeviceBuffer = undefined; + var words: [2]Words = undefined; + var made: usize = 0; + errdefer for ([_]*DeviceBuffer{ + &q[0], &q[1], &k[0], &k[1], &v[0], &v[1], &gate[0], &gate[1], &beta[0], &beta[1], + &kept[0], &kept[1], &serial[0], &serial[1], &planned[0], &planned[1], &ys[0], &ys[1], &snaps[0], &snaps[1], + }) |b| b.free(); + defer for (words[0..made]) |*w| w.free(); + for (0..2) |i| { + const n = counts[i]; + const hq = try gpa.alloc(f32, n * kw); + defer gpa.free(hq); + const hk = try gpa.alloc(f32, n * kw); + defer gpa.free(hk); + const hv = try gpa.alloc(f32, n * vw); + defer gpa.free(hv); + const hg = try gpa.alloc(f32, n * value_heads); + defer gpa.free(hg); + const hb = try gpa.alloc(f32, n * value_heads); + defer gpa.free(hb); + for ([_][]f32{ hq, hk, hv }) |set| { + for (set) |*e| e.* = rng.unit(); + } + for (hg) |*e| e.* = 0.5 + 0.25 * @abs(rng.unit()); + for (hb) |*e| e.* = 0.5 + 0.25 * @abs(rng.unit()); + const hk_ = try gpa.alloc(f32, n * sections); + defer gpa.free(hk_); + @memcpy(hk_[0 .. n * kw], hk); + @memcpy(hk_[n * kw .. n * (kw + vw)], hv); + @memcpy(hk_[n * (kw + vw) .. n * (kw + vw) + n * value_heads], hg); + @memcpy(hk_[n * (kw + vw) + n * value_heads ..], hb); + q[i] = try t.upload(hq); + k[i] = try t.upload(hk); + v[i] = try t.upload(hv); + gate[i] = try t.upload(hg); + beta[i] = try t.upload(hb); + kept[i] = try t.upload(hk_); + serial[i] = try t.upload(hs); + planned[i] = try t.upload(hs); + ys[i] = try t.alloc(n * vw * 4); + snaps[i] = try t.alloc(n * snap * 4); + words[i] = try Words.make( + t, + &.{ 0, 0 }, + &.{ @intCast(n), 0 }, + &.{ .{ 0, 0, 0, at(planned[i]) }, Words.none() }, + .{ 0, at(kept[i]) }, + ); + made += 1; + // the stream's serial run, a snapshot a row, and its own plan's, in the same launch + try o.gatedDelta(at(q[i]), at(k[i]), at(v[i]), at(gate[i]), at(beta[i]), at(serial[i]), at(ys[i]), n, key_heads, value_heads, dk, dv, at(snaps[i])); + try plan_ops.gatedDelta(o, at(q[i]), at(k[i]), at(v[i]), at(gate[i]), at(beta[i]), at(ys[i]), words[i].ref(n), 0, key_heads, value_heads, dk, dv); + } + try t.stream.synchronize(); + for (0..2) |i| { + const n = counts[i]; + for (1..n + 1) |j| { + try words[i].keep.upload(0, std.mem.sliceAsBytes(&[_]i32{ @intCast(j), -1 })); + try planned[i].upload(0, std.mem.sliceAsBytes(hs)); + try plan_ops.gdnReplay(o, at(words[i].keep), words[i].ref(n), key_heads, value_heads, dk, dv, 1); + try t.stream.synchronize(); + const want = try rig.download(u32, .{ .r = snaps[i].r, .ptr = @ptrFromInt(at(snaps[i]) + (j - 1) * snap * 4), .len = snap * 4 }); + defer gpa.free(want); + const got = try rig.download(u32, planned[i]); + defer gpa.free(got); + if (!std.mem.eql(u32, want, got)) std.debug.print("stream {d} kept {d} of {d}: state differs\n", .{ i, j, n }); + try std.testing.expectEqualSlices(u32, want, got); + } + } +} From c35068d4618af68690c06bb9a05c03800976f662 Mon Sep 17 00:00:00 2001 From: gilby Date: Tue, 10 Nov 2026 03:39:07 +0100 Subject: [PATCH 045/186] changelog: Flash Next's plan verify windows under Unreleased --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e65663831a4..9a2eaef823d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -230,6 +230,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## 1.2.6 - Prompt cache: long conversations survive 8 to 16 agents sending follow-ups at once. Redundant states are evicted before each conversation's newest one, admission never evicts a state a queued request resumes, concurrent keeps reserve their room, and the default budget scales with the room left after streams on Metal and unified-memory CUDA (discrete CUDA hosts keep 8 GiB, Flash Next's batch host 16 GiB). On Qwen3.8-27B with DFlash2 at 16 agents, a follow-up's median first token drops from 546 s to 15.9 s on an M3 Ultra and from 105 s to 25.8 s on a GB10, with every reply unchanged (#2699). +- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Evidence index: reads saved native capabilities, including unknown and XPU chip labels; advertised facts stay unqualified, and invalid supplied input stops with a bounded error (#2708). From 76e8ea89803649a40dd57701b21fb418681d380e Mon Sep 17 00:00:00 2001 From: gilby Date: Wed, 11 Nov 2026 18:20:20 +0100 Subject: [PATCH 046/186] test(hip): run Flash Next's conv split with norms and gates on, both paths The kernel tests now reach the family's model-free side through kernel_reach, so a conv split test can build a real plan and compare today's launch with the plan's q, k, v, gate and beta bit for bit. --- zig/build/hip.zig | 6 +- zig/src/families/flashnext/kernel_reach.zig | 11 ++++ zig/src/hip/kernel_tests.zig | 33 +++++++++-- zig/src/hip/kernel_tests/gdn_plan.zig | 66 ++++++++++++++++++++- 4 files changed, 109 insertions(+), 7 deletions(-) create mode 100644 zig/src/families/flashnext/kernel_reach.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 4fca8f58c8e..968da0cbea9 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -125,7 +125,11 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil kernel_module.addIncludePath(.{ .cwd_relative = hip_include }); kernel_module.addCSourceFile(.{ .file = b.path("zig/src/hip/device_arch.c"), .flags = &.{"-D__HIP_PLATFORM_AMD__"} }); kernel_module.addImport("hip_kernels", objects); - kernel_module.addImport("core", coreModule(b, target)); + const core = coreModule(b, target); + kernel_module.addImport("core", core); + kernel_module.addImport("cuda_memory", b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_memory.zig"), .target = target })); + const reach_module = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/kernel_reach.zig"), .target = target, .link_libc = true, .imports = &.{ .{ .name = "hip", .module = kernel_module }, .{ .name = "core", .module = core } } }); + kernel_module.addImport("flashnext_kernel", reach_module); const kernel_test = b.addTest(.{ .root_module = kernel_module }); b.step("hip-kernel-build", "Compile the model-free kernels and their GPU tests without running them").dependOn(&kernel_test.step); b.step("hip-kernel-test", "Every model-free kernel against its host reference on the GPU").dependOn(&b.addRunArtifact(kernel_test).step); diff --git a/zig/src/families/flashnext/kernel_reach.zig b/zig/src/families/flashnext/kernel_reach.zig new file mode 100644 index 00000000000..503d8bd8299 --- /dev/null +++ b/zig/src/families/flashnext/kernel_reach.zig @@ -0,0 +1,11 @@ +//! Flash Next for the kernel tests: the dimensions, the caches a stream holds, the window plan and the mixers. +const cfg_mod = @import("config.zig"); + +pub const config = cfg_mod; +pub const cache = @import("hip/cache.zig"); +pub const gdn_replay = @import("hip/gdn_replay.zig"); +pub const forward = @import("hip/forward.zig"); +pub const mixers = @import("hip/mixers.zig"); +pub const model = @import("hip/model.zig"); +pub const tables = @import("hip/tables.zig"); +pub const ngram = @import("hip/ngram.zig"); diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 5e8109dd613..237b73a4bab 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -1,9 +1,34 @@ //! The model-free kernels on a real GPU: every kernel resolves, and each launch matches its host reference. const std = @import("std"); -const Runtime = @import("runtime.zig").Runtime; -const Context = @import("context.zig").Context; -const kernels = @import("kernels.zig"); -const Launcher = @import("launches.zig").Launcher; + +/// Flash Next's model-free side: its config, caches, window plan and mixers, on this module as `hip`. +pub fn flashnext() type { + return flashnext_mod; +} +const flashnext_mod = @import("flashnext_kernel"); +pub const cache = flashnext_mod.cache; +pub const gdn_replay = flashnext_mod.gdn_replay; +pub const forward = flashnext_mod.forward; +pub const mixers = flashnext_mod.mixers; +pub const model = flashnext_mod.model; + +const hipmod = @import("root.zig"); +pub const Launcher = hipmod.Launcher; +pub const Runtime = hipmod.Runtime; +pub const Context = hipmod.Context; +pub const kernels = hipmod.kernels; +pub const raw = hipmod.raw; +pub const usage = hipmod.usage; +pub const ops = hipmod.ops; +pub const plan_ops = hipmod.plan_ops; +pub const counting = hipmod.counting; +pub const moe = hipmod.moe; +pub const collective = hipmod.collective; +pub const abi = hipmod.abi; +pub const runtime = hipmod.runtime; +pub const Stream = hipmod.Stream; +pub const DeviceBuffer = hipmod.DeviceBuffer; +pub const HostBuffer = hipmod.HostBuffer; test "the code objects load on this device and every kernel the launchers use resolves" { var r = try Runtime.open(); diff --git a/zig/src/hip/kernel_tests/gdn_plan.zig b/zig/src/hip/kernel_tests/gdn_plan.zig index f96040147aa..3379dba3032 100644 --- a/zig/src/hip/kernel_tests/gdn_plan.zig +++ b/zig/src/hip/kernel_tests/gdn_plan.zig @@ -20,6 +20,8 @@ const channels = 2 * key_heads * dk + value_heads * dv; const kw = key_heads * dk; const vw = value_heads * dv; const sections = kw + vw + 2 * value_heads; +const vh = value_heads; +const eps = 1e-6; /// A plan's words: a slot's first row, its rows, its kept rows and its descriptor, then an empty slot. const Words = struct { @@ -99,6 +101,7 @@ test "a window's split from its plan gives the per-row launch's q, k and v bits, const t = try Rig.open(1 << 20); defer t.close(); const o = t.ops(&t.on); + const act_kind: @import("../ops/types.zig").Kind = if (o.bf16) .bf16 else .f16; inline for (1..5) |rows| { var rng: rig.Rng = .{ .state = 0xC04 + rows }; const hx = try gpa.alloc(u16, rows * channels); @@ -136,6 +139,45 @@ test "a window's split from its plan gives the per-row launch's q, k and v bits, defer tail.free(); var words = try Words.make(t, &.{ 0, 0 }, &.{ @intCast(rows), 0 }, &.{ .{ 0, 0, at(tail), 0 }, Words.none() }, .{ at(samples), 0 }); defer words.free(); + // the head norm's weights, and the gate's a, b, a_log and dt_bias, which it writes gate and beta from + const hq_norm = try gpa.alloc(f32, kw); + defer gpa.free(hq_norm); + for (hq_norm) |*e| e.* = rng.unit() + 1; + const hk_norm = try gpa.alloc(f32, kw); + defer gpa.free(hk_norm); + for (hk_norm) |*e| e.* = rng.unit() + 1; + const ga = try gpa.alloc(u16, rows * vh); + defer gpa.free(ga); + for (ga) |*e| e.* = t.bits(rng.unit() * 2 - 1); + const gb = try gpa.alloc(u16, rows * vh); + defer gpa.free(gb); + for (gb) |*e| e.* = t.bits(rng.unit() * 2 - 1); + const a_log = try gpa.alloc(f32, vh); + defer gpa.free(a_log); + for (a_log) |*e| e.* = rng.unit() - 1; + const dt_bias = try gpa.alloc(f32, vh); + defer gpa.free(dt_bias); + for (dt_bias) |*e| e.* = rng.unit(); + var qw = try t.upload(hq_norm); + defer qw.free(); + var kw_w = try t.upload(hk_norm); + defer kw_w.free(); + var ag = try t.upload(ga); + defer ag.free(); + var bg = try t.upload(gb); + defer bg.free(); + var alog = try t.upload(a_log); + defer alog.free(); + var dt = try t.upload(dt_bias); + defer dt.free(); + var gate_a = try t.alloc(rows * vh * 4); + defer gate_a.free(); + var beta_a = try t.alloc(rows * vh * 4); + defer beta_a.free(); + var gate_b = try t.alloc(rows * vh * 4); + defer gate_b.free(); + var beta_b = try t.alloc(rows * vh * 4); + defer beta_b.free(); // today's window: one launch over the rows, each row's conv tail kept aside try o.l.tf_conv_split(.{ .x = at(x), @@ -151,6 +193,18 @@ test "a window's split from its plan gives the per-row launch's q, k and v bits, .rows = @intCast(rows), .kw = kw, .vw = vw, + .qw = at(qw), + .kw_w = at(kw_w), + .eps = eps, + .norm = 128, + .ga = at(ag), + .gb = at(bg), + .a_log = at(alog), + .dt_bias = at(dt), + .gate = at(gate_a), + .beta = at(beta_a), + .gcount = @intCast(rows * vh), + .heads = vh, }, o.stream); // the plan's: the same launch from the window's words, which leaves the state where it found it try plan_ops.convSplit( @@ -166,8 +220,8 @@ test "a window's split from its plan gives the per-row launch's q, k and v bits, conv, kw, vw, - null, - null, + .{ .q = at(qw), .k = at(kw_w), .eps = eps }, + .{ .a = .{ .ptr = at(ag), .kind = act_kind }, .b = .{ .ptr = at(bg), .kind = act_kind }, .a_log = at(alog), .dt_bias = at(dt), .gate = at(gate_b), .beta = at(beta_b), .count = rows * vh, .heads = vh }, ); try t.stream.synchronize(); for ([_][]const u8{ "q", "k", "v" }, 0..) |what, i| { @@ -178,6 +232,14 @@ test "a window's split from its plan gives the per-row launch's q, k and v bits, if (!std.mem.eql(u32, want, got)) std.debug.print("{s} differs at {d} rows\n", .{ what, rows }); try std.testing.expectEqualSlices(u32, want, got); } + for ([_][]const u8{ "gate", "beta" }, [_]DeviceBuffer{ gate_a, beta_a }, [_]DeviceBuffer{ gate_b, beta_b }) |what, want_b, got_b| { + const want = try rig.download(u32, want_b); + defer gpa.free(want); + const got = try rig.download(u32, got_b); + defer gpa.free(got); + if (!std.mem.eql(u32, want, got)) std.debug.print("{s} differs at {d} rows\n", .{ what, rows }); + try std.testing.expectEqualSlices(u32, want, got); + } // the kept rows' tails come from the conv inputs: a keep takes row n's tail from the samples it kept for (1..rows + 1) |n| { try words.mark(o, @intCast(n)); From 53144740d88f2553e80553cb00d6c1df12b5dc1d Mon Sep 17 00:00:00 2001 From: gilby Date: Wed, 11 Nov 2026 18:20:29 +0100 Subject: [PATCH 047/186] engine(flashnext): free a plan's slots when opening fails, count its bytes once The slots follow the rest of the open path's errdefers, and streamBytes no longer adds the plan: the slots are made for the streams asked for before room is measured (#2728). Thanks to @gilby. --- CHANGELOG.md | 1 + zig/src/families/flashnext/hip/engine.zig | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a2eaef823d..d66ed367a92 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -75,6 +75,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. +- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. diff --git a/zig/src/families/flashnext/hip/engine.zig b/zig/src/families/flashnext/hip/engine.zig index 9f5029b3dfe..027c5c1d558 100644 --- a/zig/src/families/flashnext/hip/engine.zig +++ b/zig/src/families/flashnext/hip/engine.zig @@ -189,6 +189,7 @@ pub const Engine = struct { const slots = try gpa.create(gdn.Slots); errdefer gpa.destroy(slots); slots.* = try gdn.Slots.init(gpa, &e.driver, e.weights.model.dims, @max(o.want, 1)); + errdefer slots.deinit(); e.gdn = slots; e.weights.model.gdn = slots; } @@ -254,6 +255,7 @@ pub const Engine = struct { /// Device bytes a stream takes at `cap` positions: its caches, window snapshots, the head's caches and its rows. pub fn streamBytes(e: *const Engine, cap: usize) usize { + // a stream's plan is made for its slot when the engine opens, before any room is measured var buf: [config.max_layers]bool = undefined; const ks = e.kinds(&buf) catch return std.math.maxInt(usize); const l = cache.Layout.of(e.weights.model.dims, ks, cap, e.act); @@ -263,8 +265,7 @@ pub const Engine = struct { head = hl.bytes() + hl.cinBytes() + 4; } return l.bytes() + l.windowBytes() + l.cinBytes() + head + - cache.max_window * e.weights.model.dims.wide() * e.act.size() + - if (gdn.supported(e.weights.model.dims, e.cfg.layers)) gdn.Plan.bytes(e.weights.model.dims) else 0; + cache.max_window * e.weights.model.dims.wide() * e.act.size(); } /// The scratch for `o.chunk` rows, or the fallback's when the larger chunk admits fewer of the streams asked for. From 3ae4a10efae684a1b1b71d78a94914d3cb2fb05d Mon Sep 17 00:00:00 2001 From: gilby Date: Wed, 11 Nov 2026 20:30:27 +0100 Subject: [PATCH 048/186] test(hip): a Flash Next window through its plan, fill to keep --- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/gdn_window.zig | 530 ++++++++++++++++++++++++ 2 files changed, 531 insertions(+) create mode 100644 zig/src/hip/kernel_tests/gdn_window.zig diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 237b73a4bab..f5e44b272e0 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -46,6 +46,7 @@ test { _ = @import("kernel_tests/gdn.zig"); _ = @import("kernel_tests/gdn_bits.zig"); _ = @import("kernel_tests/gdn_plan.zig"); + _ = @import("kernel_tests/gdn_window.zig"); _ = @import("kernel_tests/mamba.zig"); _ = @import("kernel_tests/decode.zig"); _ = @import("kernel_tests/tail_pin.zig"); diff --git a/zig/src/hip/kernel_tests/gdn_window.zig b/zig/src/hip/kernel_tests/gdn_window.zig new file mode 100644 index 00000000000..03eec626849 --- /dev/null +++ b/zig/src/hip/kernel_tests/gdn_window.zig @@ -0,0 +1,530 @@ +//! A Flash Next window through gdn_replay: fill, send, the plan's conv split and delta rule, then keep(n). +const std = @import("std"); +const rig = @import("rig.zig"); +const hip = @import("../root.zig"); +const flashnext = @import("../kernel_tests.zig").flashnext(); +const cache = flashnext.cache; +const gdn = flashnext.gdn_replay; +const model = flashnext.model; +const mixers = flashnext.mixers; +const forward = flashnext.forward; + +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; + +const key_heads = 1; +const value_heads = 8; +const key_dim = 128; +const value_dim = 128; +const conv = 4; +const taps = conv - 1; +const kw = key_heads * key_dim; +const vw = value_heads * value_dim; +const channels = 2 * kw + vw; +const hidden = 32; +const vocab = 8; +const cap = 8; +/// Rows the linear attention's products take a window of: the pending token and three drafts. +const window_rows = 4; +const eps: f32 = 1e-6; +const state_words = value_heads * key_dim * value_dim; + +/// Flash Next's dimensions at this shape, one linear layer then a sparse one so both descriptor forms run. +fn dimsOf() model.Dims { + return .{ + .hidden = hidden, + .streams = 1, + .low = 1, + .vocab = vocab, + .layers = 2, + .eps = eps, + .heads = 1, + .kv_heads = 1, + .head_dim = 2, + .rotary = 1, + .theta = 10_000, + .key_heads = key_heads, + .value_heads = value_heads, + .key_dim = key_dim, + .value_dim = value_dim, + .conv = conv, + .index_heads = 1, + .index_dim = 2, + .index_ratio = 4, + .index_top = 1, + .ple_dim = 1, + .ple_taps = 2, + .ngram = 1, + .vocab_rows = vocab, + .gate = .sigmoid, + }; +} + +fn dense(weight: u64, n: usize, k: usize) model.Projection { + return .{ .n = @intCast(n), .k = @intCast(k), .handle = .{ .dense = .{ .weight = weight, .dtype = .f32 } } }; +} + +/// `n` rows of `k` weights, scaled so a row's product stays order one. +fn matrix(t: *Rig, rng: *rig.Rng, n: usize, k: usize) !rig.Buffer { + const host = try gpa.alloc(f32, n * k); + defer gpa.free(host); + const scale: f32 = @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(k)))); + for (host) |*e| e.* = rng.unit() * scale; + return rig.Buffer.fromHost(t, host); +} + +fn vector(t: *Rig, rng: *rig.Rng, n: usize, scale: f32) !rig.Buffer { + const host = try gpa.alloc(f32, n); + defer gpa.free(host); + for (host) |*e| e.* = rng.unit() * scale; + return rig.Buffer.fromHost(t, host); +} + +/// The linear layer's weight fields, in the order `free` takes them. +const names = struct { + const all = [_][]const u8{ "qkv", "z", "a", "b", "out", "conv", "a_log", "dt_bias", "norm", "qw", "kw_w" }; +}; + +const Weights = struct { + qkv: rig.Buffer, + z: rig.Buffer, + a: rig.Buffer, + b: rig.Buffer, + out: rig.Buffer, + conv: rig.Buffer, + a_log: rig.Buffer, + dt_bias: rig.Buffer, + norm: rig.Buffer, + qw: rig.Buffer, + kw_w: rig.Buffer, + + fn make(t: *Rig, rng: *rig.Rng) !Weights { + const shapes = [_][2]usize{ + .{ channels, hidden }, .{ vw, hidden }, + .{ value_heads, hidden }, .{ value_heads, hidden }, + .{ hidden, vw }, + }; + var w: Weights = undefined; + var made: usize = 0; + // every buffer made so far, in the order `free` takes them + errdefer inline for (names.all, 0..) |name, i| { + if (i < made) @field(w, name).free(); + }; + inline for (names.all[0..5], 0..) |name, i| { + @field(w, name) = try matrix(t, rng, shapes[i][0], shapes[i][1]); + made += 1; + } + w.conv = try vector(t, rng, channels * conv, 0.25); + w.a_log = try vector(t, rng, value_heads, 0.5); + w.dt_bias = try vector(t, rng, value_heads, 1.0); + w.norm = try vector(t, rng, value_dim, 0.5); + w.qw = try vector(t, rng, key_dim, 0.5); + w.kw_w = try vector(t, rng, key_dim, 0.5); + return w; + } + + fn free(w: *Weights) void { + inline for (names.all) |name| @field(w, name).free(); + } + + fn linear(w: *const Weights) model.Linear { + return .{ + .qkv = dense(w.qkv.ptr, channels, hidden), + .z = dense(w.z.ptr, vw, hidden), + .a = dense(w.a.ptr, value_heads, hidden), + .b = dense(w.b.ptr, value_heads, hidden), + .out = dense(w.out.ptr, hidden, vw), + .conv = w.conv.ptr, + .a_log = w.a_log.ptr, + .dt_bias = w.dt_bias.ptr, + .norm = w.norm.ptr, + }; + } + + /// The model this layer runs under: the linear mixer, the plan's slots (null: the per-row window), no sparse layer. + fn modelOf(w: *const Weights, act: model.Kind, slots: ?*gdn.Slots) model.Model { + return .{ + .dims = dimsOf(), + .qk = .{ .q = w.qw.ptr, .k = w.kw_w.ptr, .eps = eps }, + .act = act, + .embed = dense(w.qkv.ptr, hidden, hidden), + .layers = &.{}, + .mix = .{ .scale = 0, .down = dense(w.z.ptr, vw, hidden), .up = dense(w.out.ptr, hidden, vw) }, + .head = dense(w.out.ptr, vocab, hidden), + .ple = .{ + .layer = 0, + .kv = dense(w.z.ptr, vw, hidden), + .key_scale = 0, + .query_scale = 0, + .conv_scale = 0, + .conv = 0, + }, + .gdn = slots, + }; + } +}; + +/// A stream's caches over the two layers, its plan slots, and the window's start state and conv tail. +const Window = struct { + c: cache.Caches, + slots: gdn.Slots, + hs: []f32, + hc: []f32, + + fn open(t: *Rig, rng: *rig.Rng, act: model.Kind) !Window { + const kinds = [_]bool{ true, false }; + var win: Window = undefined; + win.c = try cache.Caches.init(gpa, &t.r, dimsOf(), &kinds, cap, act, 0); + errdefer win.c.deinit(gpa); + win.slots = try gdn.Slots.init(gpa, &t.r, dimsOf(), 1); + errdefer win.slots.deinit(); + win.hs = try gpa.alloc(f32, state_words); + errdefer gpa.free(win.hs); + win.hc = try gpa.alloc(f32, taps * channels); + for (win.hs) |*e| e.* = 0.1 * rng.unit(); + for (win.hc) |*e| e.* = rng.unit(); + return win; + } + + fn deinit(w: *Window) void { + gpa.free(w.hc); + gpa.free(w.hs); + w.slots.deinit(); + w.c.deinit(gpa); + } + + /// The window's start state and conv tail under the linear layer, the state's snapshots a nonzero byte. + fn reset(w: *Window, mark: u8) !void { + const l0 = w.c.linearAt(0); + try w.c.mem.upload(l0.state - w.c.mem.base(), std.mem.sliceAsBytes(w.hs)); + try w.c.mem.upload(l0.conv - w.c.mem.base(), std.mem.sliceAsBytes(w.hc)); + try w.c.window.fill8(mark); + } +}; + +/// A view of `b`'s bytes at `addr`, for reading what a launch wrote there. +fn bytesAt(b: hip.DeviceBuffer, addr: u64, n: usize) hip.DeviceBuffer { + return .{ .r = b.r, .ptr = @ptrFromInt(addr), .len = n }; +} + +/// `words` 32-bit words of `b` at `addr`, for reading what a launch wrote there. +fn readWords(b: hip.DeviceBuffer, addr: u64, words: usize) ![]u32 { + const out = try gpa.alloc(u32, words); + errdefer gpa.free(out); + try bytesAt(b, addr, words * 4).download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// The per-row window's snapshots of the rows it ran: each row's state and conv tail, and the window's start. +const Snaps = struct { + states: [][]u32, + convs: [][]u32, + start: []u32, + start_conv: []u32, + + fn read(w: *const Window, rows: usize) !Snaps { + var s: Snaps = undefined; + s.states = try gpa.alloc([]u32, rows); + errdefer gpa.free(s.states); + s.convs = try gpa.alloc([]u32, rows); + var made: usize = 0; + errdefer for (s.convs[0..made]) |b| gpa.free(b); + errdefer for (s.states[0..made]) |a| gpa.free(a); + for (0..rows) |r| { + s.states[r] = try readWords(w.c.window, w.c.snapshot(0, r).state, state_words); + made += 1; + s.convs[r] = try readWords(w.c.window, w.c.snapshot(0, r).conv, taps * channels); + } + s.start = try gpa.alloc(u32, state_words); + errdefer gpa.free(s.start); + s.start_conv = try gpa.alloc(u32, taps * channels); + for (s.start, w.hs) |*e, v| e.* = @bitCast(v); + for (s.start_conv, w.hc) |*e, v| e.* = @bitCast(v); + return s; + } + + fn deinit(s: *Snaps) void { + gpa.free(s.start_conv); + gpa.free(s.start); + for (s.states, s.convs) |a, b| { + gpa.free(a); + gpa.free(b); + } + gpa.free(s.convs); + gpa.free(s.states); + } +}; + +test "a window's plan: fill writes every slot's words and the linear layer's descriptor where the kernels read them" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x2728 }; + var win = try Window.open(t, &rng, if (t.bf16) .bf16 else .f16); + defer win.deinit(); + for ([_]usize{ 1, 2, 4 }) |rows| try fillCase(t, &win, rows); +} + +/// The plan's words after a fill for `rows`: where every kernel reads them. +fn fillCase(t: *Rig, win: *Window, rows: usize) !void { + const p = win.slots.planFor(0).?; + const l = p.layout(); + const d = dimsOf(); + try p.fill(&win.c, rows); + try p.send(t.stream.handle); + try t.stream.synchronize(); + const bytes = try gpa.alloc(u8, p.dev.len); + defer gpa.free(bytes); + try p.dev.download(0, bytes); + const held = std.mem.bytesAsSlice(u64, @as([]align(8) u8, @alignCast(bytes))); + const words = std.mem.bytesAsSlice(u32, @as([]align(4) u8, @alignCast(bytes))); + // the window slot's rows from its first, kept as asked; the empty slot runs and keeps nothing + try std.testing.expectEqual(@as(u32, 0), words[l.first]); + try std.testing.expectEqual(@as(u32, @intCast(rows)), words[l.count]); + try std.testing.expectEqual(@as(u32, @bitCast(@as(i32, -1))), words[l.keep]); + try std.testing.expectEqual(@as(u32, 0), words[l.first + 1]); + try std.testing.expectEqual(@as(u32, 0), words[l.count + 1]); + try std.testing.expectEqual(@as(u32, @bitCast(@as(i32, -1))), words[l.keep + 1]); + // the slot pointers name the plan's own descriptors, one after another + const d0 = p.dev.base() + 4 * std.mem.alignForward(usize, l.words, 2); + try std.testing.expectEqual(d0, held[l.ptr / 2]); + try std.testing.expectEqual(d0 + 4 * gdn.Layout.desc(d.layers), held[l.ptr / 2 + 1]); + // the linear layer's descriptor: no kept final row, then its conv and state; the sparse layer's zeros beside + const desc = std.mem.bytesAsSlice(u64, @as([]align(8) u8, @alignCast(bytes))); + const first = std.mem.alignForward(usize, l.words, 2) / 2; + const l0 = win.c.linearAt(0); + try std.testing.expectEqual(@as(u64, 0), desc[first]); + try std.testing.expectEqual(@as(u64, 0), desc[first + 1]); + try std.testing.expectEqual(l0.conv, desc[first + 2]); + try std.testing.expectEqual(l0.state, desc[first + 3]); + try std.testing.expectEqual(@as(u64, 0), desc[first + 4]); + try std.testing.expectEqual(@as(u64, 0), desc[first + 5]); + // the samples and kept rows the linear layer keeps, and the sparse layer's zeros + try std.testing.expectEqual(p.samplesOf(0), desc[l.snaps / 2]); + try std.testing.expectEqual(p.kept.base(), desc[l.snaps / 2 + 1]); + try std.testing.expectEqual(@as(u64, 0), desc[l.snaps / 2 + 2]); + try std.testing.expectEqual(@as(u64, 0), desc[l.snaps / 2 + 3]); +} + +test "a window's conv split on its plan, norms and gates on, writes the per-row launch's q, k, v, gate and beta" { + const t = try Rig.open(1 << 22); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x2729 }; + var w = try Weights.make(t, &rng); + defer w.free(); + var win = try Window.open(t, &rng, if (t.bf16) .bf16 else .f16); + defer win.deinit(); + for ([_]usize{ 1, 2, 3, 4 }) |rows| try convSplitCase(t, &w, &win, &rng, rows); +} + +/// Both paths' conv split over `rows`: today's launch keeps the state a row, the plan's the input a row. +fn convSplitCase(t: *Rig, w: *const Weights, win: *Window, rng: *rig.Rng, rows: usize) !void { + const o = t.ops(&t.on); + const act = w.modelOf(if (t.bf16) .bf16 else .f16, null).act; + const p = win.slots.planFor(0).?; + const l0 = win.c.linearAt(0); + const hx = try gpa.alloc(u16, rows * channels); + defer gpa.free(hx); + for (hx) |*e| e.* = t.bits(rng.unit()); + const ha = try gpa.alloc(u16, rows * value_heads); + defer gpa.free(ha); + for (ha) |*e| e.* = t.bits(rng.unit()); + const hb = try gpa.alloc(u16, rows * value_heads); + defer gpa.free(hb); + for (hb) |*e| e.* = t.bits(rng.unit()); + const ng = rows * value_heads; + var x = try rig.Buffer.fromHost(t, hx); + defer x.free(); + var ga = try rig.Buffer.fromHost(t, ha); + defer ga.free(); + var gb = try rig.Buffer.fromHost(t, hb); + defer gb.free(); + // the gate's values both launches write over the same a and b + var gate = try rig.Buffer.alloc(t, ng * 4); + defer gate.free(); + var beta = try rig.Buffer.alloc(t, ng * 4); + defer beta.free(); + try o.gdnGatePrefill( + .{ .ptr = ga.ptr, .kind = act }, + .{ .ptr = gb.ptr, .kind = act }, + w.a_log.ptr, + w.dt_bias.ptr, + gate.ptr, + beta.ptr, + ng, + value_heads, + ); + // the plan writes q beside its rows and k, v by plan_ops.convSplit's signature: only q has a buffer of its own + var planned: rig.Buffer = try rig.Buffer.alloc(t, rows * kw * 4); + defer planned.free(); + var today: [3]rig.Buffer = .{ undefined, undefined, undefined }; + var made: usize = 0; + defer for (today[0..made]) |*a| a.free(); + for ([_]usize{ rows * kw, rows * kw, rows * vw }) |n| { + today[made] = try rig.Buffer.alloc(t, n * 4); + made += 1; + } + try win.reset(0x5A); + try o.l.tf_conv_split(.{ + .x = x.ptr, + .kind = @backingInt(act), + .weight = w.conv.ptr, + .state = l0.conv, + .states = win.c.snapshot(win.c.slot[0], 0).conv, + .qn = today[0].ptr, + .kn = today[1].ptr, + .v = today[2].ptr, + .channels = channels, + .kernel = conv, + .rows = @intCast(rows), + .kw = kw, + .vw = vw, + .qw = w.qw.ptr, + .kw_w = w.kw_w.ptr, + .eps = eps, + .norm = 128, + .ga = ga.ptr, + .gb = gb.ptr, + .a_log = w.a_log.ptr, + .dt_bias = w.dt_bias.ptr, + .gate = gate.ptr, + .beta = beta.ptr, + .gcount = @intCast(ng), + .heads = value_heads, + }, o.stream); + try t.stream.synchronize(); + // the plan's: the same launch from the window's words, which keeps the rows' conv input in place + try win.reset(0x5A); + try p.fill(&win.c, rows); + try p.send(o.stream); + const kept = p.keptOf(0, 0); + try hip.plan_ops.convSplit( + o, + .{ .ptr = x.ptr, .kind = act }, + w.conv.ptr, + planned.ptr, + kept.k, + kept.v, + p.ref(), + 0, + channels, + conv, + kw, + vw, + .{ .q = w.qw.ptr, .k = w.kw_w.ptr, .eps = eps }, + .{ + .a = .{ .ptr = ga.ptr, .kind = act }, + .b = .{ .ptr = gb.ptr, .kind = act }, + .a_log = w.a_log.ptr, + .dt_bias = w.dt_bias.ptr, + .gate = kept.gate, + .beta = kept.beta, + .count = ng, + .heads = value_heads, + }, + ); + try t.stream.synchronize(); + for ([_]struct { what: []const u8, n: usize, a: hip.DeviceBuffer, b: hip.DeviceBuffer }{ + .{ .what = "q", .n = rows * kw, .a = today[0].buf, .b = planned.buf }, + .{ .what = "k", .n = rows * kw, .a = today[1].buf, .b = bytesAt(p.kept, kept.k, rows * kw * 4) }, + .{ .what = "v", .n = rows * vw, .a = today[2].buf, .b = bytesAt(p.kept, kept.v, rows * vw * 4) }, + .{ .what = "gate", .n = ng, .a = gate.buf, .b = bytesAt(p.kept, kept.gate, ng * 4) }, + .{ .what = "beta", .n = ng, .a = beta.buf, .b = bytesAt(p.kept, kept.beta, ng * 4) }, + }) |c2| { + const n = c2.n; + const want = try readWords(c2.a, at(c2.a), n); + defer gpa.free(want); + const got = try readWords(c2.b, at(c2.b), n); + defer gpa.free(got); + try std.testing.expectEqualSlices(u32, want, got); + } +} + +test "a window through its plan equals the per-row window at every keep n: y, the state and the conv tail" { + const t = try Rig.open(1 << 24); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x2730 }; + var w = try Weights.make(t, &rng); + defer w.free(); + var win = try Window.open(t, &rng, if (t.bf16) .bf16 else .f16); + defer win.deinit(); + var tokens: [window_rows]u32 = undefined; + for (&tokens, 0..) |*tok, i| tok.* = @intCast(i); + for ([_]usize{ 1, 2, 3, 4 }) |rows| try planCase(t, &w, &win, &rng, tokens[0..rows], rows); +} + +/// A window of `rows` both ways: today's forward, then the plan's, then every keep n against today's snapshots. +fn planCase(t: *Rig, w: *const Weights, win: *Window, rng: *rig.Rng, tokens: []const u32, rows: usize) !void { + const o = t.ops(&t.on); + const act: model.Kind = if (t.bf16) .bf16 else .f16; + const p = win.slots.planFor(0).?; + const l0 = win.c.linearAt(0); + const m = w.modelOf(act, null); + const planned = w.modelOf(act, &win.slots); + const hx = try gpa.alloc(u16, rows * channels); + defer gpa.free(hx); + for (hx) |*e| e.* = t.bits(rng.unit()); + var x = try rig.Buffer.fromHost(t, hx); + defer x.free(); + const xt: hip.ops.Tensor = .{ .ptr = x.ptr, .kind = act }; + const seg = [_]forward.Seg{.{ .c = &win.c, .tokens = tokens, .pos = 0 }}; + // today's window: one conv split and delta rule a stream, a snapshot a row, the state written back + try win.reset(0x5A); + o.arena.reset(); + const y0 = try mixers.linear(o, &m, w.linear(), &seg, 0, xt, rows, true); + try t.stream.synchronize(); + const y_a = try readWords(win.c.mem, y0.ptr, rows * hidden * act.size() / 4); + defer gpa.free(y_a); + var snaps = try Snaps.read(win, rows); + defer snaps.deinit(); + // the plan's: the same forward from the window's words, every row's state left to the keep + try win.reset(0x12); + try p.fill(&win.c, rows); + try p.send(o.stream); + o.arena.reset(); + const y1 = try mixers.linear(o, &planned, w.linear(), &seg, 0, xt, rows, true); + try t.stream.synchronize(); + const y_b = try readWords(win.c.mem, y1.ptr, rows * hidden * act.size() / 4); + defer gpa.free(y_b); + try std.testing.expectEqualSlices(u32, y_a, y_b); + // what the plan's round left where it found it: a one-row window writes it back, a longer one leaves the start + const plan_left = try readWords(win.c.mem, l0.state, state_words); + defer gpa.free(plan_left); + const plan_tail = try readWords(win.c.mem, l0.conv, taps * channels); + defer gpa.free(plan_tail); + // the round writes its state and tail back only for a window of one row; a longer one leaves the start for its keep + if (rows == 1) { + try std.testing.expectEqualSlices(u32, snaps.states[0], plan_left); + try std.testing.expectEqualSlices(u32, snaps.convs[0], plan_tail); + } else { + try std.testing.expectEqualSlices(u32, snaps.start, plan_left); + try std.testing.expectEqualSlices(u32, snaps.start_conv, plan_tail); + } + // a one-row window's round left its state in place; a longer one keeps from the start its round left + if (rows > 1) for (1..rows + 1) |n| try keepCase(t, p, win, o, n, snaps.states[n - 1], snaps.convs[n - 1]); +} + +/// A keep of the window's first `n` rows: the state and conv tail it leaves are today's snapshot of row n. +fn keepCase( + t: *Rig, + p: *gdn.Plan, + win: *Window, + o: hip.ops.Ops, + n: usize, + want_state: []const u32, + want_conv: []const u32, +) !void { + // the keep reads the round's samples and kept rows, which the plan's forward above wrote + try win.reset(0x12); + try p.fill(&win.c, p.rows); + try p.send(o.stream); + try p.keep(o, n); + try t.stream.synchronize(); + const l0 = win.c.linearAt(0); + const got_state = try readWords(win.c.mem, l0.state, state_words); + defer gpa.free(got_state); + try std.testing.expectEqualSlices(u32, want_state, got_state); + const got_conv = try readWords(win.c.mem, l0.conv, taps * channels); + defer gpa.free(got_conv); + try std.testing.expectEqualSlices(u32, want_conv, got_conv); +} From bc466210130747b1ee8a4295abd1166707013eea Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 12 Nov 2026 09:31:55 +0100 Subject: [PATCH 049/186] engine(flashnext): free a plan's slots when opening fails after them, and keep the kernel-test wiring out of hip.zig --- zig/build/flashnext_reach.zig | 13 +++++++++++++ zig/build/hip.zig | 6 +----- zig/src/families/flashnext/hip/engine.zig | 7 ++++++- 3 files changed, 20 insertions(+), 6 deletions(-) create mode 100644 zig/build/flashnext_reach.zig diff --git a/zig/build/flashnext_reach.zig b/zig/build/flashnext_reach.zig new file mode 100644 index 00000000000..9a4b28562f5 --- /dev/null +++ b/zig/build/flashnext_reach.zig @@ -0,0 +1,13 @@ +//! What the model-free kernel tests need to reach the Flash Next HIP kernels they check against its own plan. +const std = @import("std"); + +/// `kernel_module` gains `core`, the memory sizing, and the Flash Next kernel reach built over both. +pub fn wire(b: *std.Build, target: std.Build.ResolvedTarget, kernel_module: *std.Build.Module, core: *std.Build.Module) void { + kernel_module.addImport("core", core); + kernel_module.addImport("cuda_memory", b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_memory.zig"), .target = target })); + const reach = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/kernel_reach.zig"), .target = target, .link_libc = true, .imports = &.{ + .{ .name = "hip", .module = kernel_module }, + .{ .name = "core", .module = core }, + } }); + kernel_module.addImport("flashnext_kernel", reach); +} diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 968da0cbea9..70c3f26e360 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -125,11 +125,7 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil kernel_module.addIncludePath(.{ .cwd_relative = hip_include }); kernel_module.addCSourceFile(.{ .file = b.path("zig/src/hip/device_arch.c"), .flags = &.{"-D__HIP_PLATFORM_AMD__"} }); kernel_module.addImport("hip_kernels", objects); - const core = coreModule(b, target); - kernel_module.addImport("core", core); - kernel_module.addImport("cuda_memory", b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_memory.zig"), .target = target })); - const reach_module = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/kernel_reach.zig"), .target = target, .link_libc = true, .imports = &.{ .{ .name = "hip", .module = kernel_module }, .{ .name = "core", .module = core } } }); - kernel_module.addImport("flashnext_kernel", reach_module); + @import("flashnext_reach.zig").wire(b, target, kernel_module, coreModule(b, target)); const kernel_test = b.addTest(.{ .root_module = kernel_module }); b.step("hip-kernel-build", "Compile the model-free kernels and their GPU tests without running them").dependOn(&kernel_test.step); b.step("hip-kernel-test", "Every model-free kernel against its host reference on the GPU").dependOn(&b.addRunArtifact(kernel_test).step); diff --git a/zig/src/families/flashnext/hip/engine.zig b/zig/src/families/flashnext/hip/engine.zig index 027c5c1d558..4649fca4e00 100644 --- a/zig/src/families/flashnext/hip/engine.zig +++ b/zig/src/families/flashnext/hip/engine.zig @@ -185,14 +185,19 @@ pub const Engine = struct { e.chosen = null; e.pool = .{}; // a window runs from the plan where its shapes are the plan kernels': else a snapshot a row stays + e.gdn = null; if (gdn.supported(e.weights.model.dims, e.cfg.layers)) { const slots = try gpa.create(gdn.Slots); errdefer gpa.destroy(slots); slots.* = try gdn.Slots.init(gpa, &e.driver, e.weights.model.dims, @max(o.want, 1)); - errdefer slots.deinit(); e.gdn = slots; e.weights.model.gdn = slots; } + errdefer if (e.gdn) |p| { + p.deinit(); + gpa.destroy(p); + e.weights.model.gdn = null; + }; _ = try steps.done(.{ 0, 0 }, .first); return e; } From 244a18841331c829f4e7870405bb054565d8716c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:35:44 +0000 Subject: [PATCH 050/186] changelog: back under Unreleased --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d66ed367a92..0994294d38f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. +- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). - Chat templates: a template is held to a nesting depth at load, and its render to a step count, stack depth and scratch memory, so a crafted template can't hang or exhaust the server; real templates use a small fraction of each limit, render byte-identically, and render as fast as before (#3072). @@ -75,7 +77,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. -- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. @@ -231,7 +232,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## 1.2.6 - Prompt cache: long conversations survive 8 to 16 agents sending follow-ups at once. Redundant states are evicted before each conversation's newest one, admission never evicts a state a queued request resumes, concurrent keeps reserve their room, and the default budget scales with the room left after streams on Metal and unified-memory CUDA (discrete CUDA hosts keep 8 GiB, Flash Next's batch host 16 GiB). On Qwen3.8-27B with DFlash2 at 16 agents, a follow-up's median first token drops from 546 s to 15.9 s on an M3 Ultra and from 105 s to 25.8 s on a GB10, with every reply unchanged (#2699). -- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Evidence index: reads saved native capabilities, including unknown and XPU chip labels; advertised facts stay unqualified, and invalid supplied input stops with a bounded error (#2708). From ad0e249a69c5069175cb1492073b6937aebf1d41 Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:07:01 -0500 Subject: [PATCH 051/186] cortheon: keep a page link when a space precedes href's equals --- zig/src/server/cortheon/html.zig | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/zig/src/server/cortheon/html.zig b/zig/src/server/cortheon/html.zig index 9c5acc4d54b..d724917b7b0 100644 --- a/zig/src/server/cortheon/html.zig +++ b/zig/src/server/cortheon/html.zig @@ -307,8 +307,11 @@ pub fn hrefs(a: Allocator, page: []const u8, max: usize, budget: ?*Budget) Alloc const end = tagEnd(page, lt); from = end; // the next link starts after this tag, so an unclosed one is read once const tag = page[lt..end]; - const at = std.ascii.findIgnoreCase(tag, "href=") orelse continue; - var v = tag[at + 5 ..]; + const at = std.ascii.findIgnoreCase(tag, "href") orelse continue; // HREF is the same attribute + var after = at + 4; + while (after < tag.len and std.ascii.isWhitespace(tag[after])) after += 1; // href = is the same attribute + if (after >= tag.len or tag[after] != '=') continue; + var v = tag[after + 1 ..]; const q: u8 = if (v.len > 0 and (v[0] == '"' or v[0] == '\'')) v[0] else ' '; if (q != ' ') v = v[1..]; const stop = std.mem.indexOfAny(u8, v, if (q == ' ') " >" else &.{q}) orelse v.len; @@ -455,6 +458,15 @@ test "a data URL is not kept as a page link" { try std.testing.expectEqualStrings("/docs", got[0]); } +test "a page link keeps its target when a space precedes href's equals" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const got = try hrefs(a, "n", 10, null); + try std.testing.expectEqual(@as(usize, 1), got.len); + try std.testing.expectEqualStrings("https://b.org/x", got[0]); +} + test "a link broken onto the next line by a carriage return is kept" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From c4555f710e064b74f97d3785c83c194495ff8f0a Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:07:21 -0500 Subject: [PATCH 052/186] cortheon: keep a page link when a newline follows href's equals --- zig/src/server/cortheon/html.zig | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/zig/src/server/cortheon/html.zig b/zig/src/server/cortheon/html.zig index d724917b7b0..1e5ae3f883c 100644 --- a/zig/src/server/cortheon/html.zig +++ b/zig/src/server/cortheon/html.zig @@ -312,6 +312,7 @@ pub fn hrefs(a: Allocator, page: []const u8, max: usize, budget: ?*Budget) Alloc while (after < tag.len and std.ascii.isWhitespace(tag[after])) after += 1; // href = is the same attribute if (after >= tag.len or tag[after] != '=') continue; var v = tag[after + 1 ..]; + while (v.len > 0 and std.ascii.isWhitespace(v[0])) v = v[1..]; // a newline may sit before the quotes const q: u8 = if (v.len > 0 and (v[0] == '"' or v[0] == '\'')) v[0] else ' '; if (q != ' ') v = v[1..]; const stop = std.mem.indexOfAny(u8, v, if (q == ' ') " >" else &.{q}) orelse v.len; @@ -467,6 +468,15 @@ test "a page link keeps its target when a space precedes href's equals" { try std.testing.expectEqualStrings("https://b.org/x", got[0]); } +test "a page link keeps its target when a newline follows href's equals" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const got = try hrefs(a, "n", 10, null); + try std.testing.expectEqual(@as(usize, 1), got.len); + try std.testing.expectEqualStrings("https://b.org/y", got[0]); +} + test "a link broken onto the next line by a carriage return is kept" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From 3a15de2c57190aa7b230cdee32277683d5ebe88a Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:07:46 -0500 Subject: [PATCH 053/186] cortheon: end an unquoted page link at a tab, CR or LF --- zig/src/server/cortheon/html.zig | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/zig/src/server/cortheon/html.zig b/zig/src/server/cortheon/html.zig index 1e5ae3f883c..ec1919f9bc9 100644 --- a/zig/src/server/cortheon/html.zig +++ b/zig/src/server/cortheon/html.zig @@ -315,7 +315,7 @@ pub fn hrefs(a: Allocator, page: []const u8, max: usize, budget: ?*Budget) Alloc while (v.len > 0 and std.ascii.isWhitespace(v[0])) v = v[1..]; // a newline may sit before the quotes const q: u8 = if (v.len > 0 and (v[0] == '"' or v[0] == '\'')) v[0] else ' '; if (q != ' ') v = v[1..]; - const stop = std.mem.indexOfAny(u8, v, if (q == ' ') " >" else &.{q}) orelse v.len; + const stop = std.mem.indexOfAny(u8, v, if (q == ' ') " \t\r\n>" else &.{q}) orelse v.len; const raw = std.mem.trim(u8, v[0..stop], " \t\r\n"); if (raw.len == 0 or raw[0] == '#' or std.ascii.startsWithIgnoreCase(raw, "javascript:") or std.ascii.startsWithIgnoreCase(raw, "mailto:") or std.ascii.startsWithIgnoreCase(raw, "data:")) continue; // a data URL is not a page var text: std.ArrayList(u8) = .empty; @@ -477,6 +477,17 @@ test "a page link keeps its target when a newline follows href's equals" { try std.testing.expectEqualStrings("https://b.org/y", got[0]); } +test "an unquoted page link ends at a tab" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const got = try hrefs(a, "z", 10, null); + try std.testing.expectEqual(@as(usize, 1), got.len); + try std.testing.expectEqualStrings("https://b.org/z", got[0]); + const cr = try hrefs(a, "w", 10, null); + try std.testing.expectEqualStrings("https://b.org/w", cr[0]); +} + test "a link broken onto the next line by a carriage return is kept" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From 74beae46c74def246ec02b875ccab5dcd3565f51 Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:08:07 -0500 Subject: [PATCH 054/186] cortheon: resolve a link whose ends are a newline --- zig/src/server/cortheon/web.zig | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/zig/src/server/cortheon/web.zig b/zig/src/server/cortheon/web.zig index 4fffb831c15..240f7757560 100644 --- a/zig/src/server/cortheon/web.zig +++ b/zig/src/server/cortheon/web.zig @@ -274,7 +274,7 @@ fn hasScheme(url: []const u8) bool { /// `location` resolved against `base` (absolute, scheme-relative, root-relative or path-relative). pub fn resolveUrl(a: Allocator, base: []const u8, location: []const u8) Allocator.Error![]const u8 { - const loc = std.mem.trim(u8, location, " \t"); + const loc = std.mem.trim(u8, location, &std.ascii.whitespace); // a newline at either end is not part of the link if (std.ascii.startsWithIgnoreCase(loc, "http://") or std.ascii.startsWithIgnoreCase(loc, "https://")) return withoutDots(a, try a.dupe(u8, loc)); if (hasScheme(loc)) return a.dupe(u8, loc); // an absolute URL of another scheme stays itself const scheme_end = std.mem.indexOf(u8, base, "://") orelse return a.dupe(u8, loc); @@ -323,6 +323,15 @@ test "redirect locations resolve against the page" { try std.testing.expectEqualStrings("https://a.org/login?next=https://a.org/x", try resolveUrl(a, "https://a.org/docs/intro", "/login?next=https://a.org/x")); } +test "a link whose ends are a newline still resolves" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try std.testing.expectEqualStrings("https://b.org/x", try resolveUrl(a, "https://a.org/p/q", "\nhttps://b.org/x")); + try std.testing.expectEqualStrings("https://b.org/y", try resolveUrl(a, "https://a.org/p/q", "https://b.org/y\r")); + try std.testing.expectEqualStrings("https://a.org/root", try resolveUrl(a, "https://a.org/p/q", "\n/root")); +} + test "a link that only replaces the query or the fragment stays on its page" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From 9d5258b9f1e4c2b888df2e29c63eac8997f772b8 Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 07:08:40 -0500 Subject: [PATCH 055/186] cortheon: drop a tab or newline wrapped inside a link --- zig/src/server/cortheon/web.zig | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/zig/src/server/cortheon/web.zig b/zig/src/server/cortheon/web.zig index 240f7757560..d23394d5089 100644 --- a/zig/src/server/cortheon/web.zig +++ b/zig/src/server/cortheon/web.zig @@ -272,9 +272,16 @@ fn hasScheme(url: []const u8) bool { return false; } +/// `location` with tab and newline removed, then other whitespace trimmed off the ends. +fn cleanLocation(a: Allocator, location: []const u8) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + for (location) |ch| if (!(ch == '\t' or ch == '\n' or ch == '\r')) try out.append(a, ch); + return std.mem.trim(u8, out.items, &std.ascii.whitespace); +} + /// `location` resolved against `base` (absolute, scheme-relative, root-relative or path-relative). pub fn resolveUrl(a: Allocator, base: []const u8, location: []const u8) Allocator.Error![]const u8 { - const loc = std.mem.trim(u8, location, &std.ascii.whitespace); // a newline at either end is not part of the link + const loc = try cleanLocation(a, location); if (std.ascii.startsWithIgnoreCase(loc, "http://") or std.ascii.startsWithIgnoreCase(loc, "https://")) return withoutDots(a, try a.dupe(u8, loc)); if (hasScheme(loc)) return a.dupe(u8, loc); // an absolute URL of another scheme stays itself const scheme_end = std.mem.indexOf(u8, base, "://") orelse return a.dupe(u8, loc); @@ -332,6 +339,15 @@ test "a link whose ends are a newline still resolves" { try std.testing.expectEqualStrings("https://a.org/root", try resolveUrl(a, "https://a.org/p/q", "\n/root")); } +test "a tab or newline wrapped inside a link is not part of it" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + try std.testing.expectEqualStrings("https://b.org/path", try resolveUrl(a, "https://a.org/p/q", "https://b.org/pa\nth")); + try std.testing.expectEqualStrings("https://b.org/path", try resolveUrl(a, "https://a.org/p/q", "https://b.org/pa\tth")); + try std.testing.expectEqualStrings("https://b.org/a b", try resolveUrl(a, "https://a.org/p/q", "https://b.org/a b")); +} + test "a link that only replaces the query or the fragment stays on its page" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From 190136209bac85c6288c83a1ae01178b7f3a1c6f Mon Sep 17 00:00:00 2001 From: ppardus <138499387+ppardus@users.noreply.github.com> Date: Thu, 12 Nov 2026 10:08:08 -0500 Subject: [PATCH 056/186] cortheon: keep a page link when an earlier name only contains href --- zig/src/server/cortheon/html.zig | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/zig/src/server/cortheon/html.zig b/zig/src/server/cortheon/html.zig index ec1919f9bc9..2be2902d4de 100644 --- a/zig/src/server/cortheon/html.zig +++ b/zig/src/server/cortheon/html.zig @@ -307,11 +307,19 @@ pub fn hrefs(a: Allocator, page: []const u8, max: usize, budget: ?*Budget) Alloc const end = tagEnd(page, lt); from = end; // the next link starts after this tag, so an unclosed one is read once const tag = page[lt..end]; - const at = std.ascii.findIgnoreCase(tag, "href") orelse continue; // HREF is the same attribute - var after = at + 4; - while (after < tag.len and std.ascii.isWhitespace(tag[after])) after += 1; // href = is the same attribute - if (after >= tag.len or tag[after] != '=') continue; - var v = tag[after + 1 ..]; + var scan: usize = 0; + var eq: ?usize = null; + while (std.ascii.findIgnoreCasePos(tag, scan, "href")) |name| { // HREF is the same attribute + var after = name + 4; + while (after < tag.len and std.ascii.isWhitespace(tag[after])) after += 1; // href = is the same attribute + if (after < tag.len and tag[after] == '=') { + eq = after; + break; + } + scan = name + 1; // hreflang, or a value that only says href, is not the attribute + } + const at = eq orelse continue; + var v = tag[at + 1 ..]; while (v.len > 0 and std.ascii.isWhitespace(v[0])) v = v[1..]; // a newline may sit before the quotes const q: u8 = if (v.len > 0 and (v[0] == '"' or v[0] == '\'')) v[0] else ' '; if (q != ' ') v = v[1..]; @@ -488,6 +496,17 @@ test "an unquoted page link ends at a tab" { try std.testing.expectEqualStrings("https://b.org/w", cr[0]); } +test "a page link keeps its target when href is not the first href in the tag" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const got = try hrefs(a, "de t n", 10, null); + try std.testing.expectEqual(@as(usize, 3), got.len); + try std.testing.expectEqualStrings("https://b.org/de", got[0]); + try std.testing.expectEqualStrings("/t", got[1]); + try std.testing.expectEqualStrings("/n", got[2]); +} + test "a link broken onto the next line by a carriage return is kept" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From cc0d6f1fa758d03a60cbb974d4a4f374fc804c86 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:41 +0000 Subject: [PATCH 057/186] changelog: #3006 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0994294d38f..da520d14db6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. + - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). From 7be6d084c492e05d3dd155b919d9f9bedc2533ce Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:41 +0000 Subject: [PATCH 058/186] changelog: #3006 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index da520d14db6..ade198af08c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. + + - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). From 78140ea84fcd1ddac9d3599cc790885f3904c026 Mon Sep 17 00:00:00 2001 From: Javier <177527508+javierbits@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:49:20 +0000 Subject: [PATCH 059/186] core: a bounded ladder of intermediate prompt-cache checkpoints A long prompt keeps only its history, shared and rewind marks, so a later request that preserves a large prefix after compaction has no state inside that prefix and prefills from the last shared cut. Keep up to `intermediate_most` checkpoints at power-of-two multiples of the ladder's spacing under the history mark. They are admitted after the endpoint and rewind, inside the store's budget and slots, and are exempt from superseded eviction, so they stay across append turns while remaining evictable as ordinary oldest states. Flash Next on CUDA opts in at 8,192-token spacing with four checkpoints, behind `TF_PROMPT_INTERMEDIATE=1` until a served A/B shows the extra marks pay. Host tests cover the ladder plan, a compacted 160K-token prompt resuming at 64K against the 4K shared cut without it, an append turn, admission beside the endpoint, and eviction. --- CHANGELOG.md | 2 + zig/src/core/prompt_cache.zig | 15 +- .../core/prompt_cache_intermediate_test.zig | 149 ++++++++++++++++++ zig/src/core/prompt_cache_keep.zig | 2 +- zig/src/core/prompt_cache_policy.zig | 2 +- zig/src/core/prompt_intermediate.zig | 109 +++++++++++++ zig/src/core/prompt_rules.zig | 6 + zig/src/families/flashnext/cuda_lanes.zig | 3 + zig/src/native/cuda_snaps.zig | 13 ++ 9 files changed, 295 insertions(+), 6 deletions(-) create mode 100644 zig/src/core/prompt_cache_intermediate_test.zig create mode 100644 zig/src/core/prompt_intermediate.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index ade198af08c..f7a0a599cb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1136,6 +1136,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can let a family count repeated image or keyed prompts as repeats, under the cache's own key. No family uses it yet (#1929, part of #1927). Thanks to @laudney. +- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes at the longest checkpoint inside its retained prefix instead of the shared cut, and the checkpoints stay across append turns. Host tests cover the plan, admission, retention and eviction; the Flash Next on CUDA opt-in (`TF_PROMPT_INTERMEDIATE=1`) is not yet measured on a GPU (#1909). + - HIP Qwen3.5-family servers price shared draft rounds at the row bucket they actually run, so Qwen3.6-35B-A3B with MTP at four streams no longer drops into a slower drafting mode, and models load faster. Tokens are unchanged (#1830, fixes #1695). Thanks to @AbrahamCete. - The Living Weights view reads a shard's header with core's safetensors reader, so a malformed header, length word or offset is refused with an error instead of a panic or a 64 GiB allocation, and its sums and row counts are checked (#1923, part of #1872). Thanks to @jschmied. diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index c2d88913807..8601a009956 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -9,6 +9,7 @@ const lifetime = @import("prompt_cache_lifetime.zig"); const policy = @import("prompt_cache_policy.zig"); const cuts = @import("prompt_cuts.zig"); const interior = @import("prompt_interior.zig"); +const intermediate = @import("prompt_intermediate.zig"); const index = @import("prompt_cache_index.zig"); const forks = @import("prompt_cache_fork.zig"); pub const others = @import("prompt_cache_others.zig"); @@ -91,6 +92,7 @@ pub const Entry = struct { shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it rewind: bool = false, // selected as a rewind: a reply born with this request does not supersede it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded + intermediate: bool = false, // a ladder checkpoint under the history: kept across append turns, never superseded expired: bool = false, // hidden from lookups immediately, reclaimed only at safe host boundaries borrowers: u32 = 0, // outstanding lookups: released by resumed or release, never evicted while borrowed ext: std.ArrayList(*Entry) = .empty, // live extensions born no earlier @@ -330,16 +332,20 @@ pub const Store = struct { a.free(ordinary); return err; }; + const marked = intermediate.add(s, a, prompt, from, history_len, starts, spans, marks_, endpoint, endpoint_bytes, e) catch |err| { + a.free(marks_); + return err; + }; for (shared) |w| { // the shared cuts this pass keeps, and a short prompt's mark inside a shared block: other conversations resume them const at = if (s.rules.planned) s.startAt(starts, w) else w; - for (marks_) |m| if (m == at or (m < at and !s.rules.planned and prompt.len < s.rules.min_prompt)) s.noteShared(prompt[0 .. m + s.rules.lookahead]); + for (marked) |m| if (m == at or (m < at and !s.rules.planned and prompt.len < s.rules.min_prompt)) s.noteShared(prompt[0 .. m + s.rules.lookahead]); } var primary = @import("prompt_cache_reserve.zig").endpointFor(s, prompt, from, history_len, starts, spans); - if (history_len == 0 and primary == 0) for (marks_) |at| { + if (history_len == 0 and primary == 0) for (marked) |at| { if (at != s.active_rewind_at) primary = @max(primary, at); }; - s.reserve(prompt, born, e, marks_, spans, primary); - const fitted = try @import("prompt_cache_reserve.zig").fitting(s, a, prompt, marks_, spans, e, born, primary); + s.reserve(prompt, born, e, marked, spans, primary); + const fitted = try @import("prompt_cache_reserve.zig").fitting(s, a, prompt, marked, spans, e, born, primary); if (e) |entry| entry.borrowers += 1; return .{ .entry = e, .marks = fitted }; } @@ -510,6 +516,7 @@ test { _ = @import("prompt_cache_learn_test.zig"); _ = @import("prompt_cache_rewind_test.zig"); _ = @import("prompt_cache_interior_test.zig"); + _ = @import("prompt_cache_intermediate_test.zig"); _ = @import("prompt_cache_branch_test.zig"); _ = @import("prompt_cache_newest_test.zig"); _ = @import("prompt_cache_fork_test.zig"); diff --git a/zig/src/core/prompt_cache_intermediate_test.zig b/zig/src/core/prompt_cache_intermediate_test.zig new file mode 100644 index 00000000000..2ced4b137f7 --- /dev/null +++ b/zig/src/core/prompt_cache_intermediate_test.zig @@ -0,0 +1,149 @@ +//! The bounded ladder of intermediate checkpoints: planned, admitted beside the endpoint, resumed when compacted. +const std = @import("std"); +const pc = @import("prompt_cache.zig"); +const im = @import("prompt_intermediate.zig"); +const Fake = @import("prompt_cache_fake.zig").Fake; +const fresh = @import("prompt_cache_test.zig").fresh; + +const ladder: pc.Rules = .{ .chunk = 2048, .intermediate = 8192, .intermediate_most = 4 }; + +test "ladder takes the largest powers-of-two positions under the history" { + const gpa = std.testing.allocator; + const cases = [_]struct { rules: pc.Rules, prompt: usize, history: u32, from: u32, marks: []const u32, want: []const u32 }{ + .{ .rules = ladder, .prompt = 160 * 1024, .history = 160 * 1024 - 8, .from = 0, .marks = &.{}, .want = &.{ 32768, 65536, 131072 } }, + .{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{}, .want = &.{32768} }, + // An existing mark and the resume point drop their positions; a short prompt or the default rules plan none. + .{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{32768}, .want = &.{} }, + .{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 32768, .marks = &.{}, .want = &.{} }, + .{ .rules = ladder, .prompt = 40 * 1024, .history = 40 * 1024, .from = 0, .marks = &.{40000}, .want = &.{32768} }, + .{ .rules = ladder, .prompt = 8192, .history = 8192, .from = 0, .marks = &.{}, .want = &.{} }, + .{ .rules = .{ .chunk = 2048 }, .prompt = 160 * 1024, .history = 160 * 1024, .from = 0, .marks = &.{}, .want = &.{} }, + }; + for (cases) |c| { + const got = try im.cuts(gpa, c.rules, c.prompt, c.history, c.from, c.marks); + defer gpa.free(got); + try std.testing.expectEqualSlices(u32, c.want, got); + } +} + +test "a compacted long history resumes its longest ladder checkpoint exactly" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 160 * 1024); + for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521); + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), ladder, 1 << 20); + defer s.deinit(); + const history: u32 = @intCast(prompt.len - 8); + const p = try s.beginUser(a, prompt, history, 0, 0, &.{4096}, &.{}, null, &.{}); + try std.testing.expectEqualSlices(u32, &.{ 4096, 32768, 65536, 131072, history }, p.marks); + try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); + const compacted = try a.dupe(u32, prompt[0 .. 120 * 1024]); + for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 65521); + const q = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{4096}, &.{}, null, &.{}); + try std.testing.expectEqual(@as(u32, 65536), q.from); + try std.testing.expectEqual(fresh(compacted), f.pass(&s, compacted, q)); + try std.testing.expectEqual(@as(u64, 1), s.counts.hits); +} + +test "without the ladder a compacted long history resumes only its shared cut" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 160 * 1024); + for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521); + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), .{ .chunk = 2048 }, 1 << 20); + defer s.deinit(); + const history: u32 = @intCast(prompt.len - 8); + const p = try s.beginUser(a, prompt, history, 0, 0, &.{4096}, &.{}, null, &.{}); + try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); + const compacted = try a.dupe(u32, prompt[0 .. 120 * 1024]); + for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 65521); + const q = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{4096}, &.{}, null, &.{}); + try std.testing.expectEqual(@as(u32, 4096), q.from); +} + +test "ladder checkpoints stay across an append turn and are evictable" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 160 * 1024); + for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521); + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), ladder, 1 << 20); + defer s.deinit(); + const p = try s.beginUser(a, prompt, @intCast(prompt.len - 8), 0, 0, &.{}, &.{}, null, &.{}); + try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); + const rung = for (s.entries.items) |e| { + if (e.at == 65536) break e; + } else null; + try std.testing.expect(rung != null and rung.?.intermediate); + // The next turn extends the prompt; the checkpoint is not superseded by the later state. + const grown = try a.alloc(u32, 200 * 1024); + @memcpy(grown[0..prompt.len], prompt); + for (grown[prompt.len..], 0..) |*t, i| t.* = @intCast(i % 7); + const q = try s.beginUser(a, grown, @intCast(grown.len - 8), 0, 0, &.{}, &.{}, null, &.{}); + try std.testing.expectEqual(@as(u32, @intCast(prompt.len - 8)), q.from); + try std.testing.expectEqual(fresh(grown), f.pass(&s, grown, q)); + const compacted = try a.dupe(u32, grown[0 .. 120 * 1024]); + for (compacted[100 * 1024 ..]) |*t| t.* = @intCast((t.* + 1) % 7); + const r = try s.beginUser(a, compacted, @intCast(compacted.len - 8), 0, 0, &.{}, &.{}, null, &.{}); + try std.testing.expectEqual(@as(u32, 65536), r.from); + try std.testing.expectEqual(fresh(compacted), f.pass(&s, compacted, r)); + // Under pressure the same checkpoint is no immortal state: evicting picks it once better victims are gone. + var picked = false; + while (s.entries.items.len > 0) { + const victim = s.victimBut(null, grown, @intCast(grown.len)) orelse break; + picked = picked or s.entries.items[victim].intermediate; + s.remove(victim); + } + try std.testing.expect(picked); +} + +test "ladder admission keeps the endpoint ahead of the checkpoints" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 160 * 1024); + @memset(prompt, 5); + const history: u32 = @intCast(prompt.len - 8); + const endpoint: u64 = 100 + history; + for ([_]u64{ 0, endpoint - 1, endpoint, endpoint + 100 + 32768, 1 << 20 }) |budget| for ([_]usize{ 1, 2, 3, 6 }) |slots| { + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), ladder, budget); + defer s.deinit(); + s.most = slots; + const p = try s.beginUser(a, prompt, history, 0, 0, &.{}, &.{}, null, &.{}); + const rungs = for (p.marks) |m| { + if (im.is(ladder, prompt.len, m)) break true; + } else false; + try std.testing.expectEqual(budget >= endpoint and slots >= 1, std.mem.indexOfScalar(u32, p.marks, history) != null); + try std.testing.expectEqual(budget >= endpoint + 100 + 32768 and slots >= 2, rungs); + try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); + try std.testing.expect(s.held + s.spare() <= budget); + try std.testing.expect(s.entries.items.len <= slots); + }; +} + +test "a short prompt or a cache-off store plans no ladder" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 8000); + @memset(prompt, 7); + for ([_]pc.Rules{ ladder, .{ .chunk = 2048 } }) |rules| { + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), rules, 1 << 20); + defer s.deinit(); + const p = try s.beginUser(a, prompt, 7992, 16, 7992, &.{}, &.{}, null, &.{}); + for (p.marks) |m| try std.testing.expect(!im.is(rules, prompt.len, m)); + try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); + } +} diff --git a/zig/src/core/prompt_cache_keep.zig b/zig/src/core/prompt_cache_keep.zig index 73cdd729cf8..54870156aab 100644 --- a/zig/src/core/prompt_cache_keep.zig +++ b/zig/src/core/prompt_cache_keep.zig @@ -73,7 +73,7 @@ pub fn keepBorn(s: *Store, prompt: []const u32, born: u32, at: u32, owner: ?*any return s.fail(at, err); }; const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; - e.* = .{ .tokens = last[0..n], .images = image_ids, .image_frame = s.image_frame, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .seen = .of(s, last[0..n], decoded, starts, at), .last = last, .shared = shared, .rewind = at == s.active_rewind_at and sharedKey(last[0..n]) == s.active_rewind_key, .own = last }; + e.* = .{ .tokens = last[0..n], .images = image_ids, .image_frame = s.image_frame, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .seen = .of(s, last[0..n], decoded, starts, at), .last = last, .shared = shared, .rewind = at == s.active_rewind_at and sharedKey(last[0..n]) == s.active_rewind_key, .intermediate = @import("prompt_intermediate.zig").is(s.rules, prompt.len, at), .own = last }; owns_decoded = false; owns_images = false; index.add(s, e) catch { diff --git a/zig/src/core/prompt_cache_policy.zig b/zig/src/core/prompt_cache_policy.zig index 04648781fe6..77310d4fe43 100644 --- a/zig/src/core/prompt_cache_policy.zig +++ b/zig/src/core/prompt_cache_policy.zig @@ -21,7 +21,7 @@ pub fn victimIn(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born: oldest = i; oldest_rewind = rewind; } - if (e.borrowers > 0 or e.shared or e.branched or e == skip or rewind) continue; + if (e.borrowers > 0 or e.shared or e.branched or e.intermediate or e == skip or rewind) continue; const now = (born > e.born or (born == e.born and !e.rewind)) and prompt.len > e.tokens.len and on; if ((now or extended(s, e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } diff --git a/zig/src/core/prompt_intermediate.zig b/zig/src/core/prompt_intermediate.zig new file mode 100644 index 00000000000..4076ef48e16 --- /dev/null +++ b/zig/src/core/prompt_intermediate.zig @@ -0,0 +1,109 @@ +//! A bounded ladder of intermediate checkpoints under the history: a later, compacted prompt resumes inside its prefix. +const std = @import("std"); +const pc = @import("prompt_cache.zig"); +const modes = @import("cache_modes.zig"); +const Allocator = std.mem.Allocator; + +/// Whether `at` belongs to the ladder: a power-of-two multiple of its spacing, under the prompt's history. +pub fn is(r: pc.Rules, prompt_len: usize, at: u32) bool { + if (r.intermediate == 0 or r.intermediate_most == 0 or prompt_len < r.intermediate_min or at < r.intermediate_min) return false; + if (at >= prompt_len or at % r.intermediate != 0) return false; + return std.math.isPowerOfTwo(at / r.intermediate); +} + +/// The largest `intermediate_most` ladder positions below `history`, ascending, away from `from` and `marks`. +pub fn cuts(a: Allocator, r: pc.Rules, prompt_len: usize, history: u32, from: u32, marks: []const u32) ![]u32 { + if (r.intermediate == 0 or r.intermediate_most == 0 or prompt_len < r.intermediate_min or history == 0) return a.alloc(u32, 0); + var all: std.ArrayList(u32) = .empty; + defer all.deinit(a); + var at = r.intermediate; + while (at < history) { + const outside = at < r.intermediate_min or at <= from or at + r.lookahead > prompt_len or history - at < r.min_gap; + var skip = outside; + if (!skip) skip = std.mem.indexOfScalar(u32, marks, at) != null; + if (!skip) skip = for (marks) |m| { + if (@max(m, at) - @min(m, at) < r.min_gap) break true; + } else false; + if (!skip) try all.append(a, at); + if (at > history / 2) break; + at *= 2; + } + if (all.items.len <= @as(usize, r.intermediate_most)) return all.toOwnedSlice(a); + const first = all.items.len - r.intermediate_most; + const out = try a.dupe(u32, all.items[first..]); + return out; +} + +/// Append the ladder's marks that fit beside the endpoint, rewind and existing marks; a new slice in `a`. +pub fn add( + s: *pc.Store, + a: Allocator, + prompt: []const u32, + from: u32, + history: u32, + starts: []const u32, + spans: []const modes.Span, + marks: []const u32, + endpoint: u32, + endpoint_bytes: u64, + restored: ?*pc.Entry, +) ![]const u32 { + // The endpoint keeps its priority: no ladder beside a state that does not fit. + if (endpoint_bytes > 0 and endpoint_bytes > s.budget) return marks; + const candidates = try cuts(a, s.rules, prompt.len, history, from, marks); + defer a.free(candidates); + if (candidates.len == 0) return marks; + var required: std.ArrayList(u32) = .empty; + defer required.deinit(a); + try required.appendSlice(a, marks); + for ([_]u32{ if (endpoint_bytes > 0 and endpoint_bytes <= s.budget) endpoint else 0, s.active_rewind_at }) |p| { + if (p > 0 and std.mem.indexOfScalar(u32, required.items, p) == null) try required.append(a, p); + } + var used: u64 = 0; + var slots: usize = required.items.len; + for (required.items) |p| { + var bytes = s.family.vtable.bytes(s.family.ptr, p); + for (s.entries.items) |e| if (same(e, prompt, p, s.rules.lookahead, spans)) { + bytes = @max(bytes, e.bytes); + break; + }; + used +|= bytes; + } + // Other in-flight restores cannot be evicted. The current restore has not been borrowed by lookup yet. + for (s.entries.items) |e| if (e.borrowers > 0 or e == restored) { + const included = for (required.items) |p| { + if (same(e, prompt, p, s.rules.lookahead, spans)) break true; + } else false; + if (!included) { + used +|= e.bytes; + slots +|= 1; + } + }; + // A family without trim owns its spare storage until it reuses it; conservative admission leaves it charged. + if (s.family.vtable.trim == null) used +|= s.spare(); + var added: usize = 0; + var out: std.ArrayList(u32) = .empty; + defer out.deinit(a); + try out.appendSlice(a, marks); + for (candidates) |at| { + if (added == @as(usize, s.rules.intermediate_most) or slots >= s.most) break; + if (!s.usable(at, starts)) continue; + const bytes = s.family.vtable.bytes(s.family.ptr, at); + if (used > s.budget or bytes > s.budget - used) continue; + try out.append(a, at); + used +|= bytes; + slots +|= 1; + added += 1; + } + if (added == 0) return marks; + std.mem.sort(u32, out.items, {}, std.sort.asc(u32)); + const kept = try out.toOwnedSlice(a); + a.free(marks); + return kept; +} + +fn same(e: *const pc.Entry, prompt: []const u32, at: u32, lookahead: u32, spans: []const modes.Span) bool { + return !e.expired and e.at == at and e.tokens.len == @as(usize, at) + lookahead and + e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]) and + modes.equal(e.decode_spans, spans, at); +} diff --git a/zig/src/core/prompt_rules.zig b/zig/src/core/prompt_rules.zig index c3541f08bbc..e999872a427 100644 --- a/zig/src/core/prompt_rules.zig +++ b/zig/src/core/prompt_rules.zig @@ -30,4 +30,10 @@ pub const Rules = struct { repeat_keyed: bool = false, /// Keep one before where the prompt parts from the one it resumed (a miss: the longest kept one), however near. fork: bool = false, + /// Spacing of a bounded ladder of checkpoints under the history (0: off; only chunk-invariant families opt in). + intermediate: u32 = 0, + /// Checkpoints kept from the ladder at most, the largest positions first. + intermediate_most: u8 = 0, + /// Prompts shorter than this, and ladder positions before it, keep none. + intermediate_min: u32 = 32768, }; diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 725729ebc06..058dfb9939d 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -35,6 +35,9 @@ pub fn Namespace(comptime Tri: type) type { /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. pub const snap_bytes = snaps.bytes; pub const snap_repeat_min: u32 = 16; + /// The prompt cache's ladder of intermediate checkpoints under a long prompt's history mark. + pub const snap_intermediate: u32 = 8192; + pub const snap_intermediate_most: u8 = 4; pub const snap_save = snaps.save; pub const snap_restore = snaps.restore; pub const snap_drop = snaps.drop; diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index a0f1685a307..0b847885e7f 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -9,6 +9,10 @@ pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { if (@hasDecl(F, "warm_quiet")) rules.warm = .waits; if (@hasDecl(F, "snap_fork")) rules.fork = F.snap_fork; // opt-in, a family at a time after its device check if (@hasDecl(F, "snap_tail")) rules.tail = if (@import("core").profile.lookup(F.snap_tail)) |v| std.mem.eql(u8, v, "1") else false; + if (@hasDecl(F, "snap_intermediate")) { // a chunk-invariant family's opt-in ladder of intermediate checkpoints + rules.intermediate = F.snap_intermediate; + rules.intermediate_most = F.snap_intermediate_most; + } if (@hasDecl(F, "snap_grid")) { rules.planned = true; rules.grid = F.snap_grid; @@ -21,6 +25,8 @@ pub fn passRules(comptime F: type, prompt_rows: u32) api.prompt_cache.Rules { var rules = cacheRules(F); rules.chunk = prompt_rows; if (@hasDecl(F, "snap_repeat_min")) rules.repeat_min = if (std.c.getenv("TF_FLASHNEXT_REPEAT_CACHE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else F.snap_repeat_min) else F.snap_repeat_min; + // Opt-in until a served A/B shows the extra marks pay: TF_PROMPT_INTERMEDIATE=1 enables the family's ladder. + if (rules.intermediate > 0) rules.intermediate = if (std.c.getenv("TF_PROMPT_INTERMEDIATE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else rules.intermediate) else 0; rules.chunk_starts = false; return rules; } @@ -76,10 +82,17 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as const Headless = struct { pub const snap_lookahead: u32 = 0; }; + const Ladder = struct { + pub const snap_intermediate: u32 = 8192; + pub const snap_intermediate_most: u8 = 4; + }; try std.testing.expectEqual(@as(u32, 0), cacheRules(Headless).lookahead); const any = cacheRules(Other); try std.testing.expect(!any.planned and any.min_prompt == 0 and any.lookahead == 1 and any.warm == .off); try std.testing.expect(cacheRules(Warm).warm == .waits and cacheRules(Glm).warm == .off); + try std.testing.expect(!any.planned and any.intermediate == 0); + try std.testing.expectEqual(@as(u32, 8192), cacheRules(Ladder).intermediate); + try std.testing.expectEqual(@as(u8, 4), cacheRules(Ladder).intermediate_most); const Stub = struct { fn bytes(_: *anyopaque, _: u32) u64 { return 0; From 08265a539179926b40ad02ca5b98082b42a6659d Mon Sep 17 00:00:00 2001 From: Javier <177527508+javierbits@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:49:20 +0000 Subject: [PATCH 060/186] docs: the intermediate checkpoint ladder's GB10 receipt On one GB10 with nvidia/Qwen3.8-Flash-Next-NVFP4, --no-drafts and an 8 GiB cache, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2 (TTFT 19.9 s against 50.7 s, same tokens). A cold 165K prompt paid 0.9 s more and the cache held two more states. The follow-up resumed 65,536 because the compacted turn had evicted the larger states. --- CHANGELOG.md | 2 +- docs/recipes/flash-next-cuda.md | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f7a0a599cb0..2dde6e18ca4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1136,7 +1136,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can let a family count repeated image or keyed prompts as repeats, under the cache's own key. No family uses it yet (#1929, part of #1927). Thanks to @laudney. -- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes at the longest checkpoint inside its retained prefix instead of the shared cut, and the checkpoints stay across append turns. Host tests cover the plan, admission, retention and eviction; the Flash Next on CUDA opt-in (`TF_PROMPT_INTERMEDIATE=1`) is not yet measured on a GPU (#1909). +- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). - HIP Qwen3.5-family servers price shared draft rounds at the row bucket they actually run, so Qwen3.6-35B-A3B with MTP at four streams no longer drops into a slower drafting mode, and models load faster. Tokens are unchanged (#1830, fixes #1695). Thanks to @AbrahamCete. diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 9f9a7f520e8..2bbee279424 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -132,6 +132,10 @@ The family plugs into the core's pieces and adds no scheduler: `tensorfold:prompt_cache_states` and `tensorfold:prompt_cache_bytes`; a successful resume counts toward `tensorfold:prompt_tokens_cached_total`. The chunk and budget host fixtures exercise this policy; real-model resumed/fresh parity and TTFT for this interior state still need the GPU gate. +- `TF_PROMPT_INTERMEDIATE=1` keeps up to four intermediate states at power-of-two positions under a long prompt's + history, so a later request that keeps a compacted prefix resumes inside it. On one GB10 (NVFP4, `--no-drafts`, + an 8 GiB cache), a 104K-token request keeping its first 100K resumed 65,536 tokens instead of 2 (TTFT 19.9 s + against 50.7 s, same tokens); a cold 165K prompt paid 0.9 s more and the cache held two more states. Off by default. - Shared row kernels: the token rows (`embed_rows.cu`, groups of 32 and a copy a hyper-connection stream), the MTP head's RMS rows with their group sums (`norms.cu`), the router logits and softmax top-k (`route_rows.cu`), and the routed experts over a window's or a chunk's rows on the core's plan (`cuda/grouped.zig`, `experts_rows.cu`). From c8cb885cde785e9a90fe47d2671abf6df8ec2eda Mon Sep 17 00:00:00 2001 From: Javier <177527508+javierbits@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:49:21 +0000 Subject: [PATCH 061/186] core: share the interior key check and test the ladder's opt-in prompt_intermediate used a copy of prompt_interior's same check from before the image-identity change, so it missed cache_images.matches. Make the interior check public and call it, state that any kept state at a rung position counts as intermediate, and test that passRules leaves the ladder off with TF_PROMPT_INTERMEDIATE unset or 0. --- CHANGELOG.md | 2 +- zig/src/core/prompt_cache.zig | 2 +- zig/src/core/prompt_interior.zig | 2 +- zig/src/core/prompt_intermediate.zig | 11 +++-------- zig/src/native/cuda_snaps.zig | 15 +++++++++++++++ 5 files changed, 21 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2dde6e18ca4..7df530ce500 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1136,7 +1136,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can let a family count repeated image or keyed prompts as repeats, under the cache's own key. No family uses it yet (#1929, part of #1927). Thanks to @laudney. -- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). +- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Thanks to @javierbits. - HIP Qwen3.5-family servers price shared draft rounds at the row bucket they actually run, so Qwen3.6-35B-A3B with MTP at four streams no longer drops into a slower drafting mode, and models load faster. Tokens are unchanged (#1830, fixes #1695). Thanks to @AbrahamCete. diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index 8601a009956..6b12c53399f 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -92,7 +92,7 @@ pub const Entry = struct { shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it rewind: bool = false, // selected as a rewind: a reply born with this request does not supersede it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded - intermediate: bool = false, // a ladder checkpoint under the history: kept across append turns, never superseded + intermediate: bool = false, // any kept state at a ladder position: retained across append turns, never superseded expired: bool = false, // hidden from lookups immediately, reclaimed only at safe host boundaries borrowers: u32 = 0, // outstanding lookups: released by resumed or release, never evicted while borrowed ext: std.ArrayList(*Entry) = .empty, // live extensions born no earlier diff --git a/zig/src/core/prompt_interior.zig b/zig/src/core/prompt_interior.zig index 57ba741c373..b20bb4f7905 100644 --- a/zig/src/core/prompt_interior.zig +++ b/zig/src/core/prompt_interior.zig @@ -29,7 +29,7 @@ pub fn cut(r: pc.Rules, len: usize, from: u32, end: u32, user: u32, starts: []co return if (at > user and at > from and at < end and at <= limit) at else 0; } -fn same(s: *const pc.Store, e: *const pc.Entry, prompt: []const u32, at: u32, lookahead: u32, spans: []const modes.Span) bool { +pub fn same(s: *const pc.Store, e: *const pc.Entry, prompt: []const u32, at: u32, lookahead: u32, spans: []const modes.Span) bool { return !e.expired and e.at == at and e.tokens.len == @as(usize, at) + lookahead and e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]) and modes.equal(e.decode_spans, spans, at) and @import("cache_images.zig").matches(s, e); diff --git a/zig/src/core/prompt_intermediate.zig b/zig/src/core/prompt_intermediate.zig index 4076ef48e16..ee2a9159caf 100644 --- a/zig/src/core/prompt_intermediate.zig +++ b/zig/src/core/prompt_intermediate.zig @@ -2,6 +2,7 @@ const std = @import("std"); const pc = @import("prompt_cache.zig"); const modes = @import("cache_modes.zig"); +const interior = @import("prompt_interior.zig"); const Allocator = std.mem.Allocator; /// Whether `at` belongs to the ladder: a power-of-two multiple of its spacing, under the prompt's history. @@ -63,7 +64,7 @@ pub fn add( var slots: usize = required.items.len; for (required.items) |p| { var bytes = s.family.vtable.bytes(s.family.ptr, p); - for (s.entries.items) |e| if (same(e, prompt, p, s.rules.lookahead, spans)) { + for (s.entries.items) |e| if (interior.same(s, e, prompt, p, s.rules.lookahead, spans)) { bytes = @max(bytes, e.bytes); break; }; @@ -72,7 +73,7 @@ pub fn add( // Other in-flight restores cannot be evicted. The current restore has not been borrowed by lookup yet. for (s.entries.items) |e| if (e.borrowers > 0 or e == restored) { const included = for (required.items) |p| { - if (same(e, prompt, p, s.rules.lookahead, spans)) break true; + if (interior.same(s, e, prompt, p, s.rules.lookahead, spans)) break true; } else false; if (!included) { used +|= e.bytes; @@ -101,9 +102,3 @@ pub fn add( a.free(marks); return kept; } - -fn same(e: *const pc.Entry, prompt: []const u32, at: u32, lookahead: u32, spans: []const modes.Span) bool { - return !e.expired and e.at == at and e.tokens.len == @as(usize, at) + lookahead and - e.tokens.len <= prompt.len and std.mem.eql(u32, e.tokens, prompt[0..e.tokens.len]) and - modes.equal(e.decode_spans, spans, at); -} diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index 0b847885e7f..9fe7b17886c 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -1,6 +1,7 @@ //! A CUDA family's prompt cache in the native server: rules, budget and copies from the family's `snap_*` decls. const std = @import("std"); const api = @import("engine_api"); +const env = @import("sys").env; /// Any row, the head one token on (Nemotron's); a family's snap_grid, snap_lookahead, warm_quiet, snap_fork, snap_tail. pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { @@ -137,6 +138,20 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as try std.testing.expect(!cacheRules(Other).tail); } +test "the ladder stays off unless TF_PROMPT_INTERMEDIATE is set" { + const Ladder = struct { + pub const snap_intermediate: u32 = 8192; + pub const snap_intermediate_most: u8 = 4; + }; + _ = env.unsetenv("TF_PROMPT_INTERMEDIATE"); + try std.testing.expectEqual(@as(u32, 0), passRules(Ladder, 512).intermediate); + _ = env.setenv("TF_PROMPT_INTERMEDIATE", "0", 1); + defer _ = env.unsetenv("TF_PROMPT_INTERMEDIATE"); + try std.testing.expectEqual(@as(u32, 0), passRules(Ladder, 512).intermediate); + _ = env.setenv("TF_PROMPT_INTERMEDIATE", "1", 1); + try std.testing.expectEqual(@as(u32, 8192), passRules(Ladder, 512).intermediate); +} + test { _ = @import("cuda_snaps_test.zig"); } From 1d5e7af63f8f8ed2a4055b1894b440db1e6fe304 Mon Sep 17 00:00:00 2001 From: Javier <177527508+javierbits@users.noreply.github.com> Date: Thu, 5 Nov 2026 10:32:46 +0000 Subject: [PATCH 062/186] flashnext: keep cuda_lanes.zig at 600 lines The lean rule caps the file at 600 lines and main is already there. Move the ladder's two values into cuda_snap.zig, expose them to the family through one alias, and drop one blank line between the weight-byte wrappers. --- zig/src/families/flashnext/cuda_lanes.zig | 5 +---- zig/src/families/flashnext/cuda_snap.zig | 3 +++ zig/src/native/cuda_snaps.zig | 10 ++++------ 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 058dfb9939d..1aaf22148e8 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -35,9 +35,7 @@ pub fn Namespace(comptime Tri: type) type { /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. pub const snap_bytes = snaps.bytes; pub const snap_repeat_min: u32 = 16; - /// The prompt cache's ladder of intermediate checkpoints under a long prompt's history mark. - pub const snap_intermediate: u32 = 8192; - pub const snap_intermediate_most: u8 = 4; + pub const snap_intermediate = snap.ladder; pub const snap_save = snaps.save; pub const snap_restore = snaps.restore; pub const snap_drop = snaps.drop; @@ -102,7 +100,6 @@ pub fn Namespace(comptime Tri: type) type { pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { return bytes.deviceWeightBytes(gpa, io, dir); } - pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { return bytes.placementWeightBytes(gpa, io, dir, world, rank, drafts); } diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 4786e59d9b5..768bebd6479 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -14,6 +14,9 @@ const tail_bytes = wt.ple_tail_n * wt.wide * 2; /// The MTP head's cache: its own layer, whole on every rank. const head_heads: Heads = .full; +/// The prompt cache's ladder: tokens between rungs, and the most it keeps under one history mark. +pub const ladder = .{ .tokens = 8192, .most = 4 }; + /// A state's shape for its disk image (cuda_disk.zig, one rank, the whole layer): layer counts and bytes a row. pub const shape = .{ .linear = wt.linear_n, .attn = wt.attn_n, .conv = sizes.convBytes(.full), .state = sizes.stateBytes(.full), .kv_row = Heads.full.kvRow(), .index_row = index_row, .ratio = wt.index_ratio, .tail = tail_bytes }; diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index 9fe7b17886c..6e97d7563f9 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -11,8 +11,8 @@ pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { if (@hasDecl(F, "snap_fork")) rules.fork = F.snap_fork; // opt-in, a family at a time after its device check if (@hasDecl(F, "snap_tail")) rules.tail = if (@import("core").profile.lookup(F.snap_tail)) |v| std.mem.eql(u8, v, "1") else false; if (@hasDecl(F, "snap_intermediate")) { // a chunk-invariant family's opt-in ladder of intermediate checkpoints - rules.intermediate = F.snap_intermediate; - rules.intermediate_most = F.snap_intermediate_most; + rules.intermediate = F.snap_intermediate.tokens; + rules.intermediate_most = F.snap_intermediate.most; } if (@hasDecl(F, "snap_grid")) { rules.planned = true; @@ -84,8 +84,7 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as pub const snap_lookahead: u32 = 0; }; const Ladder = struct { - pub const snap_intermediate: u32 = 8192; - pub const snap_intermediate_most: u8 = 4; + pub const snap_intermediate = .{ .tokens = 8192, .most = 4 }; }; try std.testing.expectEqual(@as(u32, 0), cacheRules(Headless).lookahead); const any = cacheRules(Other); @@ -140,8 +139,7 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as test "the ladder stays off unless TF_PROMPT_INTERMEDIATE is set" { const Ladder = struct { - pub const snap_intermediate: u32 = 8192; - pub const snap_intermediate_most: u8 = 4; + pub const snap_intermediate = .{ .tokens = 8192, .most = 4 }; }; _ = env.unsetenv("TF_PROMPT_INTERMEDIATE"); try std.testing.expectEqual(@as(u32, 0), passRules(Ladder, 512).intermediate); From 29b2089d3754192080d388d91e3a6bd46ea91ffe Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:29:25 +0000 Subject: [PATCH 063/186] core: ladder checkpoints yield to newest states and to other marks' room Since #2699 the victim order keeps each conversation's newest state before redundant copies. A ladder checkpoint was exempt from the superseded pick, so once the plain redundant copies were gone the oldest state went next, which could be another conversation's newest. A superseded checkpoint now goes after every other redundant copy and before any newest state. Under a reservation guard, fitting took the marks in position order, so the low ladder positions could take the room the rewind state needed. Ladder positions now fit last, after the endpoint and every other mark, in the room other passes' promises leave. Co-authored-by: javierbits <177527508+javierbits@users.noreply.github.com> --- zig/src/core/prompt_cache.zig | 2 +- .../core/prompt_cache_intermediate_test.zig | 59 +++++++++++++++++++ zig/src/core/prompt_cache_policy.zig | 9 ++- zig/src/core/prompt_cache_reserve.zig | 8 +-- 4 files changed, 70 insertions(+), 8 deletions(-) diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index 6b12c53399f..ec73923dc58 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -92,7 +92,7 @@ pub const Entry = struct { shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it rewind: bool = false, // selected as a rewind: a reply born with this request does not supersede it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded - intermediate: bool = false, // any kept state at a ladder position: retained across append turns, never superseded + intermediate: bool = false, // any kept state at a ladder position: once superseded, goes after other copies but before any newest state expired: bool = false, // hidden from lookups immediately, reclaimed only at safe host boundaries borrowers: u32 = 0, // outstanding lookups: released by resumed or release, never evicted while borrowed ext: std.ArrayList(*Entry) = .empty, // live extensions born no earlier diff --git a/zig/src/core/prompt_cache_intermediate_test.zig b/zig/src/core/prompt_cache_intermediate_test.zig index 2ced4b137f7..0bb9336018e 100644 --- a/zig/src/core/prompt_cache_intermediate_test.zig +++ b/zig/src/core/prompt_cache_intermediate_test.zig @@ -147,3 +147,62 @@ test "a short prompt or a cache-off store plans no ladder" { try std.testing.expectEqual(fresh(prompt), f.pass(&s, prompt, p)); } } + +fn none(_: *anyopaque, _: *const pc.Entry) bool { + return false; +} + +fn put(f: *Fake, s: *pc.Store, words: []const u32, at: u32) !void { + _ = f.pass(s, words[0..at], .{}); + try std.testing.expect(s.keepBorn(words, @intCast(words.len), at, null, &.{}, &.{})); +} + +test "a superseded ladder checkpoint goes after other redundant copies but before any conversation's newest state" { + const gpa = std.testing.allocator; + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), .{ .min_prompt = 0, .intermediate = 8, .intermediate_most = 4, .intermediate_min = 16 }, 1 << 20); + defer s.deinit(); + var other: [40]u32 = undefined; + var own: [40]u32 = undefined; + for (&other, &own, 0..) |*o, *w, i| { + o.* = @intCast(9000 + i); + w.* = @intCast(1 + i); + } + try put(&f, &s, &other, 36); // the other conversation's newest state, and the oldest + try put(&f, &s, &own, 20); // a plain copy its own later states supersede + try put(&f, &s, &own, 32); // a ladder position: 32 is 8 times 4 + try put(&f, &s, &own, 36); // this conversation's newest state + var order: [4]u32 = undefined; + var from: [4]u32 = undefined; + for (&order, &from) |*at, *first| { + const i = s.victimBut(null, &.{}, 0).?; + at.* = s.entries.items[i].at; + first.* = s.entries.items[i].tokens[0]; + s.remove(i); + } + try std.testing.expectEqualSlices(u32, &.{ 20, 32, 36, 36 }, &order); + try std.testing.expectEqualSlices(u32, &.{ 1, 1, 9000, 1 }, &from); +} + +test "ladder checkpoints fit only in the room the endpoint, the rewind state and other passes' promises leave" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = try a.alloc(u32, 160 * 1024); + for (prompt, 0..) |*t, i| t.* = @intCast(i % 65521); + const history: u32 = @intCast(prompt.len - 8); + const rewind: u32 = history - 4096; + var f: Fake = .{ .gpa = gpa }; + var s = pc.Store.init(gpa, f.snapshots(), ladder, 1 << 20); + defer s.deinit(); + const open = try s.lookupRewind(a, prompt, @intCast(prompt.len), history, rewind, &.{}, &.{}, &.{}); + try std.testing.expectEqualSlices(u32, &.{ 32768, 65536, 131072, rewind, history }, open.marks); + // Other admitted passes hold all but the endpoint's and the rewind state's bytes, and a little more. + const left = (100 + @as(u64, history)) + (100 + @as(u64, rewind)) + 1000; + s.reservation_guard = .{ .ctx = &f, .needed = none, .promised = .{ .bytes = s.budget - left } }; + defer s.reservation_guard = null; + const tight = try s.lookupRewind(a, prompt, @intCast(prompt.len), history, rewind, &.{}, &.{}, &.{}); + try std.testing.expectEqualSlices(u32, &.{ rewind, history }, tight.marks); + try std.testing.expectEqual(@as(u64, 3), s.counts.trimmed); +} diff --git a/zig/src/core/prompt_cache_policy.zig b/zig/src/core/prompt_cache_policy.zig index 77310d4fe43..476c11e143e 100644 --- a/zig/src/core/prompt_cache_policy.zig +++ b/zig/src/core/prompt_cache_policy.zig @@ -9,6 +9,7 @@ pub fn victimBut(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born /// `victimBut`'s order over every state, or with `own_only` over the prompt's own conversation's states alone. pub fn victimIn(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born: u32, own_only: bool) ?usize { var best: ?usize = null; + var rung: ?usize = null; // a superseded ladder checkpoint: after every other redundant copy, before any newest state var oldest: ?usize = null; var oldest_rewind = false; for (s.entries.items, 0..) |e, i| { // one prefix compare an entry: the superseded scan and the oldest scan share it @@ -21,11 +22,13 @@ pub fn victimIn(s: *const pc.Store, skip: ?*pc.Entry, prompt: []const u32, born: oldest = i; oldest_rewind = rewind; } - if (e.borrowers > 0 or e.shared or e.branched or e.intermediate or e == skip or rewind) continue; + if (e.borrowers > 0 or e.shared or e.branched or e == skip or rewind) continue; const now = (born > e.born or (born == e.born and !e.rewind)) and prompt.len > e.tokens.len and on; - if ((now or extended(s, e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; + if (!now and !extended(s, e)) continue; + const pick = if (e.intermediate) &rung else &best; + if (pick.* == null or e.used < s.entries.items[pick.*.?].used) pick.* = i; } - return best orelse oldest; + return best orelse rung orelse oldest; } /// A reservation may free redundant copies, but not a newest state just to plan more copies. diff --git a/zig/src/core/prompt_cache_reserve.zig b/zig/src/core/prompt_cache_reserve.zig index 593c782391d..db94132b941 100644 --- a/zig/src/core/prompt_cache_reserve.zig +++ b/zig/src/core/prompt_cache_reserve.zig @@ -118,7 +118,7 @@ pub fn reserve(s: *pc.Store, prompt: []const u32, born: u32, from: ?*pc.Entry, m } } -/// Select the endpoint first; only interior, rewind and extra shared cuts compete for room after it. +/// The endpoint first; interior, rewind and extra shared cuts compete for the room after it, ladder rungs last. pub fn fitting(s: *pc.Store, a: std.mem.Allocator, prompt: []const u32, marks: []const u32, spans: []const modes.Span, from: ?*pc.Entry, born: u32, endpoint: u32) ![]const u32 { if (s.reservation_guard == null) return marks; defer a.free(marks); @@ -134,15 +134,15 @@ pub fn fitting(s: *pc.Store, a: std.mem.Allocator, prompt: []const u32, marks: [ if (bytes > 0) room.slots -= 1; } } - for (marks) |at| { - if (end == at) continue; + for ([_]bool{ false, true }) |rungs| for (marks) |at| { // ladder checkpoints take only what every other mark leaves + if (end == at or @import("prompt_intermediate.zig").is(s.rules, prompt.len, at) != rungs) continue; const bytes = if (exists(s, prompt, at, spans)) 0 else s.family.vtable.bytes(s.family.ptr, at); if (bytes == 0 or (bytes <= room.bytes and room.slots > 0)) { try out.append(a, at); room.bytes -= bytes; if (bytes > 0) room.slots -= 1; } - } + }; const fitted = try out.toOwnedSlice(a); s.counts.trimmed += marks.len - fitted.len; if (s.reservation_guard.?.promise) |promise| { From 8e98927ee702f1207b11090d53abce414ad5c7a9 Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:29:25 +0000 Subject: [PATCH 064/186] cuda: read TF_PROMPT_INTERMEDIATE through the strict environment table The switch is a CUDA-scope flag now, checked when the backend starts like the other prompt-cache switches, and documented. The test's env handle moves inside the test, since main's snap_tail test declares its own. Co-authored-by: javierbits <177527508+javierbits@users.noreply.github.com> --- docs/environment.md | 1 + zig/src/native/cuda_snaps.zig | 4 ++-- zig/src/platform/environment.zig | 1 + zig/src/platform/environment_test.zig | 2 ++ 4 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/environment.md b/docs/environment.md index 5b937ab5861..0af40a6f12c 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -28,6 +28,7 @@ Diagnostics quote and escape the supplied text so it cannot create another termi | `TF_DRAFT_NODES` | Whole native-size nonnegative count, capped by the drafter's window; zero disables trees | | `TF_QMMF_HALF_BELOW` | Whole native-size nonnegative block count; zero disables half tiles; default 96 | | `TENSORFOLD_PROMPT_CACHE_KEEP_OTHERS` | Exactly `0` or `1`, checked when the CUDA backend starts; default off ([profiles](profiles.md)) | +| `TF_PROMPT_INTERMEDIATE` | Exactly `0` or `1`, checked when the CUDA backend starts; default off. `1` keeps Flash Next's ladder of intermediate prompt-cache checkpoints under a long prompt's history | | `TF_HEAD_GEOMETRY` | Whole native-size nonnegative geometry count in the head benchmark | | `TF_ROCE_GID_INDEX` | Whole index 0 through 2147483647; unset searches for a RoCE v2 IPv4 GID | | `TF_GEMMA4_CHUNK_SOLO` | Gemma 4 on CUDA: a multiple of 512 from 512 through 4096; unset, empty or 512 is off. Checked when the model loads, with the family's own message | diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index 6e97d7563f9..df4527771b9 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -1,7 +1,6 @@ //! A CUDA family's prompt cache in the native server: rules, budget and copies from the family's `snap_*` decls. const std = @import("std"); const api = @import("engine_api"); -const env = @import("sys").env; /// Any row, the head one token on (Nemotron's); a family's snap_grid, snap_lookahead, warm_quiet, snap_fork, snap_tail. pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { @@ -27,7 +26,7 @@ pub fn passRules(comptime F: type, prompt_rows: u32) api.prompt_cache.Rules { rules.chunk = prompt_rows; if (@hasDecl(F, "snap_repeat_min")) rules.repeat_min = if (std.c.getenv("TF_FLASHNEXT_REPEAT_CACHE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else F.snap_repeat_min) else F.snap_repeat_min; // Opt-in until a served A/B shows the extra marks pay: TF_PROMPT_INTERMEDIATE=1 enables the family's ladder. - if (rules.intermediate > 0) rules.intermediate = if (std.c.getenv("TF_PROMPT_INTERMEDIATE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else rules.intermediate) else 0; + if (rules.intermediate > 0 and !(@import("sys").environment.read(bool, "TF_PROMPT_INTERMEDIATE", false) catch false)) rules.intermediate = 0; rules.chunk_starts = false; return rules; } @@ -138,6 +137,7 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as } test "the ladder stays off unless TF_PROMPT_INTERMEDIATE is set" { + const env = @import("sys").env; const Ladder = struct { pub const snap_intermediate = .{ .tokens = 8192, .most = 4 }; }; diff --git a/zig/src/platform/environment.zig b/zig/src/platform/environment.zig index 884df7d5e35..3ad7ecb4941 100644 --- a/zig/src/platform/environment.zig +++ b/zig/src/platform/environment.zig @@ -49,6 +49,7 @@ pub const rules = [_]Rule{ .{ .name = "TF_DRAFT_NODES", .scope = .cuda, .kind = .size, .forms = "a whole nonnegative native-size node count (capped by the drafter's window)" }, .{ .name = "TF_QMMF_HALF_BELOW", .scope = .cuda, .kind = .size, .forms = "a whole nonnegative native-size block count (0 disables half tiles)" }, .{ .name = "TENSORFOLD_PROMPT_CACHE_KEEP_OTHERS", .scope = .cuda, .kind = .flag, .forms = "0 or 1" }, + .{ .name = "TF_PROMPT_INTERMEDIATE", .scope = .cuda, .kind = .flag, .forms = "0 or 1 (default 0)" }, .{ .name = "TF_ROCE_GID_INDEX", .scope = .cuda, .kind = .u32, .forms = "a whole GID index from 0 to 2147483647" }, .{ .name = "TF_GEMMA4_CHUNK_SOLO", .scope = .gemma4, .kind = .size, .forms = "a solo prompt pass's rows, a multiple of 512 up to 4096" }, .{ .name = "TF_DRAMA_PREFILL_ROWS", .scope = .drama3, .kind = .u32, .trim = " ", .forms = "whole prompt rows a chunk, 1 through 1024 (checked when a Drama 3 checkpoint opens)" }, diff --git a/zig/src/platform/environment_test.zig b/zig/src/platform/environment_test.zig index 04093602fc8..aab423a84b2 100644 --- a/zig/src/platform/environment_test.zig +++ b/zig/src/platform/environment_test.zig @@ -65,6 +65,8 @@ const valid = [_]Case{ .{ .name = "TF_QMMF_HALF_BELOW", .raw = "0", .n = 0 }, .{ .name = "TENSORFOLD_PROMPT_CACHE_KEEP_OTHERS", .raw = "0", .n = 0 }, .{ .name = "TENSORFOLD_PROMPT_CACHE_KEEP_OTHERS", .raw = "1", .n = 1 }, + .{ .name = "TF_PROMPT_INTERMEDIATE", .raw = "0", .n = 0 }, + .{ .name = "TF_PROMPT_INTERMEDIATE", .raw = "1", .n = 1 }, .{ .name = "TF_ROCE_GID_INDEX", .raw = "2147483647", .n = 2147483647 }, .{ .name = "TF_DRAMA_PREFILL_ROWS", .raw = " 128 ", .n = 128 }, .{ .name = "TF_DRAMA_REFERENCE_CACHE", .raw = "0", .n = 0 }, From 68cde09769b683e02bafbb241df431e8f549fcb9 Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:29:25 +0000 Subject: [PATCH 065/186] changelog: the ladder's line under Unreleased The rebase put it under 1.2.2. It now also says how its checkpoints are evicted and fitted. Co-authored-by: javierbits <177527508+javierbits@users.noreply.github.com> --- CHANGELOG.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7df530ce500..e9bb73e4fa3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -78,6 +78,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. + +- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. + - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. @@ -1136,8 +1139,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can let a family count repeated image or keyed prompts as repeats, under the cache's own key. No family uses it yet (#1929, part of #1927). Thanks to @laudney. -- The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Thanks to @javierbits. - - HIP Qwen3.5-family servers price shared draft rounds at the row bucket they actually run, so Qwen3.6-35B-A3B with MTP at four streams no longer drops into a slower drafting mode, and models load faster. Tokens are unchanged (#1830, fixes #1695). Thanks to @AbrahamCete. - The Living Weights view reads a shard's header with core's safetensors reader, so a malformed header, length word or offset is refused with an error instead of a panic or a 64 GiB allocation, and its sums and row counts are checked (#1923, part of #1872). Thanks to @jschmied. From e3062d283e4eedea3a3116cc42e67aaeb4d8a42a Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:27:55 +0000 Subject: [PATCH 066/186] tests: re-pin cuda_lanes.zig and cuda_snap.zig for #1953's checkpoint ladder --- zig/tests/fixtures/gradient-serving.tsv | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/tests/fixtures/gradient-serving.tsv b/zig/tests/fixtures/gradient-serving.tsv index bb4da22d6c5..60b83c49883 100644 --- a/zig/tests/fixtures/gradient-serving.tsv +++ b/zig/tests/fixtures/gradient-serving.tsv @@ -129,7 +129,7 @@ zig/src/families/flashnext/cuda_hc.zig 01f897ac65d6edd2cd6ac76d97cb4d96619c25ffa zig/src/families/flashnext/cuda_head.zig fc6410e8b17d559a12d81fd4bee9ab4e65f6f8b0999796c68baba58efb8fe818 raw zig/src/families/flashnext/cuda_images.zig 5bd8a6c06cb69440522f4392794d9ef52b3760195b97d96f2316b1acac271bd4 raw zig/src/families/flashnext/cuda_images_test.zig df984d54d9de1bb732544785cddf82a124de2a7019041ed1f7483b2a4a2d7c99 raw -zig/src/families/flashnext/cuda_lanes.zig df7c17cb079b4a00ff16964c54bd3895bf997ac45942d051290bddc117e54b8f raw main 4f6017ced (marks inside the pass), fbca90714 (shared solo-chunk rule, #2391), 80e0fd72c (#2753: non-finite draft rows refused by name; finite rows unchanged), dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) +zig/src/families/flashnext/cuda_lanes.zig 1ebf522d252b795b15d58a15eb38d9d3227c341df80522ad8a89d5ff3cfbdac4 raw main 4f6017ced (marks inside the pass), fbca90714 (shared solo-chunk rule, #2391), 80e0fd72c (#2753: non-finite draft rows refused by name; finite rows unchanged), dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning), #1953: the ladder's values through snap_intermediate zig/src/families/flashnext/cuda_lanes_admission_test.zig 4b8052db793ba81f0e73e1b9bd77b1780bbc4fd550d90976d017da47e76faaa2 raw zig/src/families/flashnext/cuda_lanes_load.zig 241d091d4b5199325665d25ea5d4def81fd15660351a199773c596ca2025fa18 raw main dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) zig/src/families/flashnext/cuda_lanes_restore_test.zig f9947521a6d9c7d7380ca1b784a93a43b364c6fa5861cad15efd3538d2b8ceff raw @@ -166,4 +166,4 @@ zig/src/families/flashnext/cuda_rank_wire.zig ee1ea4202cb6b383e03a3c909d002fbdef zig/src/families/flashnext/cuda_round_graph.zig 7fa3a54e93c81de36cf3962d0a68eb72a7dd81b8869f47b6d9ed8b9a6cc325cd raw main dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) zig/src/families/flashnext/cuda_rows.zig 735fc8915036e47d09f533f4e5948e5fc4b61556172ebe2c6bfdaad3e0373216 raw main dfd9ac5ef4 (#2585: TF_FLASHNEXT_* read through platform/environment.zig, malformed text refused at startup; valid values keep their meaning) zig/src/families/flashnext/cuda_select_main.zig 5c012c2d59fa6dd573424fce7b6d8b162a06cb9f6103530a0fc0b03dbdb17bd3 raw -zig/src/families/flashnext/cuda_snap.zig 9a90f5761f063e08e568b4c86901b3d5196c6897da655cc5b864742cc2b663a6 raw main 4f6017ced (prompt cache marks kept inside the pass) +zig/src/families/flashnext/cuda_snap.zig bbc66ed7e1aed820b4b60c148b1bbe47b2a5c7e7240c7da7f0bd8504080a6b56 raw main 4f6017ced (prompt cache marks kept inside the pass), #1953: the prompt cache ladder's values From 1a368ecc561cb25560e2bf29dacb02c6f231cf94 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:41 +0000 Subject: [PATCH 067/186] changelog: #3006 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e9bb73e4fa3..be9ced2627c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel + + - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). From f986a778cd0857429a80b744e910c6db2c574445 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 06:38:22 -0800 Subject: [PATCH 068/186] qwen27 metal: retain the long-copy driver for a lone greedy stream The narrow tensor-unit default installs the existing greedy callback again. Sampled and shared rounds retain their lane path, and explicit lane overrides keep their meaning. Every copied token still goes through target verification. Refs #3013 --- docs/optimization-ledger.md | 2 ++ docs/recipes/qwen3.8-27b.md | 7 +++++-- zig/src/native/qwen27_served.zig | 2 +- zig/src/native/qwen27_window.zig | 16 +++++++++------- zig/tests/qwen27/reuse.zig | 2 +- 5 files changed, 18 insertions(+), 11 deletions(-) diff --git a/docs/optimization-ledger.md b/docs/optimization-ledger.md index 88e18546774..828c041aca0 100644 --- a/docs/optimization-ledger.md +++ b/docs/optimization-ledger.md @@ -375,6 +375,8 @@ below; these are estimates for planning a port, not results. 16. **DFlash and DFlash2 drafters.** An external block drafter reads the target's hidden taps and drafts a block (or a tree) a step. It needs a trained drafter for each target; see [training a drafter](recipes/training-a-drafter.md). + Qwen27 Metal retains its target-verified long-copy tail for an eligible lone greedy request at the M5 + 16-row default; sampled and shared requests keep the lane core ([#3013](https://github.com/ashhart/TensorFold/issues/3013)). 17. **DSpark drafters.** A drafter that ships with the release (DeepSeek-V4.1's `mtp.i` blocks, MiniMax-M3.1's, Kimi K3's) drafts from the target's taps with a confidence schedule. Applies wherever a DSpark drafter exists for the diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index 2a426004e4c..19d5212a0c4 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -33,8 +33,11 @@ A lone lane request may propose up to 15/31/63 tree nodes at 16/32/64 rows, incl by expected accepted tokens per measured millisecond in the common lane allocator, using its calibrated landing rates. The proposal keeps every candidate available until that decision. Copies compete for the same verification window; `TF_QWEN27_COPIES=0` disables them. The narrow tensor-unit default keeps copies off. -The masked drafter block is unchanged. Both device classes use the lane loop by default, including at 16 rows. -`TF_QWEN27_LANES=0` selects the legacy greedy shortcut; sampled requests retain the shared lane loop. +The masked drafter block is unchanged. At the tensor-unit 16-row default, an eligible lone greedy request +uses the long-copy driver; every copy is target-verified and arrivals yield back to the shared lane loop. +Sampled requests, shared rounds, wider windows and earlier chips retain the lane loop. +`TF_QWEN27_LANES=1` keeps every request on lanes; `0` requests the lone greedy driver on either device class. +The lane copy switches do not size the lone driver's tail; its existing cap is `TF_Q27_COPY_ROWS` (default 48). Startup names the selected lone loop and width beside the lane window and copy policy. `tf-qwen27-bench --model DIR --widths costs` measures rows 1 through 64, with GPU and wall time including keep. The wall price excludes the diagnostic's host logit hash and scan. diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index d3d47a9be4b..cf6c6ced00c 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -414,7 +414,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q h.host.memory = .{ .ctx = h, .read = readMemory }; h.host.explain = .{ .text = words }; h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off - // The narrow tensor window retains the serial loop; wide windows use the priced lane path. + // A lone greedy request keeps its long-copy tail on the narrow tensor window; arrivals yield to lanes. if (h.draft != null and serial) h.host.lone = .{ .ctx = h, .run = Host.lone }; // a reply that ends soon goes first beside a long prompt; a share someone set stays as they set it if (o.decode_share == null and std.c.getenv("TENSORFOLD_PREFILL_SHARE") == null) h.host.srpt = .withDefault(true); diff --git a/zig/src/native/qwen27_window.zig b/zig/src/native/qwen27_window.zig index 79bd7bf4198..b454161a4d9 100644 --- a/zig/src/native/qwen27_window.zig +++ b/zig/src/native/qwen27_window.zig @@ -54,9 +54,9 @@ fn flag(value: ?[]const u8, fallback: bool, comptime bad: anyerror) !bool { if (std.mem.eql(u8, v, "1")) return true; return bad; } -/// Both device classes use lanes by default; checked overrides retain the legacy greedy shortcut. -pub fn serial(_: bool, _: u32, value: ?[]const u8) !bool { - return !try flag(value, true, error.BadQwenLanes); +/// The narrow tensor window keeps the lone greedy copy tail; shared and sampled rounds still use lanes. +pub fn serial(tensor_units: bool, window: u32, value: ?[]const u8) !bool { + return !try flag(value, !(tensor_units and window == 16), error.BadQwenLanes); } pub fn serialPolicy(tensor_units: bool, value: ?[]const u8) !bool { return serial(tensor_units, 16, value); @@ -77,7 +77,7 @@ test "qwen27 window: unknown switches refuse by their startup variable names" { try std.testing.expectEqualStrings("TF_QWEN27_COPIES must be 0 or 1", problem(error.BadQwenCopies).?); try std.testing.expect(std.mem.startsWith(u8, problem(error.BadQwenWindow).?, "TF_QWEN27_WINDOW")); try std.testing.expect(std.mem.startsWith(u8, problem(error.BadQwenLanes).?, "TF_QWEN27_LANES")); - try std.testing.expect(!try serialPolicy(true, null)); + try std.testing.expect(try serialPolicy(true, null)); try std.testing.expect(!try serialPolicy(false, null)); try std.testing.expect(try serialPolicy(false, "0")); try std.testing.expect(!try serialPolicy(true, "1")); @@ -95,13 +95,15 @@ test "qwen27 window: startup reports the effective lone width beside the lane po var buf: [160]u8 = undefined; try std.testing.expectEqualStrings("lone loop=serial width=16; lane window=64, lane copies=on", try describe(&buf, true, 64, true, true)); try std.testing.expectEqualStrings("lone loop=lanes width=32; lane window=32, lane copies=off", try describe(&buf, false, 32, true, false)); - try std.testing.expectEqualStrings("lone loop=lanes width=16; lane window=16, lane copies=off", try describe(&buf, try serial(true, 16, null), 16, true, false)); + try std.testing.expectEqualStrings("lone loop=serial width=16; lane window=16, lane copies=off", try describe(&buf, try serial(true, 16, null), 16, true, false)); try std.testing.expect(std.mem.startsWith(u8, try describe(&buf, false, 64, false, false), "lone loop=plain width=1;")); } -test "qwen27 window: all widths use lanes by default and retain the checked serial override" { +test "qwen27 window: default lone copies return only on the narrow tensor window" { + try std.testing.expect(try serial(true, 16, null)); + try std.testing.expect(!try copyPolicy(true, 16, null)); for ([_]u32{ 16, 32, 64 }) |width| { - try std.testing.expect(!try serial(true, width, null)); + try std.testing.expectEqual(width == 16, try serial(true, width, null)); try std.testing.expect(!try serial(false, width, null)); try std.testing.expect(!try serial(true, width, "1")); try std.testing.expect(try serial(true, width, "0")); diff --git a/zig/tests/qwen27/reuse.zig b/zig/tests/qwen27/reuse.zig index 76ce44de5b2..691fe6a156e 100644 --- a/zig/tests/qwen27/reuse.zig +++ b/zig/tests/qwen27/reuse.zig @@ -91,7 +91,7 @@ const Run = struct { } }; -/// Drafted requests on the lane core: the host left its lone serial loop off (the default; TF_QWEN27_LANES=0 sets it). +/// Whether this attached host sends drafted requests through the lane core rather than the lone driver. fn lanesMode(h: *host.Host) bool { return h.host.lone == null; } From 900d8eae092c9df47f1c948e08bbbc74dd098f83 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:52:32 -0800 Subject: [PATCH 069/186] qwen27 metal: retain the lane head on confident lone copy rounds Use the lane selector sibling policy for the default narrow callback, and require a confident suffix before appending copy rows. Explicit serial overrides retain their original head, admission and submissions. Keep dependent matching, taps and preparation of the same keep at the existing verify tail, preserving target operators, budgets and EOS. Refs #3013 --- .../families/qwen27/dflash/execution_test.zig | 30 ++++++++ zig/src/families/qwen27/dflash/generation.zig | 73 +++++++++++++++++-- .../families/qwen27/dflash/runtime_model.zig | 9 ++- zig/src/native/draft_sources_metal_test.zig | 26 +++++++ zig/src/native/qwen27_lanes.zig | 1 + zig/src/native/qwen27_served.zig | 4 +- zig/src/native/qwen27_slots.zig | 8 +- zig/src/native/qwen27_window.zig | 14 ++++ 8 files changed, 153 insertions(+), 12 deletions(-) diff --git a/zig/src/families/qwen27/dflash/execution_test.zig b/zig/src/families/qwen27/dflash/execution_test.zig index 52f61000913..d21f3866a2e 100644 --- a/zig/src/families/qwen27/dflash/execution_test.zig +++ b/zig/src/families/qwen27/dflash/execution_test.zig @@ -397,6 +397,36 @@ test "qwen27 window: priced siblings match exhaustive expected accepted tokens p try std.testing.expectError(error.BadDraftCosts, selector.propose(a, &graph, lattice, 3, 2, .{ .nodes = 6, .costs = &invalid })); } +test "qwen27 window: lone callback retains the lane's wide sibling candidates" { + const a = std.testing.allocator; + var store = try fixtures.Store.init(a); + defer store.deinit(); + var graph = try store.graph(a); + defer graph.deinit(); + graph.config.topk = 16; + var fake: Fake = .{}; + var lattice = try Fake.lattice(&fake, 1, undefined, .{ .handle = 1, .rows = 2, .width = graph.config.rank, .stride = graph.config.rank, .dtype = .bf16 }, graph.config); + defer lattice.deinit(); + const model = @import("runtime_model.zig").Model; + const prices: [16]f64 = @splat(1); + const options = model.loneSelection(15); + var serial = try selector.propose(a, &graph, lattice, 3, 2, options); + defer serial.deinit(); + var lane = try selector.propose(a, &graph, lattice, 3, 2, model.selection(.{ .context = null, .pending = 3, .nodes = 15, .costs = &prices })); + defer lane.deinit(); + var legacy = try selector.propose(a, &graph, lattice, 3, 2, .{ .nodes = 15 }); + defer legacy.deinit(); + try std.testing.expectEqualSlices(u32, lane.tokens, serial.tokens); + try std.testing.expectEqualSlices(i32, lane.parents, serial.parents); + try std.testing.expectEqualSlices(f64, lane.scores, serial.scores); + var roots: usize = 0; + var old_roots: usize = 0; + for (serial.parents) |parent| roots += @intFromBool(parent == -1); + for (legacy.parents) |parent| old_roots += @intFromBool(parent == -1); + try std.testing.expectEqual(@as(usize, 15), roots); + try std.testing.expectEqual(@as(usize, 4), old_roots); +} + test "qwen27 window: common-core proposals retain candidates for calibrated window pricing" { const a = std.testing.allocator; var store = try fixtures.Store.init(a); diff --git a/zig/src/families/qwen27/dflash/generation.zig b/zig/src/families/qwen27/dflash/generation.zig index d060a32d1dc..710faf31098 100644 --- a/zig/src/families/qwen27/dflash/generation.zig +++ b/zig/src/families/qwen27/dflash/generation.zig @@ -136,12 +136,7 @@ pub const Generation = struct { const e = pick.compute(.concurrent); var ended = false; errdefer if (!ended) e.end(); - const logits = g.runner.model.frame.get(.logits); - const width: u32 = @intCast(g.runner.model.config.vocab); - const shape = try g.matcher.encode(e, .{ .buf = logits.buffer, .off = logits.offset }, ids, parents, width, budget, eos); - const result = @import("core").tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; - try g.commit.taps(e, result, .{ .buf = g.runner.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = g.runner.taps.width, .capacity = g.runner.taps.capacity, .planes = @intCast(g.runner.taps.count) }); - try g.commit.head(e, result, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = shape.rows, .width = width, .stride = width }); + const shape = try g.encodeAccept(e, ids, parents, budget, eos); e.end(); ended = true; pick.commit(); @@ -152,6 +147,7 @@ pub const Generation = struct { const k = keep.compute(.concurrent); var kept = false; errdefer if (!kept) k.end(); + const result = @import("core").tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; try g.runner.keepDevice(round, k, g.commit, result); k.end(); kept = true; @@ -160,6 +156,50 @@ pub const Generation = struct { if (gpuFailed(pick, "a draft round's picks")) return error.DraftKeepGpuFailure; const accepted = try g.matcher.read(shape); // committed without a wait: the drafter must not read or write what the keep writes until finishKeep + return g.publish(round, accepted, keep); + } + fn encodeAccept(g: *Generation, e: mtl.ComputeEncoder, ids: []const u32, parents: []const i32, budget: u32, eos: []const u32) !@import("core").tree_round.Match { + const logits = g.runner.model.frame.get(.logits); + const width: u32 = @intCast(g.runner.model.config.vocab); + const shape = try g.matcher.encode(e, .{ .buf = logits.buffer, .off = logits.offset }, ids, parents, width, budget, eos); + const result = @import("core").tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; + try g.commit.taps(e, result, .{ .buf = g.runner.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = g.runner.taps.width, .capacity = g.runner.taps.capacity, .planes = @intCast(g.runner.taps.count) }); + try g.commit.head(e, result, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = shape.rows, .width = width, .stride = width }); + return shape; + } + const Tail = struct { + g: *Generation, + round: *const Round, + ids: []const u32, + parents: []const i32, + budget: u32, + eos: []const u32, + shape: ?@import("core").tree_round.Match = null, + keep: ?mtl.CommandBuffer = null, + fn deinit(t: *Tail) void { + if (t.keep) |cb| cb.release(); + } + fn encode(ptr: *anyopaque, e: mtl.ComputeEncoder) !void { + const t: *Tail = @ptrCast(@alignCast(ptr)); + try t.g.runner.taps.complete(); + t.shape = try t.g.encodeAccept(e, t.ids, t.parents, t.budget, t.eos); + const cb = t.g.runner.model.queue.commandBuffer(); + const k = cb.compute(.concurrent); + var ended = false; + errdefer if (!ended) k.end(); + try t.g.runner.prepareDeviceKeep(t.round, k, t.g.commit, .{ .buf = t.g.matcher.buffers[4] }); + k.end(); + ended = true; + t.keep = cb.retain(); + } + }; + fn acceptReady(g: *Generation, round: *const Round, shape: @import("core").tree_round.Match, keep: mtl.CommandBuffer) !@import("core").tree_round.Decoded { + errdefer g.runner.failed = true; + try g.runner.checkDeviceKeep(round); + const accepted = try g.matcher.read(shape); + return g.publish(round, accepted, keep); + } + fn publish(g: *Generation, round: *const Round, accepted: @import("core").tree_round.Decoded, keep: mtl.CommandBuffer) !@import("core").tree_round.Decoded { keep.commit(); g.keep = keep.retain(); try g.runner.completeDeviceKeep(round, accepted.path); @@ -176,6 +216,13 @@ pub const Generation = struct { } /// `extend`: a copy from the context shares the tree's agreeing nodes and adds up to `extra` rows past them. pub fn runRound(g: *Generation, remaining: u32, nodes: usize, eos: []const u32, already_delivered: bool, extend: ?Extender, extra: u32) !Result { + return g.roundImpl(remaining, nodes, eos, already_delivered, extend, extra, false); + } + /// The lane's full candidate tree and dependent picks after the verify head's barrier, before its existing wait. + pub fn runRoundTail(g: *Generation, remaining: u32, nodes: usize, eos: []const u32, already_delivered: bool, extend: ?Extender, extra: u32) !Result { + return g.roundImpl(remaining, nodes, eos, already_delivered, extend, extra, true); + } + fn roundImpl(g: *Generation, remaining: u32, nodes: usize, eos: []const u32, already_delivered: bool, extend: ?Extender, extra: u32, fused: bool) !Result { if (g.runner.failed or g.runner.active != null or g.runner.offsets[0] == 0 or g.draft.backend.failed or g.draft.backend.frame.active or g.draft.backend.transaction != null or g.draft.committed_end != g.runner.offsets[0] or g.draft.committed_end == 0 or nodes == 0 or nodes > 15 or eos.len > 16) return error.BadDraftOptions; for (eos) |token| if (token >= g.runner.model.config.vocab) return error.BadDraftOptions; if (remaining == 0) return .{ .tokens = try g.allocator.alloc(u32, 0), .rounds = 0, .verified = 0, .matched = 0 }; @@ -189,7 +236,8 @@ pub const Generation = struct { const available = g.runner.capacity - g.runner.offsets[0]; if (available == 0) return error.ContextFull; const offered: u32 = @intCast(@min(nodes, available - 1)); - var tree = if (g.proposer) |proposer| try proposer.call(proposer.ptr, pending, offered) else try g.draft.propose(pending, offered); + const selection = if (fused) Draft.loneSelection(offered) else @import("selector.zig").Options{ .nodes = offered }; + var tree = if (g.proposer) |proposer| try proposer.call(proposer.ptr, pending, offered) else try g.draft.proposeOptions(pending, selection); defer tree.deinit(); const drafted = try tree.verifyParents(g.allocator); defer g.allocator.free(drafted); @@ -220,8 +268,17 @@ pub const Generation = struct { defer round.deinit(); errdefer g.finishKeep() catch {}; try g.finishKeep(); + var tail: Tail = .{ .g = g, .round = &round, .ids = ids, .parents = parents, .budget = remaining, .eos = eos }; + defer tail.deinit(); + if (fused) { + if (g.runner.tail != null) return error.BadVerifyTail; + g.runner.tail = .{ .ptr = &tail, .encode = Tail.encode }; + } + defer if (fused) { + g.runner.tail = null; + }; try g.runner.verifyHead(&round, .all); - const accepted = try g.accept(&round, ids, parents, remaining, eos); + const accepted = if (fused) try g.acceptReady(&round, tail.shape orelse return error.MissingRoundPicks, tail.keep orelse return error.MissingRoundKeep) else try g.accept(&round, ids, parents, remaining, eos); g.pending = accepted.pending; if (accepted.path.len == 0 or accepted.tokens.len == 0 or accepted.tokens.len > @min(@import("core").tree_round.max_rows, remaining)) return error.EmptyAcceptedDraft; var landed: u32 = 0; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 0ed3fe34aa8..7e3e0a99259 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -135,7 +135,14 @@ pub const Model = struct { }; /// The lane core prices complete candidate prefixes with its calibrated landing rates. pub fn selection(p: Proposal) selector.Options { - return .{ .nodes = p.nodes, .children = if (p.costs.len > 0) 16 else 4, .costs = if (p.defer_price) &.{} else p.costs, .overhead_ms = if (p.defer_price) 0 else p.overhead_ms }; + var options = if (p.costs.len > 0) loneSelection(p.nodes) else selector.Options{ .nodes = p.nodes }; + options.costs = if (p.defer_price) &.{} else p.costs; + options.overhead_ms = if (p.defer_price) 0 else p.overhead_ms; + return options; + } + /// A lone lane's full candidate tree, before the core prices its prefixes. + pub fn loneSelection(nodes: u32) selector.Options { + return .{ .nodes = nodes, .children = 16 }; } /// Trees for 2 to 8 streams from one frame: each stream's new taps into its own caches, then one block pass. pub fn proposeBatch(m: *Model, entries: []const Proposal, out: []selector.Tree) !void { diff --git a/zig/src/native/draft_sources_metal_test.zig b/zig/src/native/draft_sources_metal_test.zig index a4a2902610d..4359c02c9f8 100644 --- a/zig/src/native/draft_sources_metal_test.zig +++ b/zig/src/native/draft_sources_metal_test.zig @@ -4,6 +4,32 @@ const check = @import("draft_sources_check.zig").check; const present = @import("draft_sources_check.zig").present; const std = @import("std"); +test "qwen27 window: weak serial copies stay off while a confident continuation remains eligible" { + const a = std.testing.allocator; + const policy = @import("qwen27_window.zig"); + const minimum = policy.serialCopyMatch(true, 16, null); + try std.testing.expectEqual((tf.lanes.proposer.Params{}).confident_match, minimum); + for ([_]u32{ 16, 32, 64 }) |width| { + try std.testing.expectEqual(@as(i64, 5), policy.serialCopyMatch(false, width, null)); + try std.testing.expectEqual(@as(i64, 5), policy.serialCopyMatch(true, width, "0")); + try std.testing.expectEqual(@as(i64, 5), policy.serialCopyMatch(true, width, "1")); + } + var history: [65]u32 = undefined; + for (history[0..32], 0..) |*token, i| token.* = @intCast(i + 10); + history[32] = 999; + @memcpy(history[33..], history[0..32]); + var legacy = try tf.lanes.SuffixLookup.init(a, .{ .min_match = 5 }); + defer legacy.deinit(); + var strict = try tf.lanes.SuffixLookup.init(a, .{ .min_match = minimum }); + defer strict.deinit(); + // The former five-token rule takes this short repeated fragment; the default requires a long match. + try std.testing.expect((try legacy.propose(history[0..42], 16)).len > 0); + try std.testing.expectEqual(@as(usize, 0), (try strict.propose(history[0..42], 16)).len); + const copied = try strict.propose(history[0..57], 16); + try std.testing.expect(copied.len >= 8 and strict.last_confident); + try std.testing.expectEqualSlices(u32, history[24..32], copied[0..8]); +} + test "draft sources: Metal Nemotron follows the loaded MTP head" { var b: tf.nemotron.backend.Metal = undefined; b.costs = .{}; diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index 801dbf922d9..394590a14fa 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -56,6 +56,7 @@ pub const Metal = struct { owners: [max_streams]?*lanes.Stream = @splat(null), // the stream each runner slot holds contexts: [max_streams]?*Context = @splat(null), // slots past 0: their taps (a kept state carries them) and drafter state draft_slots: u32 = 1, // slots below this draft their own trees while no more streams are live (TF_Q27_DRAFT_SLOTS) + serial_copy_match: i64 = 5, // the lone callback's suffix rule, chosen once with its serving policy release_hook: ?Release = null, checkpoint_hook: ?lanes.stream.MarkHook = null, // unplanned prompt states only when the cache has room grow: ?Grow = null, diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index cf6c6ced00c..05c40ddba7f 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -344,7 +344,8 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q why.* = policy.problem(err) orelse why.*; return err; }; - const serial = policy.serial(units, window, if (platform.c.getenv("TF_QWEN27_LANES")) |v| std.mem.span(v) else null) catch |err| { + const lanes_value = if (platform.c.getenv("TF_QWEN27_LANES")) |v| std.mem.span(v) else null; + const serial = policy.serial(units, window, lanes_value) catch |err| { why.* = policy.problem(err) orelse why.*; return err; }; @@ -400,6 +401,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q errdefer if (h.draft) |d| d.deinit(); h.back = try lb.Metal.init(gpa, &h.runner, h.draft); errdefer h.back.deinit(); + h.back.serial_copy_match = policy.serialCopyMatch(units, window, lanes_value); h.back.chunk = h.runner.promptRows(@min(rows, row_policy.requested(model.device.generation(), model.config.moe(), value))); h.back.grow = .{ .ctx = h, .room = Host.roomFor, .added = Host.added }; h.back.vision = o.images; diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index ede4383d82d..2df612a9944 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -468,7 +468,7 @@ pub fn lone(b: *Metal, s: *lanes.Stream, hooks: api.LoneHooks) !bool { const g = &d.generation; const cap = try copyRows(); var extra: u32 = @min(cap, 16); // grows while copies land whole, shrinks when most miss - var copies = Copies{ .lookup = try lanes.SuffixLookup.init(b.gpa, .{ .min_match = 5 }) }; + var copies = Copies{ .lookup = try lanes.SuffixLookup.init(b.gpa, .{ .min_match = b.serial_copy_match }) }; defer copies.lookup.deinit(); _ = try s.commit(b.gpa, &.{try (q.session.Session{ .runner = b.runner }).greedy()}, &.{}); hooks.committed(hooks.ctx); @@ -484,7 +484,11 @@ pub fn lone(b: *Metal, s: *lanes.Stream, hooks: api.LoneHooks) !bool { const before = b.runner.gpu_seconds; mtl.pass_metrics.begin(b.pass_profile, "serial", b.window, mtl.clock.seconds); defer mtl.pass_metrics.end(); - const result = try g.runRound(room, @min(b.nodes, 15), s.eos, true, if (extra > 0) .{ .ptr = &copies, .call = Copies.call } else null, extra); + const extend: ?q.dflash.generation.Extender = if (extra > 0) .{ .ptr = &copies, .call = Copies.call } else null; + const result = if (b.serial_copy_match == 24 and @import("qwen27_window.zig").serialTail(b.tensor_units, b.window)) + try g.runRoundTail(room, @min(b.nodes, 15), s.eos, true, extend, extra) + else + try g.runRound(room, @min(b.nodes, 15), s.eos, true, extend, extra); defer b.gpa.free(result.tokens); charge(b, s, before); if (result.extended > 0) { diff --git a/zig/src/native/qwen27_window.zig b/zig/src/native/qwen27_window.zig index b454161a4d9..638ffd81425 100644 --- a/zig/src/native/qwen27_window.zig +++ b/zig/src/native/qwen27_window.zig @@ -61,6 +61,20 @@ pub fn serial(tensor_units: bool, window: u32, value: ?[]const u8) !bool { pub fn serialPolicy(tensor_units: bool, value: ?[]const u8) !bool { return serial(tensor_units, 16, value); } +/// The narrow tensor-unit serial path uses its existing verify tail for dependent picks. +pub fn serialTail(tensor_units: bool, window: u32) bool { + return tensor_units and window == 16; +} +/// The default narrow tensor stream requires a long suffix; explicit serial overrides retain their copy policy. +pub fn serialCopyMatch(tensor_units: bool, window: u32, lanes_value: ?[]const u8) i64 { + return if (tensor_units and window == 16 and lanes_value == null) 24 else 5; +} +test "qwen27 window: only the narrow tensor serial driver fuses its dependent picks" { + for ([_]u32{ 16, 32, 64 }) |width| { + try std.testing.expectEqual(width == 16, serialTail(true, width)); + try std.testing.expect(!serialTail(false, width)); + } +} pub fn problem(err: anyerror) ?[]const u8 { return switch (err) { error.BadQwenWindow => "TF_QWEN27_WINDOW must be 16, 32 or 64", From 20c5c61b574f5a0fdec7676c224e3ee7bbd35515 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:46 +0000 Subject: [PATCH 070/186] changelog: #3069 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index be9ced2627c..c3dbf05bdfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). + - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. From 8e66c92d02842f00026fc14a69fc16bfbbaacd20 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:41 +0000 Subject: [PATCH 071/186] changelog: #3006 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c3dbf05bdfd..affef2caf34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. + - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. From edd789366544de090f6f0fd8e5d443fee724db04 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:46 +0000 Subject: [PATCH 072/186] changelog: #3069 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index affef2caf34..f79c96a33cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). + - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). From 7e6aec2fd1b4c60a65f8d58adddb3f417105d790 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Thu, 12 Nov 2026 22:38:47 +0900 Subject: [PATCH 073/186] docs: the CUDA memory plan, what each setting decides and the file that reads it --- docs/api.md | 45 ++++++++++++++++++++++++++++++++++---------- docs/recipes/cuda.md | 1 + 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/docs/api.md b/docs/api.md index 65f790f38d4..d0e7fd2987a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -706,16 +706,41 @@ A startup capacity estimate is not a measured release capacity. On a unified-memory GPU (the DGX Spark's GB10), the CUDA server's allocations come out of the host's RAM but are not charged to a container's memory limit (`docker run --memory`, cgroup `memory.max`): the limit neither caps the -model's weights and cache nor keeps them from crowding other work on the machine. The server sizes its window from the -host's `MemAvailable` less a reserve (a tenth of RAM, at least 4 GiB); to leave room for other containers, start it -with a smaller `--context` or `--parallel`. `TENSORFOLD_MEMORY_RESERVE_GIB` replaces that reserve (at least 2 GiB) -when you know the machine's headroom: a larger one leaves more for other work, a smaller one more for KV caches. The -reserve also carries CUDA context, NCCL and workspace memory the estimate does not count, and exhausting a unified -GPU's memory can freeze the host, so lower it only with room to spare. - -On a discrete card (an RTX 3090, an A100) the server sizes its window from the card's own free memory, as the driver -reports it, less a reserve of a 32nd of the card and at least 1 GiB. `TENSORFOLD_MEMORY_RESERVE_GIB` sets that one from -0.5 GiB up. +model's weights and cache nor keeps them from crowding other work on the machine. The startup budget is the host's +`MemAvailable` less a reserve (a tenth of RAM, at least 4 GiB). On native CUDA a stream's window is `--context`, or +the family's default when it is unset; on the paged families the streams and pages then fill the budget up to the +streams' windows plus the cache's pages, so once those exceed it, the free memory after load is the reserve. To leave +room for other containers, start the server with a smaller `--context` or `--parallel`. +`TENSORFOLD_MEMORY_RESERVE_GIB` replaces that reserve (at least 2 GiB) when you know the machine's headroom: a larger +one leaves more for other work, a smaller one more for KV caches. The reserve also carries CUDA context, NCCL and +workspace memory the estimate does not count, and exhausting a unified GPU's memory can freeze the host, so lower it +only with room to spare. + +On a discrete card (an RTX 3090, an A100) the budget is the card's own free memory, as the driver reports it, less a +reserve of a 32nd of the card and at least 1 GiB, and the plan is the same. `TENSORFOLD_MEMORY_RESERVE_GIB` sets that +one from 0.5 GiB up. + +With pages, the plan takes the engine's own window first, then each other stream's fixed part, then the prompt cache: +a named `--prompt-cache-gib` comes before every window but the engine's own, and an unnamed one after the whole +windows the budget holds, half of what they leave but at least 16 GiB and at most all of it (#2987 shows what that +does to `TF_KV_SHARE`'s states). The cache then holds the kept states (on GLM with `TF_KV_SHARE`, a slot a state, the +rest of it made pages), and the rest of the budget becomes pages, up to the streams' windows plus the cache's pages. +The settings that decide each part: + +| Setting | Kind | Scope | Unset | What it decides | Read in | +|---|---|---|---|---|---| +| `TENSORFOLD_MEMORY_RESERVE_GIB` | environment | the host: every CUDA native family, each rank at load | a tenth of RAM, at least 4 GiB (unified); a 32nd of the card, at least 1 GiB (discrete) | what the budget leaves free | `zig/src/native/cuda_family_memory.zig`, `cuda_memory.zig` | +| `TENSORFOLD_CUDA_MEMORY_LIMIT_GB` | environment | the host: every CUDA native family, each rank at load | none | a cap on the budget, the weights included | `zig/src/native/cuda_family_memory.zig`, `cuda_memory.zig` | +| `--context` | flag | the launch | the family's default | a stream's window | `zig/src/server/cli.zig` | +| `--parallel` | flag | the launch | `auto`: up to 8 | the streams | `zig/src/server/cli.zig` | +| `--prompt-cache-gib` | flag | the launch | after the whole windows, as above | the prompt cache's budget | `zig/src/server/cli.zig`, planned in `zig/src/core/kv_pool.zig` | +| `TF_KV_SHARE` | environment, in the profile registry | the launch, on the families of the core page adapter | off: copies | prompts kept as pool pages, with at most this many states past them | `zig/src/core/kv_prompts.zig` | +| `TF_GLM_KV_HOST_GIB` | environment | GLM | one window and 16 pages | GLM's host KV pages with `TF_GLM_KV_HOST=1`, outside the device pool | `zig/src/families/glm/cuda_kv_host.zig` | +| the pool's pages | none | the core plan | — | what is left | `zig/src/core/kv_pool.zig` (`plan`) | + +For example, GLM-5.3-Flash at TP=3 on three GB10s with `--context 1048576 --parallel 8`, `TF_KV_SHARE=32` and no +`--prompt-cache-gib`, on #2765's head `f6d2a6d5e`: eight 1M windows need more pages than the budget holds, so the pool +takes the rest, and rank 0's free memory after load was 16.4 to 16.7 GiB at a 16 GiB reserve and 20.5 GiB at 20 GiB. ## Compaction diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index 9790f3262a3..5eb9845aa5f 100644 --- a/docs/recipes/cuda.md +++ b/docs/recipes/cuda.md @@ -358,6 +358,7 @@ Unified-memory GPUs share physical RAM with host buffers and file-backed model d host's available memory, reclaimable page cache included, less a floor of a tenth of RAM (at least 4 GiB) that `TENSORFOLD_MEMORY_RESERVE_GIB` can move, and considers mapped-table residency when sizing an automatic window. It accounts for stream count and retained caches where concurrency is enabled. +The native engine's memory plan and its settings are in the [API reference](../api.md#context-and-errors). Two-rank Flash Next, Nemotron and GLM requests finish on both ranks after a client disconnects, keeping the collective sequence aligned. MLX disk snapshots and cache-budget flags do not configure these CUDA From 0fb14b5931fc385743fecdacf6f7cbdd4bbe480d Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Fri, 13 Nov 2026 08:00:54 +0900 Subject: [PATCH 074/186] docs: --context's unset and 0 windows, the files that size the unnamed cache, and #3000 for the numbers --- docs/api.md | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/docs/api.md b/docs/api.md index d0e7fd2987a..98f10a003f4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -707,10 +707,11 @@ A startup capacity estimate is not a measured release capacity. On a unified-memory GPU (the DGX Spark's GB10), the CUDA server's allocations come out of the host's RAM but are not charged to a container's memory limit (`docker run --memory`, cgroup `memory.max`): the limit neither caps the model's weights and cache nor keeps them from crowding other work on the machine. The startup budget is the host's -`MemAvailable` less a reserve (a tenth of RAM, at least 4 GiB). On native CUDA a stream's window is `--context`, or -the family's default when it is unset; on the paged families the streams and pages then fill the budget up to the -streams' windows plus the cache's pages, so once those exceed it, the free memory after load is the reserve. To leave -room for other containers, start the server with a smaller `--context` or `--parallel`. +`MemAvailable` less a reserve (a tenth of RAM, at least 4 GiB). On native CUDA a stream's window is `--context`, the +model's own window at `--context 0`, or, when it is unset, the family's default capped at the model's own window; on +the paged families the streams and pages then fill the budget up to the streams' windows plus the cache's pages, so +once those exceed it, the free memory after load is the reserve. To leave room for other containers, start the server +with a smaller `--context` or `--parallel`. `TENSORFOLD_MEMORY_RESERVE_GIB` replaces that reserve (at least 2 GiB) when you know the machine's headroom: a larger one leaves more for other work, a smaller one more for KV caches. The reserve also carries CUDA context, NCCL and workspace memory the estimate does not count, and exhausting a unified GPU's memory can freeze the host, so lower it @@ -731,16 +732,14 @@ The settings that decide each part: |---|---|---|---|---|---| | `TENSORFOLD_MEMORY_RESERVE_GIB` | environment | the host: every CUDA native family, each rank at load | a tenth of RAM, at least 4 GiB (unified); a 32nd of the card, at least 1 GiB (discrete) | what the budget leaves free | `zig/src/native/cuda_family_memory.zig`, `cuda_memory.zig` | | `TENSORFOLD_CUDA_MEMORY_LIMIT_GB` | environment | the host: every CUDA native family, each rank at load | none | a cap on the budget, the weights included | `zig/src/native/cuda_family_memory.zig`, `cuda_memory.zig` | -| `--context` | flag | the launch | the family's default | a stream's window | `zig/src/server/cli.zig` | +| `--context` | flag | the launch | the family's default, capped at the model's own window | a stream's window; `0` takes the model's own | `zig/src/server/cli.zig`, `zig/src/native/cuda_window.zig` | | `--parallel` | flag | the launch | `auto`: up to 8 | the streams | `zig/src/server/cli.zig` | -| `--prompt-cache-gib` | flag | the launch | after the whole windows, as above | the prompt cache's budget | `zig/src/server/cli.zig`, planned in `zig/src/core/kv_pool.zig` | +| `--prompt-cache-gib` | flag | the launch | after the whole windows, as above | the prompt cache's budget | `zig/src/server/cli.zig`; unset, `zig/src/native/cuda_kv_pool.zig` (`budgetOf`) and `cuda_memory.zig` | | `TF_KV_SHARE` | environment, in the profile registry | the launch, on the families of the core page adapter | off: copies | prompts kept as pool pages, with at most this many states past them | `zig/src/core/kv_prompts.zig` | | `TF_GLM_KV_HOST_GIB` | environment | GLM | one window and 16 pages | GLM's host KV pages with `TF_GLM_KV_HOST=1`, outside the device pool | `zig/src/families/glm/cuda_kv_host.zig` | | the pool's pages | none | the core plan | — | what is left | `zig/src/core/kv_pool.zig` (`plan`) | -For example, GLM-5.3-Flash at TP=3 on three GB10s with `--context 1048576 --parallel 8`, `TF_KV_SHARE=32` and no -`--prompt-cache-gib`, on #2765's head `f6d2a6d5e`: eight 1M windows need more pages than the budget holds, so the pool -takes the rest, and rank 0's free memory after load was 16.4 to 16.7 GiB at a 16 GiB reserve and 20.5 GiB at 20 GiB. +Measured numbers for one launch are in #3000. ## Compaction From 29307aeed812f2a33be1cd27df749af85b54ed5c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 10:55:49 +0000 Subject: [PATCH 075/186] changelog: #3028 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index f79c96a33cb..7fdff0c914a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. + - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. From 4a68da22656864ccbe04a759902d15871d289056 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:38:01 +0000 Subject: [PATCH 076/186] changelog: each line once under Unreleased --- CHANGELOG.md | 8 -------- 1 file changed, 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fdff0c914a..7f35734060c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,14 +10,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. -- Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). - -- Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. - - - - - - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). From c34ff562ca6d8dfb9ac33fda39880c121f7c7c80 Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:39:36 +0000 Subject: [PATCH 077/186] changelog: one line for #2728, and #1878's line without its placeholder The train that landed #2728 left two lines for it under Unreleased (one had been under 1.2.6): keep the full wording with its number. #1878 also had a second line still reading "TBD-RECEIPT%"; its measured line (about 5%) stays. --- CHANGELOG.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f35734060c..716bd8c6a93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,8 +10,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Cortheon: a page link survives whitespace around `href` or inside its URL: `href ="…"` and a quote after a newline are kept, an unquoted link ends at a tab, CR or LF, and a tab or newline wrapped inside a link is not part of the address (#3006). Thanks to @ppardus. -- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, and concurrent streams equal solo ones. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. -- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it. Thanks to @gilby. +- Qwen3.8-Flash-Next on HIP: verify windows run their DeltaNet from the round's plan, so a window writes no state snapshot a row, and its keep replays the kept rows through the same scan. Drafted and undrafted replies are unchanged: every drafted reply equals the same request with drafts off, concurrent streams equal solo ones, and a resumed prompt equals a fresh one. Decode is 2–4% faster on Strix Halo (gfx1151) at one to four streams; shapes the plan kernels cannot take keep the per-row path, and no variable selects it (#2728). Thanks to @gilby. - Slide rollback: a fact graph is removed only when its own folders file names the folder being rolled back, and its folders file goes with it; a graph whose folders file can't be read is left in place and printed (#3073). - Chat templates: a template is held to a nesting depth at load, and its render to a step count, stack depth and scratch memory, so a crafted template can't hang or exhaust the server; real templates use a small fraction of each limit, render byte-identically, and render as fast as before (#3072). @@ -84,7 +83,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). -- HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about TBD-RECEIPT% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. - Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus. From 885918b651e97106845e1d482324538137dfecb2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:16:12 +0100 Subject: [PATCH 078/186] qwen27 metal: drop prompt_glue.metal and prompt_state.metal, whose six kernels nothing launches q27_prompt_product, _silu, _decay, _qk, q27_prompt_conv and q27_prompt_recur were reached only through prompt_gpu.Kernels and prompt_state_gpu.Kernels, which nothing builds with pipelines. The host records in prompt_glue.zig, prompt_gpu.zig, prompt_state.zig and prompt_state_gpu.zig, the build and runtime source lists, the runtime analysis line and the prompt.metal write in tests/qwen27/run.zig go with them. --- zig/build/qwen27.zig | 2 +- zig/kernels/metal/qwen27/prompt_glue.metal | 92 -------------------- zig/kernels/metal/qwen27/prompt_state.metal | 82 ----------------- zig/runtime_sources/qwen27.zig | 2 +- zig/src/families/qwen27/prompt_glue.zig | 50 ----------- zig/src/families/qwen27/prompt_gpu.zig | 71 --------------- zig/src/families/qwen27/prompt_state.zig | 60 ------------- zig/src/families/qwen27/prompt_state_gpu.zig | 38 -------- zig/src/families/qwen27/qwen27.zig | 6 -- zig/tests/qwen27/run.zig | 1 - zig/tests/qwen27/runtime.zig | 1 - 11 files changed, 2 insertions(+), 403 deletions(-) delete mode 100644 zig/kernels/metal/qwen27/prompt_glue.metal delete mode 100644 zig/kernels/metal/qwen27/prompt_state.metal delete mode 100644 zig/src/families/qwen27/prompt_glue.zig delete mode 100644 zig/src/families/qwen27/prompt_gpu.zig delete mode 100644 zig/src/families/qwen27/prompt_state.zig delete mode 100644 zig/src/families/qwen27/prompt_state_gpu.zig diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 88f939e7984..2faa0416cf6 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -16,7 +16,7 @@ pub const Shared = struct { }; /// The 27B's Metal sources, embedded by name (zig/kernels/metal/qwen27). -const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg", "prompt_glue", "prompt_state" }; +const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg" }; /// The shared modules, imported into `engine` (whose `core` is itself) for the server and the release binary. pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module) Shared { diff --git a/zig/kernels/metal/qwen27/prompt_glue.metal b/zig/kernels/metal/qwen27/prompt_glue.metal deleted file mode 100644 index 3ff271a61ab..00000000000 --- a/zig/kernels/metal/qwen27/prompt_glue.metal +++ /dev/null @@ -1,92 +0,0 @@ -// Prompt intermediates round where the declared BF16 graph rounds, rather than at the lane fusion's final store. -#include -using namespace metal; - -struct Q27PromptProduct { uint width, rows, mode, pad; }; -struct Q27PromptQk { uint rows, heads, width, stride, offset; float eps, scale; uint pad; }; -struct Q27PromptDecay { uint rows, heads, weight_flags, pad; }; - -inline float q27_prompt_probability(float value) { - const float magnitude = metal::abs(value); - const float denominator = precise::exp(magnitude) + 1.0f; - const float tail = 1.0f / denominator; - if (value < 0.0f) return tail; - return 1.0f - tail; -} - -inline float q27_prompt_softplus(float value) { - const float small = metal::exp(-metal::abs(value)); - const float shifted = 1.0f + small; - float logarithm = small; - if (shifted != 1.0f) { - const float difference = shifted - 1.0f; - const float ratio = metal::log(shifted) / difference; - logarithm = ratio * small; - } - return metal::max(value, 0.0f) + logarithm; -} - -kernel void q27_prompt_product(device const bfloat* gate [[buffer(0)]], - device const bfloat* factor [[buffer(1)]], - constant Q27PromptProduct& p [[buffer(2)]], - device bfloat* result [[buffer(3)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float x = float(gate[at]); - const float probability = q27_prompt_probability(x); - float activated; - if (p.mode == 0) activated = float(bfloat(x * float(bfloat(probability)))); - else if (p.mode == 1) activated = x * probability; - else activated = x / (1.0f + metal::exp(-x)); - result[at] = bfloat(activated * float(factor[at])); -} - -kernel void q27_prompt_silu(device const bfloat* input [[buffer(0)]], - constant Q27PromptProduct& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float value = float(input[at]); - const bfloat probability = bfloat(q27_prompt_probability(value)); - result[at] = bfloat(value * float(probability)); -} - -kernel void q27_prompt_decay(device const bfloat* a [[buffer(0)]], - device const bfloat* b [[buffer(1)]], - device const void* a_log [[buffer(2)]], - device const void* time_bias [[buffer(3)]], - constant Q27PromptDecay& p [[buffer(4)]], - device float* decay [[buffer(5)]], - device bfloat* beta [[buffer(6)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.heads) return; - const uint head = at % p.heads; - const bool wide_bias = (p.weight_flags & 4u) != 0; - const float bias = wide_bias ? ((device const float*)time_bias)[head] : float(((device const bfloat*)time_bias)[head]); - const float logarithm = (p.weight_flags & 2u) != 0 ? ((device const float*)a_log)[head] : float(((device const bfloat*)a_log)[head]); - const float sum = float(a[at]) + bias; - const float shifted = wide_bias ? sum : float(bfloat(sum)); - const float raw = q27_prompt_softplus(shifted); - const float smooth = wide_bias ? raw : float(bfloat(raw)); - const float speed = precise::exp(logarithm); - decay[at] = precise::exp(-speed * smooth); - beta[at] = bfloat(q27_prompt_probability(float(b[at]))); -} - -kernel void q27_prompt_qk(device const bfloat* activation [[buffer(0)]], - constant Q27PromptQk& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint2 group [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint begin = group.y * p.stride + p.offset + group.x * p.width + lane * 4; - float values[4]; - float square = 0.0f; - for (uint part = 0; part < 4; ++part) { - values[part] = float(activation[begin + part]); - square += values[part] * values[part]; - } - const float inverse = precise::rsqrt(simd_sum(square) / float(p.width) + p.eps); - const float scale = float(bfloat(p.scale)); - const uint output = (group.y * p.heads + group.x) * p.width + lane * 4; - for (uint part = 0; part < 4; ++part) result[output + part] = bfloat(scale * float(bfloat(values[part] * inverse))); -} diff --git a/zig/kernels/metal/qwen27/prompt_state.metal b/zig/kernels/metal/qwen27/prompt_state.metal deleted file mode 100644 index 1f8808136fc..00000000000 --- a/zig/kernels/metal/qwen27/prompt_state.metal +++ /dev/null @@ -1,82 +0,0 @@ -// Prompt state advances through the core tap order and one fixed per-position reduction topology. -#include -using namespace metal; - -struct Q27PromptConv { uint rows, channels, taps, streams; }; -struct Q27PromptRecur { uint rows, nk, nv, dv, streams, pad0, pad1, pad2; }; - -kernel void q27_prompt_conv(device const bfloat* projected [[buffer(0)]], - device const bfloat* history [[buffer(1)]], - device const bfloat* weight [[buffer(2)]], - constant Q27PromptConv& p [[buffer(3)]], - device bfloat* output [[buffer(4)]], - device bfloat* next_history [[buffer(5)]], - uint3 at [[thread_position_in_grid]]) { - if (at.x >= p.channels || at.y >= p.rows || at.z >= p.streams) return; - const uint keep = p.taps - 1; - const uint old_base = at.z * keep * p.channels; - const uint new_base = at.z * p.rows * p.channels; - float sum = 0.0f; - for (uint tap = 0; tap < p.taps; ++tap) { - const uint row = at.y + tap; - const float value = row < keep ? float(history[old_base + row * p.channels + at.x]) : float(projected[new_base + (row - keep) * p.channels + at.x]); - sum += value * float(weight[at.x * p.taps + tap]); - } - output[new_base + at.y * p.channels + at.x] = bfloat(sum); - if (at.y == 0) { - for (uint tail = 0; tail < keep; ++tail) { - const uint row = p.rows + tail; - next_history[old_base + tail * p.channels + at.x] = row < keep ? history[old_base + row * p.channels + at.x] : projected[new_base + (row - keep) * p.channels + at.x]; - } - } -} - -inline float q27_prompt_fold(thread float* partial) { - for (uint stride = 1; stride < 8; stride *= 2) { - for (uint at = 0; at < 8; at += 2 * stride) partial[at] = partial[at] + partial[at + stride]; - } - float total = partial[0]; - for (uint neighbor = 1; neighbor < 4; neighbor *= 2) total = total + simd_shuffle_xor(total, neighbor); - return total; -} - -kernel void q27_prompt_recur(device const bfloat* query [[buffer(0)]], - device const bfloat* key [[buffer(1)]], - device const bfloat* value [[buffer(2)]], - device const float* decay [[buffer(3)]], - device const bfloat* beta [[buffer(4)]], - device const float* committed [[buffer(5)]], - constant Q27PromptRecur& p [[buffer(6)]], - device bfloat* result [[buffer(7)]], - device float* next_state [[buffer(8)]], - uint3 at [[thread_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint stream = at.z / p.nv; - const uint head = at.z % p.nv; - const uint dimension = at.y * 8 + lane / 4; - const uint first = (lane % 4) * 32; - const uint key_head = head / (p.nv / p.nk); - const uint state_base = ((stream * p.nv + head) * p.dv + dimension) * 128 + first; - float state[32]; - for (uint channel = 0; channel < 32; ++channel) state[channel] = committed[state_base + channel]; - for (uint row = 0; row < p.rows; ++row) { - const uint key_base = ((stream * p.rows + row) * p.nk + key_head) * 128 + first; - const uint head_at = (stream * p.rows + row) * p.nv + head; - float partial[8]; - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] * decay[head_at]; - partial[channel / 4] += state[channel] * float(key[key_base + channel]); - } - const float remembered = q27_prompt_fold(partial); - const float correction = (float(value[head_at * p.dv + dimension]) - remembered) * float(beta[head_at]); - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] + float(key[key_base + channel]) * correction; - partial[channel / 4] += state[channel] * float(query[key_base + channel]); - } - const float output = q27_prompt_fold(partial); - if (lane % 4 == 0) result[head_at * p.dv + dimension] = bfloat(output); - } - for (uint channel = 0; channel < 32; ++channel) next_state[state_base + channel] = state[channel]; -} diff --git a/zig/runtime_sources/qwen27.zig b/zig/runtime_sources/qwen27.zig index 83327917957..1cb40bfe3d3 100644 --- a/zig/runtime_sources/qwen27.zig +++ b/zig/runtime_sources/qwen27.zig @@ -3,7 +3,7 @@ const std = @import("std"); const ks = @import("qwen_runtime_sources"); const Set = @import("set.zig").Set; pub fn add(s: *Set) !void { - inline for (.{ "glue", "copy", "moe", "prompt_glue", "prompt_state" }) |name| + inline for (.{ "glue", "copy", "moe" }) |name| try s.raw("qwen27", name, @field(ks, name)); try s.raw("qwen27", "moe-rows", @import("../src/families/qwen27/moe_gpu.zig").rows_source); const quant = @import("../src/families/qwen27/quant_gpu.zig"); diff --git a/zig/src/families/qwen27/prompt_glue.zig b/zig/src/families/qwen27/prompt_glue.zig deleted file mode 100644 index ce681249b01..00000000000 --- a/zig/src/families/qwen27/prompt_glue.zig +++ /dev/null @@ -1,50 +0,0 @@ -//! Prompt operator records keep scaled Q/K epsilon and staged activation modes separate from decode. -const std = @import("std"); -const Config = @import("config.zig").Config; - -pub const Mode = enum(u32) { staged_bf16 = 0, precise_f32 = 1, divided_control = 2 }; -pub const Product = extern struct { width: u32, rows: u32, mode: Mode, pad: u32 = 0 }; -pub const Qk = extern struct { rows: u32, heads: u32, width: u32, stride: u32, offset: u32, eps: f32, scale: f32, pad: u32 = 0 }; -pub const Decay = extern struct { rows: u32, heads: u32, weight_flags: u32, pad: u32 = 0 }; - -pub fn product(width: usize, rows: usize, mode: Mode) !Product { - if (width == 0 or width > 32768 or rows == 0 or rows > 262144 or @as(u64, width) * rows > 67108864) return error.BadPromptProduct; - return .{ .width = @intCast(width), .rows = @intCast(rows), .mode = mode }; -} - -pub fn qk(c: Config, rows: usize, query: bool) !Qk { - if (c.dk != 128 or c.dv != c.dk or c.k_heads == 0 or c.k_heads > 128 or rows == 0 or rows > 2048) return error.UnsupportedPromptQk; - const width: f64 = @floatFromInt(c.dk); - return .{ .rows = @intCast(rows), .heads = @intCast(c.k_heads), .width = @intCast(c.dk), .stride = @intCast(c.gdnQkvDim()), .offset = if (query) 0 else @intCast(c.k_heads * c.dk), .eps = @floatCast(@as(f64, 1e-6) / width), .scale = @floatCast(if (query) 1.0 / width else 1.0 / std.math.sqrt(width)) }; -} - -pub fn check(a: Qk) !void { - if (a.width != 128 or a.rows == 0 or a.rows > 2048 or a.heads == 0 or a.heads > 128 or a.offset > a.stride or a.heads * a.width > a.stride - a.offset or a.pad != 0 or !std.math.isFinite(a.eps) or a.eps <= 0 or !std.math.isFinite(a.scale) or a.scale <= 0) return error.UnsupportedPromptQk; -} - -pub fn checkDecay(a: Decay) !void { - if (a.rows == 0 or a.rows > 2048 or a.heads == 0 or a.heads > 128 or a.weight_flags > 15 or a.pad != 0) return error.BadPromptDecay; -} - -test "prompt QK epsilon and ABI differ explicitly from the decode epsilon" { - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Product)); - try std.testing.expectEqual(@as(usize, 32), @sizeOf(Qk)); - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Decay)); - try std.testing.expectEqual(@as(usize, 20), @offsetOf(Qk, "eps")); - const a = Qk{ .rows = 35, .heads = 16, .width = 128, .stride = 10240, .offset = 0, .eps = 1e-6 / 128.0, .scale = 1.0 / 128.0 }; - try check(a); - try std.testing.expect(a.eps < 1e-6); - var bad = a; - bad.width = 64; - try std.testing.expectError(error.UnsupportedPromptQk, check(bad)); - bad = a; - bad.offset = 10240; - try std.testing.expectError(error.UnsupportedPromptQk, check(bad)); -} - -test "unsupported prompt product extents refuse without changing activation modes" { - try std.testing.expectEqual(Mode.precise_f32, (try product(128, 48, .precise_f32)).mode); - try std.testing.expectEqual(@as(u32, 6144), (try product(128, 128 * 48, .precise_f32)).rows); - try std.testing.expectError(error.BadPromptProduct, product(0, 1, .staged_bf16)); - try std.testing.expectError(error.BadPromptProduct, product(17408, 65536, .staged_bf16)); -} diff --git a/zig/src/families/qwen27/prompt_gpu.zig b/zig/src/families/qwen27/prompt_gpu.zig deleted file mode 100644 index 675a16d5768..00000000000 --- a/zig/src/families/qwen27/prompt_gpu.zig +++ /dev/null @@ -1,71 +0,0 @@ -//! Caller-owned prompt pipelines compose RMS and activation stages without changing decoder pipelines. -const metal = @import("metal"); -const abi = @import("prompt_glue.zig"); -const glue = @import("glue.zig"); -const Ref = @import("projection.zig").Ref; - -pub const Kernels = struct { - head_norm: metal.Pipeline, - product: metal.Pipeline, - silu: metal.Pipeline, - qk: metal.Pipeline, - decay: metal.Pipeline, - - pub fn rms(k: Kernels, e: metal.ComputeEncoder, x: Ref, gain: Ref, output: Ref, a: glue.Head) !void { - if (a.activation != .bf16 or a.gain != .bf16) return error.UnsupportedPromptStorage; - const threads = try glue.headThreads(a); - if (threads > k.head_norm.maxThreads()) return error.UnsupportedPromptGroup; - e.setPipeline(k.head_norm); - e.setBuffer(x.buffer, x.offset, 0); - e.setBuffer(gain.buffer, gain.offset, 1); - e.setValue(a, 2); - e.setBuffer(output.buffer, output.offset, 3); - e.dispatchGroups(.{ .width = a.rows }, .{ .width = threads }); - e.barrier(); - } - - pub fn activation(k: Kernels, e: metal.ComputeEncoder, gate: Ref, factor: Ref, output: Ref, a: abi.Product) !void { - _ = try abi.product(a.width, a.rows, a.mode); - e.setPipeline(k.product); - bind(e, .{ gate, factor }); - e.setValue(a, 2); - e.setBuffer(output.buffer, output.offset, 3); - e.dispatchThreads(.{ .width = @as(usize, a.width) * a.rows }, .{ .width = 256 }); - e.barrier(); - } - - pub fn convolutionActivation(k: Kernels, e: metal.ComputeEncoder, input: Ref, output: Ref, a: abi.Product) !void { - _ = try abi.product(a.width, a.rows, a.mode); - e.setPipeline(k.silu); - e.setBuffer(input.buffer, input.offset, 0); - e.setValue(a, 1); - e.setBuffer(output.buffer, output.offset, 2); - e.dispatchThreads(.{ .width = @as(usize, a.width) * a.rows }, .{ .width = 256 }); - e.barrier(); - } - - pub fn normalizedQk(k: Kernels, e: metal.ComputeEncoder, activation_rows: Ref, output: Ref, a: abi.Qk) !void { - try abi.check(a); - e.setPipeline(k.qk); - e.setBuffer(activation_rows.buffer, activation_rows.offset, 0); - e.setValue(a, 1); - e.setBuffer(output.buffer, output.offset, 2); - e.dispatchGroups(.{ .width = a.heads, .height = a.rows }, .{ .width = 32 }); - e.barrier(); - } - - pub fn gdnDecay(k: Kernels, e: metal.ComputeEncoder, a_rows: Ref, b_rows: Ref, a_log: Ref, dt: Ref, g: Ref, beta: Ref, p: abi.Decay) !void { - try abi.checkDecay(p); - e.setPipeline(k.decay); - bind(e, .{ a_rows, b_rows, a_log, dt }); - e.setValue(p, 4); - e.setBuffer(g.buffer, g.offset, 5); - e.setBuffer(beta.buffer, beta.offset, 6); - e.dispatchThreads(.{ .width = @as(usize, p.rows) * p.heads }, .{ .width = 256 }); - e.barrier(); - } -}; - -fn bind(e: metal.ComputeEncoder, refs: anytype) void { - inline for (refs, 0..) |r, index| e.setBuffer(r.buffer, r.offset, index); -} diff --git a/zig/src/families/qwen27/prompt_state.zig b/zig/src/families/qwen27/prompt_state.zig deleted file mode 100644 index 6de771d694d..00000000000 --- a/zig/src/families/qwen27/prompt_state.zig +++ /dev/null @@ -1,60 +0,0 @@ -//! Prompt convolution and packed-state contracts refuse unsupported shape or storage before dispatch. -const std = @import("std"); - -pub const Conv = extern struct { rows: u32, channels: u32, taps: u32, streams: u32 }; -pub const Recur = extern struct { rows: u32, nk: u32, nv: u32, dv: u32, streams: u32, pad0: u32 = 0, pad1: u32 = 0, pad2: u32 = 0 }; - -pub fn checkConv(a: Conv) !void { - if (a.rows == 0 or a.rows > 2048 or a.channels == 0 or a.channels > 65536 or a.taps < 2 or a.taps > 8 or a.streams == 0 or a.streams > 128) return error.BadPromptConv; - if (@as(u64, a.rows) * a.channels * a.streams > std.math.maxInt(u32)) return error.BadPromptConv; -} - -pub fn checkRecur(a: Recur) !void { - if (a.rows == 0 or a.rows > 2048 or a.nk == 0 or a.nk > 128 or a.nv == 0 or a.nv > 128 or a.nv % a.nk != 0 or a.dv < 8 or a.dv > 256 or a.dv % 8 != 0 or a.streams == 0 or a.streams > 64 or a.pad0 != 0 or a.pad1 != 0 or a.pad2 != 0) return error.BadPromptRecurrence; -} - -pub fn canonical(partials: [32]f32) f32 { - @setFloatMode(.strict); - var values = partials; - var stride: usize = 1; - while (stride < values.len) : (stride *= 2) { - var at: usize = 0; - while (at < values.len) : (at += 2 * stride) values[at] += values[at + stride]; - } - return values[0]; -} - -pub fn packedFold(partials: [32]f32) f32 { - @setFloatMode(.strict); - var values = partials; - var roots: [4]f32 = undefined; - for (0..4) |lane| { - var stride: usize = 1; - while (stride < 8) : (stride *= 2) { - var at: usize = 0; - while (at < 8) : (at += 2 * stride) values[lane * 8 + at] += values[lane * 8 + at + stride]; - } - roots[lane] = values[lane * 8]; - } - return (roots[0] + roots[1]) + (roots[2] + roots[3]); -} - -test "packed local folds preserve the canonical FP32 tree without reassociation" { - for (0..64) |seed| { - var values: [32]f32 = undefined; - for (&values, 0..) |*v, index| { - const bits: u32 = @intCast((seed * 7919 + index * 65537) & 0x7fffff); - v.* = @bitCast((if ((seed + index) % 3 == 0) @as(u32, 1 << 31) else 0) | (120 + @as(u32, @intCast(index % 16))) << 23 | bits); - } - try std.testing.expectEqual(@as(u32, @bitCast(canonical(values))), @as(u32, @bitCast(packedFold(values)))); - } -} - -test "prompt recurrence and convolution ABI and row bounds are explicit" { - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Conv)); - try std.testing.expectEqual(@as(usize, 32), @sizeOf(Recur)); - try checkConv(.{ .rows = 35, .channels = 10240, .taps = 4, .streams = 1 }); - try checkRecur(.{ .rows = 128, .nk = 16, .nv = 48, .dv = 128, .streams = 1 }); - try std.testing.expectError(error.BadPromptRecurrence, checkRecur(.{ .rows = 128, .nk = 16, .nv = 47, .dv = 128, .streams = 1 })); - try std.testing.expectError(error.BadPromptConv, checkConv(.{ .rows = 2049, .channels = 10240, .taps = 4, .streams = 1 })); -} diff --git a/zig/src/families/qwen27/prompt_state_gpu.zig b/zig/src/families/qwen27/prompt_state_gpu.zig deleted file mode 100644 index 55fddbe661d..00000000000 --- a/zig/src/families/qwen27/prompt_state_gpu.zig +++ /dev/null @@ -1,38 +0,0 @@ -//! Prompt chain output and final FP32 state are encoded without touching decoder replay arithmetic. -const metal = @import("metal"); -const abi = @import("prompt_state.zig"); -const Ref = @import("projection.zig").Ref; - -pub const Kernels = struct { - conv: metal.Pipeline, - recur: metal.Pipeline, - - pub fn convolution(k: Kernels, e: metal.ComputeEncoder, projected: Ref, history: Ref, weight: Ref, output: Ref, next_history: Ref, p: abi.Conv) !void { - try abi.checkConv(p); - e.setPipeline(k.conv); - e.setBuffer(projected.buffer, projected.offset, 0); - e.setBuffer(history.buffer, history.offset, 1); - e.setBuffer(weight.buffer, weight.offset, 2); - e.setValue(p, 3); - e.setBuffer(output.buffer, output.offset, 4); - e.setBuffer(next_history.buffer, next_history.offset, 5); - e.dispatchThreads(.{ .width = p.channels, .height = p.rows, .depth = p.streams }, .{ .width = 256 }); - e.barrier(); - } - - pub fn recurrence(k: Kernels, e: metal.ComputeEncoder, q: Ref, key: Ref, value: Ref, decay: Ref, beta: Ref, committed: Ref, output: Ref, next_state: Ref, p: abi.Recur) !void { - try abi.checkRecur(p); - e.setPipeline(k.recur); - e.setBuffer(q.buffer, q.offset, 0); - e.setBuffer(key.buffer, key.offset, 1); - e.setBuffer(value.buffer, value.offset, 2); - e.setBuffer(decay.buffer, decay.offset, 3); - e.setBuffer(beta.buffer, beta.offset, 4); - e.setBuffer(committed.buffer, committed.offset, 5); - e.setValue(p, 6); - e.setBuffer(output.buffer, output.offset, 7); - e.setBuffer(next_state.buffer, next_state.offset, 8); - e.dispatchThreads(.{ .width = 32, .height = p.dv / 8, .depth = p.nv * p.streams }, .{ .width = 32, .height = 2 }); - e.barrier(); - } -}; diff --git a/zig/src/families/qwen27/qwen27.zig b/zig/src/families/qwen27/qwen27.zig index 6b1c366da2f..de17e8b1f23 100644 --- a/zig/src/families/qwen27/qwen27.zig +++ b/zig/src/families/qwen27/qwen27.zig @@ -30,11 +30,7 @@ pub const moe_gpu = @import("moe_gpu.zig"); pub const projection = @import("projection.zig"); pub const prompt_memory = @import("prompt_memory.zig"); pub const prompt_rows = @import("prompt_rows.zig"); -pub const prompt_glue = @import("prompt_glue.zig"); -pub const prompt_gpu = @import("prompt_gpu.zig"); -pub const prompt_state = @import("prompt_state.zig"); pub const prompt_marks = @import("prompt_marks.zig"); -pub const prompt_state_gpu = @import("prompt_state_gpu.zig"); pub const quant_gpu = @import("quant_gpu.zig"); pub const ane_mlp = @import("ane_mlp.zig"); pub const round_plan = @import("round_plan.zig"); @@ -66,8 +62,6 @@ test { _ = @import("prompt_tail.zig"); _ = round_plan; _ = attention; - _ = prompt_glue; - _ = prompt_state; _ = prompt_marks; _ = config; _ = affine; diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 77be2f0e586..6513b5fb2e7 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -34,7 +34,6 @@ fn emit(a: std.mem.Allocator, io: std.Io, dir: []const u8) !void { try write(a, io, dir, "glue.metal", sources.glue ++ sources.copy); try write(a, io, dir, "attention.metal", sources.glue ++ sources.attention_mpp ++ sources.attention_io); try write(a, io, dir, "gdn.metal", gdn_source.text); - try write(a, io, dir, "prompt.metal", sources.prompt_glue ++ sources.prompt_state); for ([_][]const u8{ "norm_input", "norm_xs", "mlp_xs", "xsum" }) |name| { const template = if (std.mem.eql(u8, name, "norm_input")) sources.norm_input else if (std.mem.eql(u8, name, "norm_xs")) sources.norm_xs else if (std.mem.eql(u8, name, "mlp_xs")) sources.mlp_xs else sources.xsum; const mark = std.mem.indexOf(u8, template, "Q27_SPECIALIZE") orelse return error.BadKernelSource; diff --git a/zig/tests/qwen27/runtime.zig b/zig/tests/qwen27/runtime.zig index 4faba9f4812..c561c8ad1f5 100644 --- a/zig/tests/qwen27/runtime.zig +++ b/zig/tests/qwen27/runtime.zig @@ -13,5 +13,4 @@ pub fn main() void { inline for (.{ &q.gdn_pool_gpu.Pool.init, &q.gdn_pool_gpu.Pool.forward, &q.gdn_pool_gpu.Pool.commit, &q.gdn_pool_gpu.Pool.clearSlot, &q.gdn_hook.Hook.forward, &q.gdn_hook.Hook.commit }) |function| std.mem.doNotOptimizeAway(function); inline for (.{ &q.attention.Attention.init, &q.attention.Attention.preprocess, &q.attention.Attention.encode, &q.attention.Attention.keep, &q.attention.Scratch.init, &q.attention.Scratch.upload, &q.attention.Binding.upload, &q.attention.Binding.uploadKeep }) |function| std.mem.doNotOptimizeAway(function); inline for (.{ &q.decode_round.Runner.init, &q.decode_round.Runner.verify, &q.decode_round.Runner.keep, &q.decode_round.Runner.reset }) |function| std.mem.doNotOptimizeAway(function); - inline for (.{ &q.prompt_gpu.Kernels.rms, &q.prompt_gpu.Kernels.activation, &q.prompt_gpu.Kernels.normalizedQk, &q.prompt_gpu.Kernels.gdnDecay, &q.prompt_state_gpu.Kernels.convolution, &q.prompt_state_gpu.Kernels.recurrence }) |function| std.mem.doNotOptimizeAway(function); } From 4f12fa0805e547770d98b93ddd44a5b52f3b0d38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:17:05 +0100 Subject: [PATCH 079/186] core metal: drop tf_lane_rope and Attention.rotate, which only lanes_see calls No family calls rotate(); the shared attention's runtime callers (Gemma 4 and Qwen27 DFlash) rope their rows elsewhere. Shape loses rotary_dim and rope, which only the kernel read, and lanes_see feeds raw K and Q straight to append and attention. --- zig/kernels/metal/core/shared_attention.metal | 28 ------ zig/src/core/shared_attention.zig | 43 ++------ zig/src/families/gemma4/dflash.zig | 2 +- .../qwen27/dflash/runtime_backend.zig | 2 +- zig/tests/lanes_see.zig | 98 ++++--------------- 5 files changed, 30 insertions(+), 143 deletions(-) diff --git a/zig/kernels/metal/core/shared_attention.metal b/zig/kernels/metal/core/shared_attention.metal index 126aef9685c..11ae5ce8910 100644 --- a/zig/kernels/metal/core/shared_attention.metal +++ b/zig/kernels/metal/core/shared_attention.metal @@ -149,34 +149,6 @@ inline void tf_context_row(const device bfloat* Q, const device bfloat* K, const tf_context_row(Q, K, V, (const device uchar*)nullptr, O, 0, head, args.keys, 0, as_type(args.scale_bits), lane); } -// Position is supplied per row. Caller-supplied inverse frequencies preserve the model's RoPE convention. -[[kernel]] void tf_lane_rope(const device bfloat* X [[buffer(0)]], const device uint* positions [[buffer(1)]], - const device float* frequencies [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], - device bfloat* Y [[buffer(4)]], uint index [[thread_position_in_grid]]) { - if ((size_t)index >= (size_t)args.rows * args.heads * TF_D) return; - const uint d = index % TF_D; -#if TF_RD > 0 - if (d >= TF_RD) { Y[index] = X[index]; return; } -#if TF_INTERLEAVED - const uint pair = d / 2; - const uint first = pair * 2, second = first + 1; - const bool high = (d & 1) != 0; -#else - const uint pair = d % (TF_RD / 2); - const uint first = pair, second = first + TF_RD / 2; - const bool high = d >= TF_RD / 2; -#endif - const uint row = index / (args.heads * TF_D); - const size_t base = index - d; - const float angle = float(positions[row]) * frequencies[pair]; - const float c = precise::cos(angle), s = precise::sin(angle); - const float x = float(X[base + first]), y = float(X[base + second]); - Y[index] = bfloat(high ? fma(x, s, y*c) : fma(x, c, -y*s)); -#else - Y[index] = X[index]; -#endif -} - // Only newly produced rows are written. Prompt and previous-round cache slots are never copied or rewritten. [[kernel]] void tf_lane_append_kv(const device bfloat* K [[buffer(0)]], const device bfloat* V [[buffer(1)]], constant uint* slots [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], diff --git a/zig/src/core/shared_attention.zig b/zig/src/core/shared_attention.zig index f5d537300e4..f508c61475f 100644 --- a/zig/src/core/shared_attention.zig +++ b/zig/src/core/shared_attention.zig @@ -1,23 +1,20 @@ -//! Shared-context attention, per-row RoPE and cache append; byte masks alone decide what each row sees. +//! Shared-context attention and cache append; byte masks alone decide what each row sees. const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); pub const Output = enum { f32, bf16 }; -pub const Rope = enum { split_half, interleaved }; pub const Shape = struct { query_heads: usize, kv_heads: usize, head_dim: usize, - rotary_dim: usize, - rope: Rope = .split_half, capacity: usize, output: Output = .f32, pub fn validate(s: Shape) !void { if (s.query_heads == 0 or s.kv_heads == 0 or s.query_heads % s.kv_heads != 0 or s.query_heads > 128 or s.head_dim == 0 or s.head_dim > 256 or s.head_dim % 32 != 0 or - s.rotary_dim > s.head_dim or s.rotary_dim % 2 != 0 or s.capacity == 0 or s.capacity > std.math.maxInt(u32)) return error.InvalidAttentionShape; + s.capacity == 0 or s.capacity > std.math.maxInt(u32)) return error.InvalidAttentionShape; _ = try std.math.mul(usize, s.capacity, try std.math.mul(usize, s.kv_heads, s.head_dim)); } }; @@ -30,8 +27,8 @@ const Args = extern struct { rows: u32, keys: u32 = 0, mask_stride: u32 = 0, hea pub fn source(a: std.mem.Allocator, s: Shape) ![]u8 { try s.validate(); - return std.fmt.allocPrint(a, "#define TF_QH {d}\n#define TF_KVH {d}\n#define TF_D {d}\n#define TF_RD {d}\n#define TF_INTERLEAVED {d}\n#define TF_OUT {s}\n{s}", .{ - s.query_heads, s.kv_heads, s.head_dim, s.rotary_dim, @intFromBool(s.rope == .interleaved), if (s.output == .f32) "float" else "bfloat", sources.core_shared_attention, + return std.fmt.allocPrint(a, "#define TF_QH {d}\n#define TF_KVH {d}\n#define TF_D {d}\n#define TF_OUT {s}\n{s}", .{ + s.query_heads, s.kv_heads, s.head_dim, if (s.output == .f32) "float" else "bfloat", sources.core_shared_attention, }); } @@ -39,7 +36,6 @@ pub const Attention = struct { shape: Shape, attend: mtl.Pipeline, decode: mtl.Pipeline, - rope: mtl.Pipeline, append: mtl.Pipeline, split: mtl.Pipeline, merge: mtl.Pipeline, @@ -55,8 +51,6 @@ pub const Attention = struct { errdefer attend.deinit(); const decode = try mtl.Pipeline.init(device, lib, "tf_core_decode_attention", false); errdefer decode.deinit(); - const rope = try mtl.Pipeline.init(device, lib, "tf_lane_rope", false); - errdefer rope.deinit(); const append = try mtl.Pipeline.init(device, lib, "tf_lane_append_kv", false); errdefer append.deinit(); const split = try mtl.Pipeline.init(device, lib, "tf_shared_lane_attention_split", false); @@ -65,12 +59,11 @@ pub const Attention = struct { errdefer merge.deinit(); if (attend.simdWidth() != 32 or attend.maxThreads() < 32 or split.simdWidth() != 32 or merge.simdWidth() != 32) return error.UnsupportedAttentionPipeline; const part = try device.buffer(16 * s.query_heads * chunksOf(s.capacity) * (s.head_dim + 2) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); - return .{ .shape = s, .attend = attend, .decode = decode, .rope = rope, .append = append, .split = split, .merge = merge, .part = part }; + return .{ .shape = s, .attend = attend, .decode = decode, .append = append, .split = split, .merge = merge, .part = part }; } pub fn deinit(p: Attention) void { p.attend.deinit(); p.decode.deinit(); - p.rope.deinit(); p.append.deinit(); p.split.deinit(); p.merge.deinit(); @@ -145,25 +138,6 @@ pub const Attention = struct { e.dispatchGroups(mtl.Size.of(p.shape.query_heads, 1, 1), mtl.Size.of(32, 1, 1)); } - /// positions u32[rows], frequencies f32[rotary_dim/2], storage bf16[rows,heads,D]; prompt chunks too. - pub fn rotate(p: Attention, e: mtl.ComputeEncoder, rows: usize, heads: usize, x: Ref, positions: Ref, frequencies: Ref, y: Ref) !void { - if (rows == 0 or rows > p.shape.capacity or (heads != p.shape.query_heads and heads != p.shape.kv_heads)) return error.InvalidRopeView; - const elements = try std.math.mul(usize, rows, heads * p.shape.head_dim); - if (elements > std.math.maxInt(u32)) return error.InvalidRopeView; - try check(x, elements * 2, 2); - try check(y, elements * 2, 2); - try check(positions, rows * 4, 4); - try check(frequencies, @max(p.shape.rotary_dim / 2 * 4, 4), 4); - try disjoint(x, y); - try disjoint(y, positions); - try disjoint(y, frequencies); - e.setPipeline(p.rope); - for ([_]Ref{ x, positions, frequencies }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); - e.setValue(Args{ .rows = @intCast(rows), .heads = @intCast(heads) }, 3); - e.setBuffer(y.buf, y.off, 4); - e.dispatchThreads(mtl.Size.of(elements, 1, 1), mtl.Size.of(128, 1, 1)); - } - /// Prepared K and unmodified V go directly to unique physical slots; caller orders subsequent attention. pub fn appendKv(p: Attention, e: mtl.ComputeEncoder, slots: []const u32, k: Ref, v: Ref, cache_k: Ref, cache_v: Ref) !void { if (slots.len == 0 or slots.len > 16) return error.InvalidAppendView; @@ -202,19 +176,16 @@ fn disjoint(a: Ref, b: Ref) !void { if (a.buf.id == b.buf.id) return error.AliasedAttentionBuffer; } -test "attention shapes keep grouped-query and rotary dimensions explicit" { +test "attention shapes keep grouped-query dimensions explicit" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); - const good = Shape{ .query_heads = 4, .kv_heads = 2, .head_dim = 64, .rotary_dim = 32, .capacity = 128 }; + const good = Shape{ .query_heads = 4, .kv_heads = 2, .head_dim = 64, .capacity = 128 }; const text = try source(arena.allocator(), good); try std.testing.expect(std.mem.indexOf(u8, text, "#define TF_QH 4\n#define TF_KVH 2\n") != null); var bad = good; bad.query_heads = 3; try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); bad = good; - bad.rotary_dim = 33; - try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); - bad = good; bad.head_dim = 31; try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); } diff --git a/zig/src/families/gemma4/dflash.zig b/zig/src/families/gemma4/dflash.zig index addeedc41fd..7774b78e150 100644 --- a/zig/src/families/gemma4/dflash.zig +++ b/zig/src/families/gemma4/dflash.zig @@ -109,7 +109,7 @@ pub const Drafter = struct { d.ops = try ops.Ops.init(target.device); errdefer d.ops.deinit(); const keys_max = @max(options.capacity + block_rows, c.window - 1 + block_rows); - d.attention = try attention.Attention.init(gpa, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = keys_max, .output = .bf16 }); + d.attention = try attention.Attention.init(gpa, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .capacity = keys_max, .output = .bf16 }); errdefer d.attention.deinit(); d.picks = try round.Ops.init(target.device); errdefer d.picks.deinit(); diff --git a/zig/src/families/qwen27/dflash/runtime_backend.zig b/zig/src/families/qwen27/dflash/runtime_backend.zig index 51c2a4d3dbc..0fee0ed3906 100644 --- a/zig/src/families/qwen27/dflash/runtime_backend.zig +++ b/zig/src/families/qwen27/dflash/runtime_backend.zig @@ -77,7 +77,7 @@ pub const Backend = struct { try c.check(); const primitives = try core.draft_ops.Ops.init(target.device); errdefer primitives.deinit(); - const attention = try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = c.window + 16, .output = .bf16 }); + const attention = try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .capacity = c.window + 16, .output = .bf16 }); errdefer attention.deinit(); const topk = try core.bf16_topk_gpu.Ops.init(target.device); errdefer topk.deinit(); diff --git a/zig/tests/lanes_see.zig b/zig/tests/lanes_see.zig index 60705e85782..27301d7e572 100644 --- a/zig/tests/lanes_see.zig +++ b/zig/tests/lanes_see.zig @@ -1,4 +1,4 @@ -//! Shared context, data masks and lane-local rotary positions on small synthetic buffers only. +//! Shared context and data masks on small synthetic buffers only. const std = @import("std"); const mtl = @import("metal"); const shared = @import("shared_attention"); @@ -41,30 +41,6 @@ fn guards(b: mtl.Buffer) !void { for (b.contents()[b.length() - pad .. b.length()]) |x| if (x != 0xa5) return error.GuardOverwrite; } -fn checkRope(s: shared.Shape, input: []const u16, output: []const u16, positions: []const u32, freq: []const f32, heads: usize) !void { - for (0..positions.len) |r| for (0..heads) |h| for (0..s.head_dim) |d| { - const at = (r * heads + h) * s.head_dim + d; - if (d >= s.rotary_dim) { - if (input[at] != output[at]) return error.RopeTailChanged; - continue; - } - const pair = if (s.rope == .interleaved) d / 2 else d % (s.rotary_dim / 2); - const first = if (s.rope == .interleaved) pair * 2 else pair; - const second = if (s.rope == .interleaved) first + 1 else first + s.rotary_dim / 2; - const base = at - d; - const x: f64 = fromBf(input[base + first]); - const y: f64 = fromBf(input[base + second]); - const angle: f32 = @as(f32, @floatFromInt(positions[r])) * freq[pair]; - const c = @cos(@as(f64, angle)); - const sn = @sin(@as(f64, angle)); - const high = if (s.rope == .interleaved) d % 2 != 0 else d >= s.rotary_dim / 2; - const want = if (high) x * sn + y * c else x * c - y * sn; - const got: f64 = fromBf(output[at]); - // One final bf16 cast; allow its rounding and the fp32 trig/FMA rounding only. - if (@abs(got - want) > 0.004 * @abs(want) + 0.00001) return error.RopeMismatch; - }; -} - fn cpuAttention(s: shared.Shape, q: []const u16, k: []const u16, v: []const u16, mask: []const u8, keys: usize, row: usize, head: usize, scale: f32, out: []f64) !void { const scores = try a.alloc(f64, keys); defer a.free(scores); @@ -96,12 +72,12 @@ fn outputValue(b: mtl.Buffer, index: usize, output: shared.Output) f32 { return if (output == .f32) values(b, f32, index + 1)[index] else fromBf(values(b, u16, index + 1)[index]); } -const Counts = struct { ordinary: usize = 0, bf16_values: usize = 0, max_error: f64 = 0, max_fraction: f64 = 0, cases: usize = 0, rows: usize = 0, cpu_values: usize = 0, rope_values: usize = 0, prompt_bytes: usize = 0 }; +const Counts = struct { ordinary: usize = 0, bf16_values: usize = 0, max_error: f64 = 0, max_fraction: f64 = 0, cases: usize = 0, rows: usize = 0, cpu_values: usize = 0, prompt_bytes: usize = 0 }; -fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]usize, mode: shared.Rope, output: shared.Output, counts: *Counts) !void { +fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [3]usize, output: shared.Output, counts: *Counts) !void { const pool = mtl.objc.Pool.push(); defer pool.pop(); - const s = shared.Shape{ .query_heads = dims[0], .kv_heads = dims[1], .head_dim = dims[2], .rotary_dim = dims[3], .rope = mode, .capacity = prompt + n * 4, .output = output }; + const s = shared.Shape{ .query_heads = dims[0], .kv_heads = dims[1], .head_dim = dims[2], .capacity = prompt + n * 4, .output = output }; const p = try shared.Attention.init(a, device, s); defer p.deinit(); var c = try kv.Cache.init(a, .{ .lanes = n, .prompt_capacity = prompt, .round_capacity = 4, .kv_heads = s.kv_heads, .head_dim = s.head_dim }); @@ -117,37 +93,20 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u const max_rows = @max(n, prompt); const raw_k = try buffer(device, max_rows * width * 2); defer raw_k.deinit(); - const prepared_k = try buffer(device, max_rows * width * 2); - defer prepared_k.deinit(); const raw_v = try buffer(device, max_rows * width * 2); defer raw_v.deinit(); - const positions = try buffer(device, max_rows * 4); - defer positions.deinit(); - const frequencies = try buffer(device, @max(s.rotary_dim / 2 * 4, 4)); - defer frequencies.deinit(); - const freq = values(frequencies, f32, @max(s.rotary_dim / 2, 1)); - for (freq, 0..) |*v, d| v.* = 1.0 / @as(f32, @floatFromInt(1 + d * 7)); if (prompt > 0) { const raw = values(raw_k, u16, prompt * width); const val = values(raw_v, u16, prompt * width); - const pos = values(positions, u32, prompt); for (raw, 0..) |*v, i| v.* = data(i, 1); for (val, 0..) |*v, i| v.* = data(i, 2); - for (pos, 0..) |*v, i| v.* = @intCast(i); - const cb = queue.commandBuffer(); - const e = cb.compute(.serial); - try p.rotate(e, prompt, s.kv_heads, ref(raw_k, 0), ref(positions, 0), ref(frequencies, 0), ref(prepared_k, 0)); - e.end(); - try finish(cb); - try checkRope(s, raw, values(prepared_k, u16, prompt * width), pos, freq, s.kv_heads); - counts.rope_values += prompt * width; - @memcpy(values(ck, u16, prompt * width), values(prepared_k, u16, prompt * width)); + @memcpy(values(ck, u16, prompt * width), raw); @memcpy(values(cv, u16, prompt * width), val); const ids = try a.alloc(u32, prompt); defer a.free(ids); for (ids, 0..) |*v, i| v.* = @intCast(10 + i); - try c.appendPrompt(ids, values(prepared_k, u16, prompt * width), val); + try c.appendPrompt(ids, raw, val); } const prompt_k = try a.dupe(u16, values(ck, u16, prompt * width)); defer a.free(prompt_k); @@ -160,11 +119,9 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u var slots: [16]u32 = undefined; var owners: [16]usize = undefined; var used: usize = 0; - const pos = values(positions, u32, n); for (0..n) |lane| if (active & (@as(u16, 1) << @intCast(lane)) != 0) { owners[used] = lane; slots[used] = @intCast(try c.slot(ticket, lane)); - pos[used] = ticket.positions[lane]; for (0..width) |d| { values(raw_k, u16, n * width)[used * width + d] = data(d, 3 + round * 17 + lane); values(raw_v, u16, n * width)[used * width + d] = data(d, 5 + round * 19 + lane); @@ -173,26 +130,20 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u }; const cb = queue.commandBuffer(); const e = cb.compute(.serial); - try p.rotate(e, used, s.kv_heads, ref(raw_k, 0), ref(positions, 0), ref(frequencies, 0), ref(prepared_k, 0)); - e.barrier(); - try p.appendKv(e, slots[0..used], ref(prepared_k, 0), ref(raw_v, 0), ref(ck, 0), ref(cv, 0)); + try p.appendKv(e, slots[0..used], ref(raw_k, 0), ref(raw_v, 0), ref(ck, 0), ref(cv, 0)); e.end(); try finish(cb); - try checkRope(s, values(raw_k, u16, used * width), values(prepared_k, u16, used * width), pos[0..used], freq, s.kv_heads); - counts.rope_values += used * width; - for (0..used) |i| try c.put(ticket, owners[i], @intCast(100 + round * 16 + owners[i]), values(prepared_k, u16, used * width)[i * width ..][0..width], values(raw_v, u16, used * width)[i * width ..][0..width]); + for (0..used) |i| try c.put(ticket, owners[i], @intCast(100 + round * 16 + owners[i]), values(raw_k, u16, used * width)[i * width ..][0..width], values(raw_v, u16, used * width)[i * width ..][0..width]); try std.testing.expectEqualSlices(u16, prompt_k, values(ck, u16, prompt * width)); try std.testing.expectEqualSlices(u16, prompt_v, values(cv, u16, prompt * width)); for (0..used) |i| { const off = @as(usize, slots[i]) * width; - try std.testing.expectEqualSlices(u16, values(prepared_k, u16, used * width)[i * width ..][0..width], values(ck, u16, s.capacity * width)[off..][0..width]); + try std.testing.expectEqualSlices(u16, values(raw_k, u16, used * width)[i * width ..][0..width], values(ck, u16, s.capacity * width)[off..][0..width]); try std.testing.expectEqualSlices(u16, values(raw_v, u16, used * width)[i * width ..][0..width], values(cv, u16, s.capacity * width)[off..][0..width]); } if (round < 2) try c.commit(ticket) else last = ticket; } const view = try c.current(last); - const raw_q = try buffer(device, n * qwidth * 2); - defer raw_q.deinit(); const q = try buffer(device, n * qwidth * 2); defer q.deinit(); const masks = try buffer(device, n * view.keys); @@ -202,10 +153,7 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u defer out.deinit(); const serial = try buffer(device, n * qwidth * elem); defer serial.deinit(); - const raw = values(raw_q, u16, n * qwidth); - for (raw, 0..) |*v, i| v.* = data(i, 73); - const qp = values(positions, u32, n); - @memcpy(qp, last.positions[0..n]); + for (values(q, u16, n * qwidth), 0..) |*v, i| v.* = data(i, 73); const mask = values(masks, u8, n * view.keys); try c.visibility(view, n, mask); if (n > 1) { @@ -219,15 +167,11 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(s.head_dim))); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - try p.rotate(e, n, s.query_heads, ref(raw_q, 0), ref(positions, 0), ref(frequencies, 0), ref(q, 0)); - e.barrier(); try p.encode(e, n, view.keys, view.keys, scale, ref(q, 0), ref(ck, 0), ref(cv, 0), ref(masks, 0), ref(out, 0)); for (0..n) |row| try p.encode(e, 1, view.keys, view.keys, scale, ref(q, row * qwidth * 2), ref(ck, 0), ref(cv, 0), ref(masks, row * view.keys), ref(serial, row * qwidth * elem)); if (n == 1) try p.ordinary(e, view.keys, scale, ref(q, 0), ref(ck, 0), ref(cv, 0), ref(serial, 0)); e.end(); try finish(cb); - try checkRope(s, raw, values(q, u16, n * qwidth), qp, freq, s.query_heads); - counts.rope_values += n * qwidth; try std.testing.expectEqualSlices(u8, out.contents()[pad..][0 .. n * qwidth * elem], serial.contents()[pad..][0 .. n * qwidth * elem]); const expected = try a.alloc(f64, s.head_dim); defer a.free(expected); @@ -259,12 +203,12 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u try std.testing.expectEqual(fromBf(values(cv, u16, s.capacity * width)[((last.base + n - 1) * s.kv_heads + kh) * s.head_dim + d]), outputValue(out, head * s.head_dim + d, output)); try std.testing.expectEqual(@as(f32, 0), outputValue(out, ((n - 1) * s.query_heads + head) * s.head_dim + d, output)); }; - for ([_]mtl.Buffer{ ck, cv, raw_k, prepared_k, raw_v, positions, frequencies, raw_q, q, masks, out, serial }) |b| try guards(b); + for ([_]mtl.Buffer{ ck, cv, raw_k, raw_v, q, masks, out, serial }) |b| try guards(b); if (n == 1) counts.ordinary += 1; counts.prompt_bytes += prompt * width * 4; counts.rows += n; counts.cases += 1; - std.debug.print("PASS N{d} P{d} QH{d}/KVH{d} D{d} RoPE{d} {s} {s}: data masks, peer-current visibility, serial-row equality, prompt unchanged\n", .{ n, prompt, s.query_heads, s.kv_heads, s.head_dim, s.rotary_dim, @tagName(mode), @tagName(output) }); + std.debug.print("PASS N{d} P{d} QH{d}/KVH{d} D{d} {s}: data masks, peer-current visibility, serial-row equality, prompt unchanged\n", .{ n, prompt, s.query_heads, s.kv_heads, s.head_dim, @tagName(output) }); } /// Compare split attention to serial attention and fp64 at draft shapes and trained-block masks. @@ -272,7 +216,7 @@ fn splitCase(device: mtl.Device, queue: mtl.Queue, keys: usize, counts: *Counts) const pool = mtl.objc.Pool.push(); defer pool.pop(); const rows = 16; - const s = shared.Shape{ .query_heads = 32, .kv_heads = 8, .head_dim = 128, .rotary_dim = 128, .capacity = 2064, .output = .f32 }; + const s = shared.Shape{ .query_heads = 32, .kv_heads = 8, .head_dim = 128, .capacity = 2064, .output = .f32 }; const p = try shared.Attention.init(a, device, s); defer p.deinit(); const qw = s.query_heads * s.head_dim; @@ -339,14 +283,14 @@ pub fn main(_: std.process.Init) !void { const queue = try device.queue(); defer queue.deinit(); var counts: Counts = .{}; - for ([_][4]usize{ .{ 2, 1, 32, 0 }, .{ 4, 2, 64, 32 }, .{ 8, 2, 128, 128 }, .{ 2, 2, 256, 64 } }) |dims| for ([_]shared.Rope{ .split_half, .interleaved }) |mode| for (1..17) |n| { - try run(device, queue, n, if (n % 3 == 0) 0 else if (n % 3 == 1) 5 else 33, dims, mode, .f32, &counts); - }; - for ([_][4]usize{ .{ 2, 1, 32, 0 }, .{ 4, 2, 64, 32 }, .{ 8, 2, 128, 128 }, .{ 2, 2, 256, 64 } }) |dims| for ([_]shared.Rope{ .split_half, .interleaved }) |mode| { - for ([_]usize{ 0, 1, 33 }) |prompt| try run(device, queue, 1, prompt, dims, mode, .f32, &counts); - try run(device, queue, 1, 5, dims, mode, .bf16, &counts); - try run(device, queue, 8, 5, dims, mode, .bf16, &counts); + for ([_][3]usize{ .{ 2, 1, 32 }, .{ 4, 2, 64 }, .{ 8, 2, 128 }, .{ 2, 2, 256 } }) |dims| for (1..17) |n| { + try run(device, queue, n, if (n % 3 == 0) 0 else if (n % 3 == 1) 5 else 33, dims, .f32, &counts); }; + for ([_][3]usize{ .{ 2, 1, 32 }, .{ 4, 2, 64 }, .{ 8, 2, 128 }, .{ 2, 2, 256 } }) |dims| { + for ([_]usize{ 0, 1, 33 }) |prompt| try run(device, queue, 1, prompt, dims, .f32, &counts); + try run(device, queue, 1, 5, dims, .bf16, &counts); + try run(device, queue, 8, 5, dims, .bf16, &counts); + } for ([_]usize{ 300, 1100, 2063 }) |keys| try splitCase(device, queue, keys, &counts); - std.debug.print("PASS {d} cases, {d} ordinary N=1 comparisons, {d} batch/serial rows, {d} fp32 CPU attention values, {d} bf16 values, {d} RoPE values, {d} prompt bytes preserved, max fp32 error {d:.9}, max bound fraction {d:.6}\n", .{ counts.cases, counts.ordinary, counts.rows, counts.cpu_values, counts.bf16_values, counts.rope_values, counts.prompt_bytes, counts.max_error, counts.max_fraction }); + std.debug.print("PASS {d} cases, {d} ordinary N=1 comparisons, {d} batch/serial rows, {d} fp32 CPU attention values, {d} bf16 values, {d} prompt bytes preserved, max fp32 error {d:.9}, max bound fraction {d:.6}\n", .{ counts.cases, counts.ordinary, counts.rows, counts.cpu_values, counts.bf16_values, counts.prompt_bytes, counts.max_error, counts.max_fraction }); } From e751f295828c6bfac02b8d0d88546d68c28d7d3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:13:06 +0100 Subject: [PATCH 080/186] runtime sources: drop the Rope loop and rotary fields core.zig still passed to shared_attention; changelog line --- CHANGELOG.md | 1 + zig/runtime_sources/core.zig | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 716bd8c6a93..fbe77abbd46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -82,6 +82,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. +- Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. diff --git a/zig/runtime_sources/core.zig b/zig/runtime_sources/core.zig index dbc6cb491fd..6472609c1bc 100644 --- a/zig/runtime_sources/core.zig +++ b/zig/runtime_sources/core.zig @@ -46,10 +46,10 @@ fn projections(s: *Set) !void { try s.fragments("core-lane", "prepared", try lane.source(s.a, l)); }; const attention = @import("../src/core/shared_attention.zig"); - for (std.enums.values(attention.Rope)) |rope| for (std.enums.values(attention.Output)) |out| { - const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .rotary_dim = 64, .capacity = 4096, .rope = rope, .output = out }); + for (std.enums.values(attention.Output)) |out| { + const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .capacity = 4096, .output = out }); try s.add("core", "shared-attention", text, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); - }; + } const route = @import("../src/core/moe_route.zig"); for ([_]route.Shape{ .{ .hidden = 4096, .experts = 288, .topk = 8 }, .{ .hidden = 4096, .experts = 128, .topk = 8, .out_bf16 = true }, .{ .hidden = 6144, .experts = 128, .topk = 4 } }) |shape| try s.raw("core", "route", try route.source(s.a, shape)); From 93ad638272872289ab5fc0e8dadfd261ebde11ae Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 09:05:30 +0000 Subject: [PATCH 081/186] deepseek_v41: tf-dsv41-assets builds ReleaseSafe whatever -Doptimize says dsv41_cuda.module pinned the family, kernel and runtime modules to Debug, and dsv41-assets built tf-dsv41-assets on it, so the AOT planner always ran unoptimized. Defilan measured one aot-needs enumeration at --context 524288 on a GB10: about 45 minutes as built, under 7 with the family modules at ReleaseSafe, planning the same 138 variants (#2270). module and runtime keep their Debug signatures for every other caller; moduleAt and runtimeAt take the mode, and dsv41-assets builds the tool on moduleAt(.ReleaseSafe), so its runtime and kernel modules go to ReleaseSafe too. ReleaseSafe keeps the safety checks, and no code under the tool reads the optimize mode, so the plan is unchanged. test-dsv41-assets-host stays on the Debug module. Refs #2270. Co-authored-by: Defilan <7034124+Defilan@users.noreply.github.com> --- zig/build/dsv41_assets.zig | 5 +++-- zig/build/dsv41_cuda.zig | 17 ++++++++++++----- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/zig/build/dsv41_assets.zig b/zig/build/dsv41_assets.zig index 06c612faa61..aa950c2ffa5 100644 --- a/zig/build/dsv41_assets.zig +++ b/zig/build/dsv41_assets.zig @@ -2,7 +2,8 @@ const std = @import("std"); pub fn steps(b: *std.Build) void { const step = b.step("dsv41-assets", "Install CPU-only DeepSeek token-map and AOT-plan generators"); - const asset_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false); + // ReleaseSafe whatever -Doptimize says: a 512K aot-needs pass took ~45 min in Debug on a GB10, under 7 here (#2270) + const asset_module = @import("dsv41_cuda.zig").moduleAt(b, b.graph.host, .ReleaseSafe, "asset_main.zig", false); const assets = b.addExecutable(.{ .name = "tf-dsv41-assets", .root_module = asset_module }); step.dependOn(&b.addInstallArtifact(assets, .{}).step); const token_map = b.addExecutable(.{ .name = "tf-dsv41-token-map", .root_module = b.createModule(.{ @@ -12,6 +13,6 @@ pub fn steps(b: *std.Build) void { .link_libc = true, }) }); step.dependOn(&b.addInstallArtifact(token_map, .{}).step); - const tests = b.addRunArtifact(b.addTest(.{ .root_module = asset_module })); + const tests = b.addRunArtifact(b.addTest(.{ .root_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false) })); b.step("test-dsv41-assets-host", "DeepSeek asset commands refuse GPU operations and malformed arguments").dependOn(&tests.step); } diff --git a/zig/build/dsv41_cuda.zig b/zig/build/dsv41_cuda.zig index 6b64f7d380e..b3e19f0cc64 100644 --- a/zig/build/dsv41_cuda.zig +++ b/zig/build/dsv41_cuda.zig @@ -2,11 +2,14 @@ const std = @import("std"); const shared = @import("dsv41_shared.zig"); const units = @import("dsv41_units.zig"); pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *std.Build.Module { + return runtimeAt(b, target, .Debug, images); +} +pub fn runtimeAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, images: bool) *std.Build.Module { const options = @import("cuda_runtime.zig").kernelOptions(b, images, null); - const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); cuda.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); cuda.addOptions("kernel_options", options); - cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = .Debug })); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); if (images) { const empty = b.addWriteFiles().add("compile-only.fatbin", "SYNTHETIC-COMPILE-ONLY"); const all = @import("cuda.zig").kernels; @@ -15,16 +18,20 @@ pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *s return cuda; } pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8, images: bool) *std.Build.Module { - const cuda = runtime(b, target, images); + return moduleAt(b, target, .Debug, file, images); +} +/// `module` with the runtime, kernel and family modules at `optimize`. +pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, images: bool) *std.Build.Module { + const cuda = runtimeAt(b, target, optimize, images); const opts = b.addOptions(); opts.addOption(bool, "with_kernels", images); opts.addOption(u32, "family_scope", 41); - const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = optimize, .link_libc = true }); dk.addImport("cuda", cuda); dk.addOptions("dsv41_kernel_options", opts); const empty = b.addWriteFiles().add("compile-only-dsv41.fatbin", "SYNTHETIC-COMPILE-ONLY"); for (units.units ++ units.optional) |k| dk.addAnonymousImport(b.fmt("dsv41_fatbin_{s}", .{k.name}), .{ .root_source_file = empty }); - return assemble(b, target, .Debug, file, cuda, dk); + return assemble(b, target, optimize, file, cuda, dk); } /// The family's root `file` over the runtime `cuda` and kernel module `dk`, every family module at `optimize`. pub fn assemble(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, cuda: *std.Build.Module, dk: *std.Build.Module) *std.Build.Module { From d6ddf2b007ebba93079f027a7a515d37afa4a6aa Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:02 +0000 Subject: [PATCH 082/186] changelog: #3007 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index fbe77abbd46..c3007ea6f1f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. + - Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). From ad3f7384c78f8957461583a841962e2a6156aeb6 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:30:14 +0000 Subject: [PATCH 083/186] cuda experts: a bf16 format, a checkpoint's own weights in mma fragment blocks (relu^2 up, fp32 or bf16 down) --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/experts_bf16.cu | 114 +++++++++++++++ zig/src/cuda/experts.zig | 12 +- zig/src/cuda/experts_bf16.zig | 94 ++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/tests/cuda/bf16_experts_tests.zig | 201 ++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 2 + 7 files changed, 421 insertions(+), 4 deletions(-) create mode 100644 zig/kernels/cuda/experts_bf16.cu create mode 100644 zig/src/cuda/experts_bf16.zig create mode 100644 zig/tests/cuda/bf16_experts_tests.zig diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index c0c8b503336..950d2d5586f 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -97,6 +97,7 @@ pub const kernels = [_]Kernel{ .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order .{ .name = "nvfp4_experts_tile", .flags = &.{"-O3"}, .min_sm = 80 }, + .{ .name = "experts_bf16", .flags = &.{"-O3"} }, // ours: grouped bf16 experts on the plan, mma fragment blocks .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "nvfp4_experts32", .flags = &.{"-O3"} }, // ours: GLM's routed experts at block 32 with csf0 scales .{ .name = "rank_sum", .flags = &.{"-O3"} }, // ours: the rank-order sum of a reduce-scatter (rank_scatter.zig) diff --git a/zig/kernels/cuda/experts_bf16.cu b/zig/kernels/cuda/experts_bf16.cu new file mode 100644 index 00000000000..df6d756ec18 --- /dev/null +++ b/zig/kernels/cuda/experts_bf16.cu @@ -0,0 +1,114 @@ +// Ours: grouped bf16 experts on mma-fragment blocks; a row's sums chain in input order, whatever its tile or form. + +#include +#include + +#include "experts.cuh" + +namespace tf_bf16_experts { + +constexpr int WV = 4; // uint4 a lane a block + +__device__ __forceinline__ uint4 ld_w(const uint4* p) { return tf_common::ld_nc(p); } + +// RT row tiles of 16 pairs share each block's fragments; a unit is one item's 32-column block. +template +__device__ __forceinline__ void body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint4* __restrict__ W, int KG, int NB, const int* __restrict__ items, + const int* __restrict__ counts, const int* __restrict__ members, + void* __restrict__ out, int N, int skip) { + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x * WARPS + (threadIdx.x >> 5); unit < units; unit += gridDim.x * WARPS) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + if (e == skip) continue; + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (32 * WV) + lane; + for (int r0 = 0; r0 < cnt; r0 += 16 * RT) { + bool v[RT][2]; + int pr[RT][2]; + const __nv_bfloat16* x[RT][2]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = r0 + 16 * r + gq + 8 * h; + v[r][h] = i < cnt; + pr[r][h] = v[r][h] ? __ldg(members + first + i) : 0; + x[r][h] = X + (size_t)(slots ? pr[r][h] / slots : pr[r][h]) * x_stride + 4 * t; + } + float acc[1][RT][NTW][4]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[0][r][j][0] = acc[0][r][j][1] = acc[0][r][j][2] = acc[0][r][j][3] = 0.f; + const uint2 zero = make_uint2(0u, 0u); + for (int g = 0; g < KG; ++g) { + uint4 w[NTW]; +#pragma unroll + for (int j = 0; j < NTW; ++j) w[j] = ld_w(blk + (size_t)g * (32 * WV) + 32 * j); +#pragma unroll + for (int r = 0; r < RT; ++r) { + uint2 xa[2], xb[2]; +#pragma unroll + for (int q = 0; q < 2; ++q) { + xa[q] = v[r][0] ? __ldg(reinterpret_cast(x[r][0] + 32 * g + 16 * q)) : zero; + xb[q] = v[r][1] ? __ldg(reinterpret_cast(x[r][1] + 32 * g + 16 * q)) : zero; + } +#pragma unroll + for (int q = 0; q < 2; ++q) +#pragma unroll + for (int j = 0; j < NTW; ++j) + mma(acc[0][r][j], xa[q].x, xb[q].x, xa[q].y, xb[q].y, q ? w[j].z : w[j].x, q ? w[j].w : w[j].y); + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr[r][0], pr[r][1], v[r][0], v[r][1], 0.f); + } + } +} + +} // namespace tf_bf16_experts + +// Decode (items of up to 16 pairs) and prompt (up to 64) entries per epilogue: the NVFP4 entries' arguments, no scales. +#define TF_BF16_EXPERTS(NAME, EPI, RT, WARPS) \ + extern "C" __global__ void __launch_bounds__(WARPS * 32) \ + NAME(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, int NB, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int skip) { \ + tf_bf16_experts::body(X, x_stride, slots, W, KG, NB, items, counts, members, out, N, skip); \ + } +TF_BF16_EXPERTS(tf_bf16_experts_up, 1, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_f32, 0, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_bf16, 3, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_up, 1, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_f32, 0, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_bf16, 3, 4, 4) + +// The repack, one thread a uint4: matrix e's [k][n] (in_out, x @ W) or [n][k] (nn.Linear) rows into its blocks. +extern "C" __global__ void tf_bf16_experts_pack(const __nv_bfloat16* __restrict__ src, long long e_stride, int ld, + int n0, int k0, int in_out, int n, int k, uint4* __restrict__ dst) { + const int e = blockIdx.y; + const long long per = (long long)(n / 32) * (k / 32) * 32 * tf_bf16_experts::WV; + const long long q = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (q >= per) return; + const int lane = (int)(q % 32), j = (int)(q / 32 % tf_bf16_experts::WV); + const long long b = q / (32 * tf_bf16_experts::WV); + const int kg = k / 32, g = (int)(b % kg), cb = (int)(b / kg); + const int row = n0 + cb * 32 + 8 * j + (lane >> 2), col = k0 + g * 32 + 4 * (lane & 3); + const __nv_bfloat16* s = src + e * e_stride; + uint16_t h[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int c = col + (i & 3) + 16 * (i >> 2); + const __nv_bfloat16 v = in_out ? s[(size_t)c * ld + row] : s[(size_t)row * ld + c]; + h[i] = __bfloat16_as_ushort(v); + } + uint4 o; + o.x = h[0] | (uint32_t)h[1] << 16; + o.y = h[2] | (uint32_t)h[3] << 16; + o.z = h[4] | (uint32_t)h[5] << 16; + o.w = h[6] | (uint32_t)h[7] << 16; + dst[e * per + q] = o; +} diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 3d70527c1d8..3f7bb4de962 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -14,15 +14,16 @@ const Function = module.Function; const Stream = stream_.Stream; const Plan = grouped.Plan; const Affine3Experts = @import("experts3.zig").Affine3Experts; +pub const Bf16Experts = @import("experts_bf16.zig").Bf16Experts; -pub const Format = enum { affine4, affine3, fp8g, nvfp4, nvfp4_csf0, exl3, iq2xxs }; +pub const Format = enum { affine4, affine3, fp8g, nvfp4, nvfp4_csf0, exl3, iq2xxs, bf16 }; /// One layer's experts on the device as each format's make() lays them out; a family stacks its shared ones last. pub const Layer = struct { format: Format, exl3: ?@import("exl3_experts.zig").Layer = null, - up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; nvfp4_csf0 [E][NI/32] blocks (experts32.zig); iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig) - down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; nvfp4_csf0 [E][D/32] blocks; iq2xxs [E][D/16][NI/256][1056] + up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; nvfp4_csf0 [E][NI/32] blocks (experts32.zig); iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig); bf16 [E][NI/32][D/32][4][32] uint4 + down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; nvfp4_csf0 [E][D/32] blocks; iq2xxs [E][D/16][NI/256][1056]; bf16 [E][D/32][NI/32][4][32] up_scale: u64 = 0, // fp8g [E][2][NI/128][D/128] fp32; nvfp4 [E][2] global scales; affine4's are in the blocks down_scale: u64 = 0, // fp8g [E][1][D/128][NI/128]; nvfp4 [E][1] width: usize, // NI @@ -52,6 +53,7 @@ pub const Grouped = struct { nvfp4: ?*const Nvfp4Experts = null, nvfp4_csf0: ?*const @import("experts32.zig").Nvfp4b32Experts = null, iq2xxs: ?*const Iq2xxsExperts = null, + bf16: ?*const Bf16Experts = null, /// Decode rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1. pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void { @@ -77,6 +79,7 @@ pub const Grouped = struct { .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), .nvfp4_csf0 => return (g.nvfp4_csf0 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, false, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, false, epi, in, l, p, out, items, skip), } } @@ -104,6 +107,7 @@ pub const Grouped = struct { .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), .nvfp4_csf0 => return (g.nvfp4_csf0 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, true, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, true, epi, in, l, p, out, items, skip), } } }; @@ -111,7 +115,7 @@ pub const Grouped = struct { /// Pairs a prompt item holds, per format (the Python engines' prompt tiles). pub fn promptTile(f: Format) usize { return switch (f) { - .affine4, .affine3, .fp8g, .nvfp4, .nvfp4_csf0, .exl3, .iq2xxs => 64, + .affine4, .affine3, .fp8g, .nvfp4, .nvfp4_csf0, .exl3, .iq2xxs, .bf16 => 64, }; } diff --git a/zig/src/cuda/experts_bf16.zig b/zig/src/cuda/experts_bf16.zig new file mode 100644 index 00000000000..76ed554c2f7 --- /dev/null +++ b/zig/src/cuda/experts_bf16.zig @@ -0,0 +1,94 @@ +//! bf16 experts (experts_bf16.cu): a checkpoint's bf16 weights in mma fragment blocks, relu^2 up, fp32 or bf16 down. +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const grouped = @import("grouped.zig"); +const driver = @import("driver.zig"); +const experts = @import("experts.zig"); + +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; +const Plan = grouped.Plan; +const Epi = experts.Epi; +const Rows = experts.Rows; +const Layer = experts.Layer; + +pub const Bf16Experts = struct { + mod: Module, + fns: [2][3]Function, // decode, prompt; by Epi + resident: [3]usize, // decode blocks the grid may hold: per SM times SMs + pack_fn: Function, + + pub const cols = 32; // output columns a block + pub const warps = 4; + /// Bytes of a (32 columns, 32 inputs) block. + pub const block = 32 * 32 * 2; + + pub const symbols = struct { + pub const decode = [3][:0]const u8{ "tf_bf16_experts_up", "tf_bf16_experts_down_f32", "tf_bf16_experts_down_bf16" }; + pub const prompt = [3][:0]const u8{ "tf_bf16_experts_prompt_up", "tf_bf16_experts_prompt_down_f32", "tf_bf16_experts_prompt_down_bf16" }; + pub const pack = "tf_bf16_experts_pack"; + }; + + pub fn load(d: *const driver.Driver, sms: usize) !Bf16Experts { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.experts_bf16); + errdefer mod.unload(); + var e: Bf16Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .pack_fn = try mod.function(symbols.pack) }; + for (symbols.decode, symbols.prompt, 0..) |dn, pn, i| { + e.fns[0][i] = try mod.function(dn); + e.fns[1][i] = try mod.function(pn); + e.resident[i] = @max(1, try e.fns[0][i].occupancy(warps * 32, 0)) * sms; + } + return e; + } + + pub fn unload(e: *Bf16Experts) void { + e.mod.unload(); + } + + /// One call on a plan of at most `items` items: 16 pairs an item in decode, 64 in a prompt. + pub fn run(e: *const Bf16Experts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + if (!l.relu2 or l.swiglu or l.limit != 0) return error.OptionNotBuilt; + const up = epi == .up; + const kg = (if (up) l.dims else l.width) / 32; + const nb = (if (up) l.width else l.dims) / cols; + const ei = @backingInt(epi); + const units = items * nb; + const grid = if (prompt) (units + warps - 1) / warps else @min((units + warps - 1) / warps, e.resident[ei]); + if (grid < 1) return; + var a: launch_.Args = .{}; + a.add(in.x); + a.add(@as(i32, @intCast(in.stride))); + a.add(@as(i32, @intCast(if (up) in.slots else 0))); + a.add(if (up) l.up else l.down); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(nb))); + for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); + a.add(@as(i32, @intCast(if (up) l.width else l.dims))); + a.add(skip); + try launch_.launch(e.fns[@intFromBool(prompt)][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = warps * 32, .y = 1, .z = 1 } }, s, &a); + } + + /// Bytes of `count` experts' blocks for n outputs from k inputs. + pub fn bytes(count: usize, n: usize, k: usize) usize { + return count * (n / cols) * (k / 32) * block; + } + + /// A source matrix: experts `stride` apart, rows of `ld`, its first output and input, [k][n] (x @ W) or [n][k]. + pub const Source = struct { stride: usize = 0, ld: usize, n0: usize = 0, k0: usize = 0, in_out: bool }; + + /// `count` bf16 matrices of n outputs from k inputs at `src` into their blocks at `dst` (bytes(count, n, k)). + pub fn pack(e: *const Bf16Experts, s: Stream, src: u64, from: Source, dst: u64, count: usize, n: usize, k: usize) !void { + if (n % cols != 0 or k % 32 != 0) return error.UnexpectedShape; + const per = bytes(1, n, k) / 16; + var a: launch_.Args = .{}; + a.add(src); + a.add(@as(i64, @intCast(from.stride))); + for ([_]usize{ from.ld, from.n0, from.k0, @intFromBool(from.in_out), n, k }) |v| a.add(@as(i32, @intCast(v))); + a.add(dst); + try launch_.launch(e.pack_fn, .{ .grid = .{ .x = @intCast((per + 255) / 256), .y = @intCast(count), .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 1ff9f6a4489..c8eb892ca12 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -51,6 +51,7 @@ pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").b pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; pub const nvfp4_experts_tile: []const u8 = if (available) &Blob("fatbin_nvfp4_experts_tile").bytes else &.{}; +pub const experts_bf16: []const u8 = if (available) &Blob("fatbin_experts_bf16").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const nvfp4_experts32: []const u8 = if (available) &Blob("fatbin_nvfp4_experts32").bytes else &.{}; pub const rank_sum: []const u8 = if (available) &Blob("fatbin_rank_sum").bytes else &.{}; diff --git a/zig/tests/cuda/bf16_experts_tests.zig b/zig/tests/cuda/bf16_experts_tests.zig new file mode 100644 index 00000000000..323178e9476 --- /dev/null +++ b/zig/tests/cuda/bf16_experts_tests.zig @@ -0,0 +1,201 @@ +//! Grouped bf16 experts against float64: the repack, up and downs, a row's bits the same at any row count and form. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +const E = 6; // routed experts; the shared expert's halves are experts E and E + 1 +const D = 256; +const W = 96; +const K = 3; // routed picks a row; a row has K + 2 pairs + +fn bf(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +fn f(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Host = struct { + routed_up: []u16, // [E][D][W]: x @ W, as the stacked checkpoint tensors + routed_down: []u16, // [E][W][D] + shared_up: []u16, // [2W][D]: nn.Linear + shared_down: []u16, // [D][2W] + + /// Expert e's weight from input i to output o, up or down, as the kernels must read it. + fn at(h: Host, up: bool, e: usize, o: usize, i: usize) f64 { + if (e < E) return f(if (up) h.routed_up[(e * D + i) * W + o] else h.routed_down[(e * W + i) * D + o]); + const half = e - E; + return f(if (up) h.shared_up[(half * W + o) * D + i] else h.shared_down[o * 2 * W + half * W + i]); + } +}; + +pub fn experts(gpu: Gpu) !void { + const a = gpu.gpa; + var rng = std.Random.DefaultPrng.init(0x7f16); + const r = rng.random(); + var host: Host = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, .{ E * D * W, E * W * D, 2 * W * D, D * 2 * W }) |name, n| { + @field(host, name) = try a.alloc(u16, n); + for (@field(host, name)) |*v| v.* = bf((r.float(f32) - 0.5) * 0.25); + } + defer inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }) |name| a.free(@field(host, name)); + + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + defer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + + const e = E + 2; + var up = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, W, D)); + defer up.free(); + var down = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, D, W)); + defer down.free(); + var src: [4]cuda.DeviceBuffer = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, 0..) |name, i| src[i] = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(@field(host, name))); + defer for (&src) |*b| b.free(); + const one_up = Bf16.bytes(1, W, D); + const one_down = Bf16.bytes(1, D, W); + try kx.pack(s, src[0].ptr, .{ .stride = D * W, .ld = W, .in_out = true }, up.ptr, E, W, D); + try kx.pack(s, src[1].ptr, .{ .stride = W * D, .ld = D, .in_out = true }, down.ptr, E, D, W); + for (0..2) |h| { + try kx.pack(s, src[2].ptr, .{ .ld = D, .n0 = h * W, .in_out = false }, up.ptr + (E + h) * one_up, 1, W, D); + try kx.pack(s, src[3].ptr, .{ .ld = 2 * W, .k0 = h * W, .in_out = false }, down.ptr + (E + h) * one_down, 1, D, W); + } + try s.synchronize(); + const layer: ex_.Layer = .{ .format = .bf16, .up = up.ptr, .down = down.ptr, .width = W, .dims = D, .experts = e, .relu2 = true }; + const g: ex_.Grouped = .{ .bf16 = &kx }; + + // 37 rows: three 16-pair decode items for the busiest expert and a partial prompt tile; then each row alone + const rows = 37; + const slots = K + 2; + const x = try a.alloc(u16, rows * D); + defer a.free(x); + for (x) |*v| v.* = bf(r.float(f32) * 2 - 1); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + for (0..rows) |row| { + var chosen: [K]i32 = undefined; + for (0..K) |k| { + var p: i32 = undefined; + while (true) { + p = @intCast(r.uintLessThan(usize, E)); + if (std.mem.indexOfScalar(i32, chosen[0..k], p) == null) break; + } + chosen[k] = p; + picks[row * slots + k] = p; + } + picks[row * slots + K] = E; + picks[row * slots + K + 1] = E + 1; + } + + // the float64 reference: act = bf16(relu(bf16(x . up))^2), y = act . down + const act_ref = try a.alloc(f64, rows * slots * W); + defer a.free(act_ref); + const y_ref = try a.alloc(f64, rows * slots * D); + defer a.free(y_ref); + const mag = try a.alloc(f64, rows * slots * D); + defer a.free(mag); + for (0..rows * slots) |pr| { + const ex: usize = @intCast(picks[pr]); + const row = pr / slots; + for (0..W) |o| { + var acc: f64 = 0; + for (0..D) |i| acc += f(x[row * D + i]) * host.at(true, ex, o, i); + const u = @max(f(bf(@floatCast(acc))), 0); + act_ref[pr * W + o] = f(bf(@floatCast(u * u))); + } + for (0..D) |o| { + var acc: f64 = 0; + var m: f64 = 0; + for (0..W) |i| { + const t = act_ref[pr * W + i] * host.at(false, ex, o, i); + acc += t; + m += @abs(t); + } + y_ref[pr * D + o] = acc; + mag[pr * D + o] = m; + } + } + + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var whole: [3][]u8 = undefined; // decode form, all rows: act, y fp32, y bf16 + var got = [3]bool{ false, false, false }; + defer for (whole, got) |w, ok| if (ok) a.free(w); + for ([_]bool{ false, true }) |prompt| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr, picks, rows, slots, prompt); + defer for (run) |b| a.free(b); + const act_h = std.mem.bytesAsSlice(u16, run[0]); + var worst: f64 = 0; + for (act_h, act_ref) |gv, want| worst = @max(worst, @abs(f(gv) - want) / @max(@abs(want), 1e-3)); + try check.expect(worst <= 1.0 / 64.0, "bf16 experts {s}: up off by {e} (relative)", .{ if (prompt) "prompt" else "decode", worst }); + const y32 = std.mem.bytesAsSlice(f32, run[1]); + var worst_y: f64 = 0; + for (y32, y_ref, mag) |gv, want, m| worst_y = @max(worst_y, @abs(@as(f64, gv) - want) / @max(m, 1e-6)); + try check.expect(worst_y <= 1e-5, "bf16 experts {s}: fp32 down off by {e} of its terms' magnitude", .{ if (prompt) "prompt" else "decode", worst_y }); + if (!prompt) { + for (run, &whole, &got) |b, *w, *ok| { + w.* = try a.dupe(u8, b); + ok.* = true; + } + } else for (run, whole, [_][]const u8{ "prompt up", "prompt fp32 down", "prompt bf16 down" }) |b, w, what| try check.sameBytes(what, b, w); + check.pass("bf16 experts: {d} rows x {d} pairs ({s}) up within {e}, fp32 down within {e} of the float64 reference", .{ rows, slots, if (prompt) "prompt, items of 64" else "decode, items of 16", worst, worst_y }); + } + // each row alone gives the bytes it gave beside the others + for ([_]usize{ 0, 17, rows - 1 }) |row| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr + row * D * 2, picks[row * slots ..][0..slots], 1, slots, false); + defer for (run) |b| a.free(b); + for (run, whole, [_]usize{ W * 2, D * 4, D * 2 }, [_][]const u8{ "row up", "row fp32 down", "row bf16 down" }) |b, w, per, what| try check.sameBytes(what, b, w[row * slots * per ..][0 .. slots * per]); + } + check.pass("bf16 experts: a row alone, beside 36 others and in the prompt form gives the same bytes", .{}); +} + +/// Plan, up and both downs over `rows` rows; returns act, fp32 y and bf16 y. +fn oneCase(gpu: Gpu, router: grouped.Router, g: ex_.Grouped, s: cuda.Stream, l: ex_.Layer, x: u64, picks: []const i32, rows: usize, slots: usize, prompt: bool) ![3][]u8 { + const pairs = rows * slots; + const tile: usize = if (prompt) ex_.promptTile(.bf16) else 16; + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer dpicks.free(); + const bound = grouped.maxItems(pairs, l.experts, tile); + var bufs: [5]cuda.DeviceBuffer = undefined; + const big = pairs > grouped.small; + const sizes = [5]usize{ pairs * 4, bound * 12, 8, (if (big) pairs else 1) * 4, (if (big) (pairs + 1023) / 1024 * l.experts else 1) * 4 }; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, n); + made += 1; + } + const plan: grouped.Plan = .{ .members = bufs[0].ptr, .items = bufs[1].ptr, .counts = bufs[2].ptr, .rank = bufs[3].ptr, .hist = bufs[4].ptr }; + try router.route(s, dpicks.ptr, pairs, l.experts, tile, plan); + var act = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.width * 2); + defer act.free(); + var y32 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 4); + defer y32.free(); + var y16 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 2); + defer y16.free(); + for ([_]*cuda.DeviceBuffer{ &act, &y32, &y16 }) |b| try b.fill8(0xff, s.handle); // a call that writes nothing fails + const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; + try call(g, s, .up, .{ .x = x, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, pairs, tile, -1); + try call(g, s, .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y32.ptr, bound, pairs, tile, -1); + try call(g, s, .down_bf16, .{ .x = act.ptr, .stride = l.width }, l, plan, y16.ptr, bound, pairs, tile, -1); + try s.synchronize(); + var out: [3][]u8 = undefined; + out[0] = try check.download(gpu, act); + errdefer gpu.gpa.free(out[0]); + out[1] = try check.download(gpu, y32); + errdefer gpu.gpa.free(out[1]); + out[2] = try check.download(gpu, y16); + return out; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f9f39206584..ebd3f880790 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -215,6 +215,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ bf16-experts grouped bf16 experts against float64: repack, up, downs, rows and forms byte for byte \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) @@ -430,6 +431,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "bf16-experts")) return @import("bf16_experts_tests.zig").experts(gpu); if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); From d0ff0dcb519d1870bbda1bd6d68467a72972fd33 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:30:14 +0000 Subject: [PATCH 084/186] kernel contract: the bf16 experts (experts_bf16.cu), relu^2 up, fp32 and bf16 down, gates 1-5 with two controls --- docs/kernel-contract.md | 4 + zig/kernels/cuda/contract_controls_set.cu | 1 + .../cuda/contract_experts_bf16_controls.cuh | 24 ++ .../core/contract/experts_bf16_records.zig | 37 ++ zig/src/core/contract/registry_joins.zig | 5 +- zig/src/core/contract/registry_sets_test.zig | 1 + zig/tests/cuda/contract_experts_bf16.zig | 319 ++++++++++++++++++ zig/tests/cuda/contract_family_sets.zig | 5 + zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/main.zig | 4 +- 10 files changed, 398 insertions(+), 4 deletions(-) create mode 100644 zig/kernels/cuda/contract_experts_bf16_controls.cuh create mode 100644 zig/src/core/contract/experts_bf16_records.zig create mode 100644 zig/tests/cuda/contract_experts_bf16.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 6b4d03873ff..520a89bc5d5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -401,6 +401,10 @@ CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 down, each bit-equal to the decode instance it replaces; the scaled-code prompt tile is the planted control). +CUDA also checks the bf16 experts (`experts_bf16.cu`, a checkpoint's own bf16 weights: Nemotron 3.5 Transcribe's routed +and shared experts): relu^2 up, fp32 down and bf16 down, in the decode and prompt forms, against float64 (two controls: +the fp32 down launched one 32-input group short, and a down summing in bf16). nvcc 13.2's prompt forms and 13.3's +decode forms on sm_120 and sm_121 take other registers than 13.0's and have their own records. CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's diff --git a/zig/kernels/cuda/contract_controls_set.cu b/zig/kernels/cuda/contract_controls_set.cu index e4608a5b9f7..3624ac10b03 100644 --- a/zig/kernels/cuda/contract_controls_set.cu +++ b/zig/kernels/cuda/contract_controls_set.cu @@ -3,3 +3,4 @@ #include "contract_controls.cu" #include "contract_js_controls.cuh" #include "contract_dense_controls.cuh" +#include "contract_experts_bf16_controls.cuh" diff --git a/zig/kernels/cuda/contract_experts_bf16_controls.cuh b/zig/kernels/cuda/contract_experts_bf16_controls.cuh new file mode 100644 index 00000000000..1af8126ec29 --- /dev/null +++ b/zig/kernels/cuda/contract_experts_bf16_controls.cuh @@ -0,0 +1,24 @@ +// Gate 4's control for experts_bf16.cu: its fp32 down with each pair's sum rounded to bf16 at every input. +#pragma once + +extern "C" __global__ void tf_ctl_bf16_experts_bf16_sums(const __nv_bfloat16* X, int x_stride, const uint4* W, int KG, int NB, + const int* items, const int* counts, const int* members, float* out, + int N) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= N) return; + const int cb = col / 32, j = (col % 32) / 8; + for (int it = blockIdx.y; it < counts[0]; it += gridDim.y) { + const int e = items[3 * it], first = items[3 * it + 1], cnt = items[3 * it + 2]; + const __nv_bfloat16* blk = reinterpret_cast(W + (static_cast(e) * NB + cb) * KG * 128); + for (int i = 0; i < cnt; ++i) { + const int pr = members[first + i]; + float s = 0.0f; + for (int k = 0; k < KG * 32; ++k) { + const int q = k % 16, lane = 4 * (col % 8) + q / 4; + const float w = __bfloat162float(blk[(static_cast(k / 32) * 128 + 32 * j + lane) * 8 + 4 * (k % 32 / 16) + q % 4]); + s = __bfloat162float(__float2bfloat16_rn(s + tf_ctl_x(X, static_cast(pr) * x_stride + k) * w)); + } + out[static_cast(pr) * N + col] = s; + } + } +} diff --git a/zig/src/core/contract/experts_bf16_records.zig b/zig/src/core/contract/experts_bf16_records.zig new file mode 100644 index 00000000000..7b4c4bf7784 --- /dev/null +++ b/zig/src/core/contract/experts_bf16_records.zig @@ -0,0 +1,37 @@ +//! The bf16 experts' CUDA launches (experts_bf16.cu): decode and prompt rows by epilogue, and the repack. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_up", .symbol = "tf_bf16_experts_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .symbol = "tf_bf16_experts_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .symbol = "tf_bf16_experts_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .symbol = "tf_bf16_experts_prompt_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .symbol = "tf_bf16_experts_prompt_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .symbol = "tf_bf16_experts_prompt_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_pack", .symbol = "tf_bf16_experts_pack", .image = "experts_bf16", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, +}; + +/// nvcc 13.2's prompt forms on sm_120 and sm_121 (164 registers where 13.0 takes 162); the rest compile as 13.0's. +pub const cuda_132 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, +}; + +/// nvcc 13.3's decode forms on sm_120 and sm_121 (70 registers where 13.0 and 13.2 take 72). +pub const cuda_133 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_up", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, +}; + +test "bf16 experts CUDA: every launch is in the live registry and fits every SM" { + const t = @import("std").testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(@import("std").mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/registry_joins.zig b/zig/src/core/contract/registry_joins.zig index 615190a1678..782618b89bd 100644 --- a/zig/src/core/contract/registry_joins.zig +++ b/zig/src/core/contract/registry_joins.zig @@ -13,6 +13,7 @@ pub const cuda_after = @import("registry_media.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("nvfp4_experts32_records.zig").cuda ++ + @import("experts_bf16_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @@ -30,6 +31,7 @@ pub const nvcc_132_after = @import("registry_learning.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ + @import("experts_bf16_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132; pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @@ -37,7 +39,8 @@ pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133 ++ @import("minimax_cuda_records.zig").cuda_133 ++ @import("nvfp4_experts32_records.zig").cuda_133 ++ - @import("paging_pending.zig").cuda_133; + @import("paging_pending.zig").cuda_133 ++ + @import("experts_bf16_records.zig").cuda_133; pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("exl3_prompt_pending.zig").records ++ @import("cuda_lane_pending.zig").records ++ diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index 743c4b1a8cf..43979dd1ab3 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -54,4 +54,5 @@ test { _ = @import("exl3_grouped_records.zig"); _ = @import("proposal_sampling_records.zig"); _ = @import("pocket_records.zig"); + _ = @import("experts_bf16_records.zig"); } diff --git a/zig/tests/cuda/contract_experts_bf16.zig b/zig/tests/cuda/contract_experts_bf16.zig new file mode 100644 index 00000000000..4eb411f98b9 --- /dev/null +++ b/zig/tests/cuda/contract_experts_bf16.zig @@ -0,0 +1,319 @@ +//! bf16 experts (experts_bf16.cu) under the contract on Transcribe's shapes, decode and prompt forms; two controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +pub const Kind = enum { up, down_f32, down_bf16 }; + +/// A planted fault: the fp32 down launched one 32-input group short, or contract_controls.cu's bf16-summing down. +pub const Fault = enum { none, short, bf16_sums }; + +pub const Entry = struct { name: []const u8, kind: Kind, prompt: bool, fault: Fault = .none }; + +pub const entries = [_]Entry{ + .{ .name = "cuda bf16 experts up tf_bf16_experts_up (relu^2)", .kind = .up, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_f32", .kind = .down_f32, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_bf16", .kind = .down_bf16, .prompt = false }, + .{ .name = "cuda bf16 experts prompt up tf_bf16_experts_prompt_up (relu^2)", .kind = .up, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_f32", .kind = .down_f32, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_bf16", .kind = .down_bf16, .prompt = true }, + .{ .name = "control: bf16 experts fp32 down launched one 32-input group short", .kind = .down_f32, .prompt = false, .fault = .short }, + .{ .name = "control: bf16 experts fp32 down summing in bf16", .kind = .down_f32, .prompt = false, .fault = .bf16_sums }, +}; + +const experts = 4; +const slots = 8; // pairs a token row: Transcribe's 6 routed picks and its shared expert's two halves + +/// An mma k16 step a link of each row's chain, plus one step's error over its own products (attention_range's 36). +fn depth(s: c.Shape) usize { + return s.k / 16 + core.attention_range.mma_step; +} + +/// Transcribe's routed experts, up 2,688 to 1,856 and back; decode windows of 1 to 32 rows, prompt items of 2 and 9. +const ups = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 1856, .k = 2688 }, .{ .rows = 32, .n = 1856, .k = 2688 } }, &.{ .{ .rows = 40, .n = 1856, .k = 2688 }, .{ .rows = 72, .n = 1856, .k = 2688 } } }; +const downs = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 2688, .k = 1856 }, .{ .rows = 32, .n = 2688, .k = 1856 } }, &.{ .{ .rows = 40, .n = 2688, .k = 1856 }, .{ .rows = 72, .n = 2688, .k = 1856 } } }; + +pub fn spec(e: Entry) c.Spec { + const served = (if (e.kind == .up) ups else downs)[@intFromBool(e.prompt)]; + const tiles = [3]c.Tile{ .{ .size = (if (e.prompt) ex_.promptTile(.bf16) else 16) / slots }, .{ .size = 64, .step = 32, .min = 32 }, .{ .size = 64, .step = 32, .min = 32 } }; + return switch (e.kind) { + // bf16(acc)'s rounding enters the square twice + .up => .{ .name = e.name, .accum = .f32, .depth = depth, .ops = 1, .mids = 2, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .f16 = 30, .bf16 = 30, .why = "a sum that collapses below 0 gives relu^2 0, which misses by want itself, and the bf16 roundings already allow 3 u of it: at most 1 / (3 u) = 85x is provable" } }, + .down_f32 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .f32, .tiles = tiles, .served = served }, + .down_bf16 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .bf16, .tiles = tiles, .served = served }, + }; +} + +/// shape.rows token rows, `slots` pairs each over `experts` experts; up reads a token's row, down a pair's. +pub const Experts = struct { + rig: rig_.Rig, + entry: Entry, + kx: *const Bf16, + router: grouped.Router, + ctl: ?cuda.Function = null, + shape: c.Shape = undefined, + x: []u16 = &.{}, + w: []u16 = &.{}, // [experts][n][k], as nn.Linear stores them + picks: []i32 = &.{}, + xd: u64 = 0, + wd: u64 = 0, + plan: grouped.Plan = undefined, + out: u64 = 0, + + pub fn deinit(self: *Experts) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Experts) void { + const a = self.rig.gpu.gpa; + a.free(self.x); + a.free(self.w); + a.free(self.picks); + self.x = &.{}; + self.w = &.{}; + self.picks = &.{}; + } + + pub fn kernel(self: *Experts) g.Kernel { + const control: ?g.Gate = switch (self.entry.fault) { + .none => null, + .short => .shapes, + .bf16_sums => .sums, + }; + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Experts); + + fn up(self: *const Experts) bool { + return self.entry.kind == .up; + } + + fn pairs(self: *const Experts) usize { + return self.shape.rows * slots; + } + + fn xRows(self: *const Experts) usize { + return if (self.up()) self.shape.rows else self.pairs(); + } + + fn tile(self: *const Experts) usize { + return if (self.entry.prompt) ex_.promptTile(.bf16) else 16; + } + + fn items(self: *const Experts) usize { + return grouped.maxItems(self.pairs(), experts, self.tile()); + } + + /// Random bf16 rows and weights; non-finite: NaN, +Inf and -Inf in the first rows; stress: mirrored integers. + pub fn prepare(self: *Experts, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0 or shape.k == 0 or shape.k % 32 != 0) return false; + self.free(); + self.rig.reset(); + self.shape = shape; + const a = self.rig.gpu.gpa; + const n = shape.n; + const k = shape.k; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.x = try a.alloc(u16, self.xRows() * k); + self.w = try a.alloc(u16, experts * n * k); + self.picks = try a.alloc(i32, self.pairs()); + for (self.x) |*v| v.* = c.toBf16(rng.float(f32) * 2 - 1); + for (self.w) |*v| v.* = c.toBf16((rng.float(f32) - 0.5) * 0.25); + for (self.picks, 0..) |*p, i| p.* = @intCast(if (i < experts) i else rng.uintLessThan(usize, experts)); + switch (fill) { + .random => {}, + .nonfinite => { + // a row holding NaN, one holding +Inf, one holding both infinities + const planted = [_]struct { usize, usize, f32 }{ .{ 0, k / 3, std.math.nan(f32) }, .{ 1, k / 2, std.math.inf(f32) }, .{ 2, 1, std.math.inf(f32) }, .{ 2, k - 1, -std.math.inf(f32) } }; + for (planted) |p| if (p[0] < self.xRows()) { + self.x[p[0] * k + p[1]] = c.toBf16(p[2]); + }; + }, + .stress => self.stress(rng), + } + self.xd = try self.rig.input(std.mem.sliceAsBytes(self.x)); + const src = try self.rig.scratch(self.w.len * 2); + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = src, .len = self.w.len * 2 }).upload(0, std.mem.sliceAsBytes(self.w)); + const blocks = Bf16.bytes(experts, n, k); + const zeros = try a.alloc(u8, @max(blocks, self.pairs() * 4, self.items() * 12)); + defer a.free(zeros); + @memset(zeros, 0); + self.wd = try self.rig.input(zeros[0..blocks]); + try self.kx.pack(self.rig.s, src, .{ .stride = n * k, .ld = k, .in_out = false }, self.wd, experts, n, k); + const ids = try self.rig.input(std.mem.sliceAsBytes(self.picks)); + const hist = (self.pairs() + 1023) / 1024; + self.plan = .{ .members = try self.rig.input(zeros[0 .. self.pairs() * 4]), .items = try self.rig.input(zeros[0 .. self.items() * 12]), .counts = try self.rig.input(zeros[0..8]), .rank = try self.rig.scratch(self.pairs() * 4), .hist = try self.rig.scratch(hist * experts * 4) }; + try self.router.route(self.rig.s, ids, self.pairs(), experts, self.tile(), self.plan); + try self.rig.s.synchronize(); + self.out = try self.rig.output(self.pairs() * n * c.size(spec(self.entry).out)); + return true; + } + + /// Gate 4: mirrored integer weights and inputs, every partial sum an exact integer under 2^20, the row's sum small. + fn stress(self: *Experts, rng: std.Random) void { + const k = self.shape.k; + const h = k / 2; + for (0..experts * self.shape.n) |j| for (0..h) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 2)); + self.w[j * k + i] = c.toBf16(v); + self.w[j * k + h + i] = c.toBf16(v); + }; + for (0..self.xRows()) |r| for (0..h) |i| { + const v: f32 = @floatFromInt(4 * (1 + 2 * rng.uintLessThan(u32, 64))); + self.x[r * k + i] = c.toBf16(v); + self.x[r * k + h + i] = c.toBf16(if (i < 8) 4 - v else -v); + }; + } + + pub fn run(self: *Experts) !void { + const n = self.shape.n; + const k = self.shape.k; + const in: ex_.Rows = .{ .x = self.xd, .stride = k, .slots = if (self.up()) slots else 0 }; + switch (self.entry.fault) { + .none => { + const epi: ex_.Epi = switch (self.entry.kind) { + .up => .up, + .down_f32 => .down_f32, + .down_bf16 => .down_bf16, + }; + const l: ex_.Layer = .{ .format = .bf16, .up = self.wd, .down = self.wd, .width = if (self.up()) n else k, .dims = if (self.up()) k else n, .experts = experts, .relu2 = true }; + try self.kx.run(self.rig.s, self.entry.prompt, epi, in, l, self.plan, self.out, self.items(), -1); + }, + .short => try self.launch(self.kx.fns[0][1], k / 32 - 1, .{ .x = @intCast((self.items() * n / Bf16.cols + Bf16.warps - 1) / Bf16.warps) }, Bf16.warps * 32), + .bf16_sums => try self.launch(self.ctl.?, k / 32, .{ .x = @intCast((n + 127) / 128), .y = @intCast(self.items()) }, 128), + } + try self.rig.s.synchronize(); + } + + /// A down kernel's launch with its own input groups: experts_bf16.cu's arguments, less `slots` for the control. + fn launch(self: *Experts, f: cuda.Function, kg: usize, grid: struct { x: u32, y: u32 = 1 }, threads: u32) !void { + var a: cuda.Args = .{}; + a.add(self.xd); + a.add(@as(i32, @intCast(self.shape.k))); + if (self.entry.fault == .short) a.add(@as(i32, 0)); + a.add(self.wd); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(self.shape.n / Bf16.cols))); + for ([_]u64{ self.plan.items, self.plan.counts, self.plan.members, self.out }) |v| a.add(v); + a.add(@as(i32, @intCast(self.shape.n))); + if (self.entry.fault == .short) a.add(@as(i32, -1)); + try cuda.launch.launch(f, .{ .grid = .{ .x = grid.x, .y = grid.y, .z = 1 }, .block = .{ .x = threads } }, self.rig.s, &a); + } + + pub fn count(self: *Experts) usize { + return self.pairs() * self.shape.n; + } + + /// Pair p's products with column j of its expert, in input order. + fn terms(self: *const Experts, p: usize, j: usize, t: []f64) void { + const k = self.shape.k; + const xr = (if (self.up()) p / slots else p) * k; + const wr = (@as(usize, @intCast(self.picks[p])) * self.shape.n + j) * k; + for (t, 0..) |*o, i| o.* = c.fromBf16(self.x[xr + i]) * c.fromBf16(self.w[wr + i]); + } + + /// relu^2 for up, as float64 keeps a NaN; down's sum as it is. + fn epilogue(self: *const Experts, sum: f64) f64 { + if (!self.up()) return sum; + const u = if (std.math.isNan(sum)) sum else @max(sum, 0); + return u * u; + } + + pub fn reference(self: *Experts, want: []f64, mag: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + var sum: f64 = 0; + var abs: f64 = 0; + for (t) |v| { + sum += v; + abs += @abs(v); + } + const o = p * self.shape.n + j; + want[o] = self.epilogue(sum); + // a sum's error reaches relu^2 through its derivative 2 u + mag[o] = if (self.up()) 2 * @sqrt(want[o]) * abs else abs; + }; + } + + pub fn simulate(self: *Experts, acc: c.Accum, order: g.Order, out: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + out[p * self.shape.n + j] = self.epilogue(contract.rows.sumIn(acc, order, t)); + }; + } +}; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + kx: Bf16, + mod: cuda.Module, + ctl: cuda.Module, + items: [entries.len]Experts = undefined, + + const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; + }; + + /// The set, its kernels appended to `list` (each control after the kernels it breaks). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + errdefer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + errdefer mod.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .kx = kx, .mod = mod, .ctl = ctl }; + const router = try grouped.Router.resolve(set.mod); + const sums = try set.ctl.function("tf_ctl_bf16_experts_bf16_sums"); + for (&set.items, entries) |*x, e| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .kx = &set.kx, .router = router, .ctl = if (e.fault == .bf16_sums) sums else null }; + for (&set.items) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.items) |*x| x.deinit(); + set.ctl.unload(); + set.mod.unload(); + set.kx.unload(); + set.gpa.destroy(set); + } +}; + +/// The bf16 experts' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x62663136, &t); + std.debug.print("SUMMARY bf16 experts: {d} kernels, {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "bf16-experts contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} + +test "every control is named one, and only the controls plant a fault" { + for (entries) |e| try std.testing.expectEqual(e.fault != .none, std.mem.startsWith(u8, e.name, "control: ")); +} diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index 666968e5ba2..ed1bdc8d345 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -10,6 +10,7 @@ const Grouped = @import("contract_gb10_many.zig").Registrations; pub const Sets = struct { gpa: std.mem.Allocator, + bf16_experts: *@import("contract_experts_bf16.zig").Set, pocket: *@import("contract_pocket.zig").Set, minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, @@ -23,6 +24,8 @@ pub const Sets = struct { /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { + const bf16_experts = try @import("contract_experts_bf16.zig").Set.init(gpu, s, a, list); + errdefer bf16_experts.deinit(); const pocket = try @import("contract_pocket.zig").Set.init(gpu, s, a, list); errdefer pocket.deinit(); const minimax = try @import("contract_minimax.zig").Set.init(gpu, s, a, list); @@ -57,6 +60,7 @@ pub const Sets = struct { const joint_schema = try @import("contract_joint_schema.zig").Set.init(gpu, s, a, list); return .{ .gpa = gpu.gpa, + .bf16_experts = bf16_experts, .pocket = pocket, .minimax = minimax, .moss = moss, @@ -84,5 +88,6 @@ pub const Sets = struct { sets.moss.deinit(); sets.minimax.deinit(); sets.pocket.deinit(); + sets.bf16_experts.deinit(); } }; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 02f9ef56dc2..1f970b29532 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -30,7 +30,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple", "gemma4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "experts_bf16", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; return error.UnknownContractImage; } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index ebd3f880790..7c113ec4aec 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -215,7 +215,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes - \\ bf16-experts grouped bf16 experts against float64: repack, up, downs, rows and forms byte for byte + \\ bf16-experts [contract] grouped bf16 experts against float64 (repack, up, downs, rows, forms); contract: gates 1-4, controls red \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) @@ -431,7 +431,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); - if (std.mem.eql(u8, cmd, "bf16-experts")) return @import("bf16_experts_tests.zig").experts(gpu); + if (std.mem.eql(u8, cmd, "bf16-experts")) return if (rest.len == 0) @import("bf16_experts_tests.zig").experts(gpu) else if (std.mem.eql(u8, rest[0], "contract")) @import("contract_experts_bf16.zig").run(gpu) else error.UnknownCommand; if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); From 3033a7ccd112ba40afc046aa9ef2a29bff98713d Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:06:01 +0000 Subject: [PATCH 085/186] cuda experts: bf16 run refuses a width or dims off the 32-wide blocks, as pack does, with a test --- zig/src/cuda/experts_bf16.zig | 1 + zig/src/cuda/experts_host_test.zig | 15 +++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/zig/src/cuda/experts_bf16.zig b/zig/src/cuda/experts_bf16.zig index 76ed554c2f7..f585ba0875a 100644 --- a/zig/src/cuda/experts_bf16.zig +++ b/zig/src/cuda/experts_bf16.zig @@ -52,6 +52,7 @@ pub const Bf16Experts = struct { /// One call on a plan of at most `items` items: 16 pairs an item in decode, 64 in a prompt. pub fn run(e: *const Bf16Experts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { if (!l.relu2 or l.swiglu or l.limit != 0) return error.OptionNotBuilt; + if (l.width % 32 != 0 or l.dims % 32 != 0) return error.UnexpectedShape; // whole blocks, as pack takes const up = epi == .up; const kg = (if (up) l.dims else l.width) / 32; const nb = (if (up) l.width else l.dims) / cols; diff --git a/zig/src/cuda/experts_host_test.zig b/zig/src/cuda/experts_host_test.zig index 635eea49d2f..65f4a7f7952 100644 --- a/zig/src/cuda/experts_host_test.zig +++ b/zig/src/cuda/experts_host_test.zig @@ -38,3 +38,18 @@ test "the clamp-after order: exl3 takes it on both paths, every other format ref try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1)); }; } + +test "bf16 experts refuse dims or widths off the 32-wide blocks before any launch" { + const b: experts.Bf16Experts = undefined; // never reached: the refusal comes first + const g: Grouped = .{ .bf16 = &b }; + const s: Stream = undefined; + const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; + const in: experts.Rows = .{ .x = 0, .stride = 64 }; + for ([_][2]usize{ .{ 48, 64 }, .{ 64, 48 }, .{ 2688, 1880 } }) |wd| { + const l: Layer = .{ .format = .bf16, .up = 0, .down = 0, .width = wd[0], .dims = wd[1], .experts = 2, .relu2 = true }; + for ([_]experts.Epi{ .up, .down_f32, .down_bf16 }) |epi| { + try std.testing.expectError(error.UnexpectedShape, g.decode(s, epi, in, l, p, 0, 1, 1, 16, -1)); + try std.testing.expectError(error.UnexpectedShape, g.prompt(s, epi, in, l, p, 0, 1, 1, 64, -1)); + } + } +} From 294f28be2396007244cbb6589792411be70ea6df Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:54:19 +0000 Subject: [PATCH 086/186] tests: re-pin experts.zig and kernels.zig for #1630 #1630 adds the bf16 experts format to Grouped (experts.zig) and registers the experts_bf16 image (kernels.zig); no existing kernel source changed. Both rows are recomputed on main after #3011, which re-pinned experts.zig for #2843. --- zig/tests/fixtures/gradient-serving.tsv | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/tests/fixtures/gradient-serving.tsv b/zig/tests/fixtures/gradient-serving.tsv index 60b83c49883..3d3c09e7466 100644 --- a/zig/tests/fixtures/gradient-serving.tsv +++ b/zig/tests/fixtures/gradient-serving.tsv @@ -89,11 +89,11 @@ zig/src/cuda/cublaslt.zig c6e3b8892cad23bd08c7d9843da8117670e3a881e0708a89cdf72f zig/src/cuda/driver.zig f20856bd06afe9365ab9ca1bc39800ccb260ade2704197c033e1344a0c8dd742 raw commits e9bf05ffd, b923e554c, b2c75a98a; series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/embed_rows.zig 09fb2e612f43297c19d2cb77d45139dd98212a20075da0a15069a70544f72f6b raw zig/src/cuda/exlr8.zig 1e2164096ae949a8a9a4903df82c7c0c184a6f5e1ba5ba4423038a42b0633fc6 raw -zig/src/cuda/experts.zig 1450c0f7874f0cb5b9ae3a8bebf845f283214863b60cf401a26fdb1898ee7271 raw main 3768bc0168 (#2843: grouped EXL3 experts launched and sized at the caller's tile and pair count; no kernel source change) +zig/src/cuda/experts.zig 8bfc2f15a742cceba87dd7eb8e80f3f1b4e3f7eb323b557eb1bb4bf5a9fed440 raw main 3768bc0168 (#2843: grouped EXL3 experts launched and sized at the caller's tile and pair count; no kernel source change), #1630 (a bf16 format: its Grouped member, decode and prompt arms and prompt tile; other formats unchanged) zig/src/cuda/fp8.zig a9caf312cb451dc95572b493e69ae211dd1e8b5be5b2c48cae3febd5d9327143 raw series patches 65-pr-1024.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/graph.zig 17639aaef0298c934b3f5ad3511925169c096ae57ede8ebee655faa838844f7b raw zig/src/cuda/grouped.zig 5fa1f39394b701eab59383bc7ed214ea3e17652a1229236df49f5281e6e764b5 raw commits e1be0806d; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate -zig/src/cuda/kernels.zig db8a02c5b1469baa67e29b915b3e63cee0d0ec8b76009b9472b9307dadf12ebd raw main 42735f94f0 (#2335: the mla image registered), f35900e60c (#2709: the dots_mocr image registered) +zig/src/cuda/kernels.zig 34361df431dc6419474f5460158fd29bcba90f2a3792092b25308b0cb413fd69 raw main 42735f94f0 (#2335: the mla image registered), f35900e60c (#2709: the dots_mocr image registered), #1630 (the experts_bf16 image registered) zig/src/cuda/launch.zig 9ec64848834d82ddcea85a6123e975c95a5ae2bf85c2a71772244794b0bcd705 raw series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/memory.zig 7d2a91bfca8a08ae59df16ddf30d22226aba773cc814dd106930840ab344fc52 raw commits 5dcf77d7c, f630f1a90, e9bf05ffd, 881cccaa4, a34cb6ded; series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/modelopt.zig 648c9ba58197af0ce905f8792383502fc44a52c45d1413fe9e21d28d04c22253 raw From 62a9d5e17681d5e5dcaad2fc036aa41c879a442c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:04 +0000 Subject: [PATCH 087/186] changelog: #1630 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c3007ea6f1f..b66d1e88107 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. + - DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. - Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. From c22079b13f31e5cacfd96af78dba03e036a06b9f Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:56:28 +0000 Subject: [PATCH 088/186] glm: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says --- THIRD_PARTY_NOTICES.md | 3 +-- docs/recipes/glm-5.3-flash-cuda.md | 2 +- docs/recipes/glm-5.3-flash.md | 4 ++-- zig/src/families/glm/cuda_native.zig | 2 +- zig/src/families/glm/dflash.zig | 2 +- zig/src/families/glm/dflash_load.zig | 2 +- 6 files changed, 7 insertions(+), 8 deletions(-) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index c894ea02bd1..7a72107791f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -292,8 +292,7 @@ below); `tools/zig/bolt_ggml.py` reads them from that file. The IQ2_XXS and Q8_0 Bolt's GGML checkpoints (`zig/kernels/cuda/ggml.cu`, `zig/src/cuda/ggml.zig`) are TensorFold's own. TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card states Apache-2.0. -The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use -without derivatives; only the Python engine reads that drafter. Each checkpoint keeps its own license. +The optional `incoai/GLM-5.3-Flash-DFlash2` model card states Apache-2.0 too. Each checkpoint keeps its own license. ## Structured output: libtfgrammar diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 8d5efda9499..93f671bd8fe 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -162,7 +162,7 @@ concurrent one its run alone. | DFlash2 | Prose, 16 streams | 526.0 | 583.6 | One-stream runs are `tools/bench_openai.py` at 512 tokens and stream runs `tools/bench_concurrent.py` at 256. DFlash2's -weights are CC BY-NC-ND 4.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts +weights are Apache-2.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts are unchanged at about 4,430 tok/s from 8K to 131K tokens, and a repeated code reply reaches about 850 tok/s with windows of up to 32 rows. Other GPUs and rank counts are not covered here. diff --git a/docs/recipes/glm-5.3-flash.md b/docs/recipes/glm-5.3-flash.md index 538adf74bdc..f2e7a68c365 100644 --- a/docs/recipes/glm-5.3-flash.md +++ b/docs/recipes/glm-5.3-flash.md @@ -42,8 +42,8 @@ tensorfold serve TensorFold/GLM-5.3-Flash-MLX-4bit-MTP --tp 2 --rank 0 --master ``` Rank 1 starts first and rank 0 serves HTTP. Use the same context and drafting settings on both ranks. -The optional `incoai/GLM-5.3-Flash-DFlash2` model has CC BY-NC-ND 4.0 terms; pull it on both ranks only -when those terms fit the intended use. The CLI uses it automatically once it has been pulled. +The optional `incoai/GLM-5.3-Flash-DFlash2` model is Apache-2.0; pull it on both ranks to use it. +The CLI uses it automatically once it has been pulled. Without it, the engine uses MTP drafts; `--drafter none` explicitly selects MTP-only drafting. Give both ranks the same drafter setting. `--no-drafts` disables all drafting for the serial reference. A checkpoint with neither an MTP head nor a supplied DFlash2 model is refused unless drafts are disabled. diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 25efdd50087..904d7cb72aa 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -52,7 +52,7 @@ pub const Options = struct { mtp_confidence: ?f64 = null, // --mtp-confidence: rank 0's lanes stop an MTP chain where its chance falls under it learn_conversations: bool = false, // --learn-conversations: refused with TF_GLM_REQUEST_STARTS=1 (cuda_env) }; -/// --drafter: DFlash2 (incoai/GLM-5.3-Flash-DFlash2, non-commercial weights) drafts instead of the MTP head. +/// --drafter: DFlash2 (incoai/GLM-5.3-Flash-DFlash2, Apache-2.0 weights) drafts instead of the MTP head. pub const takes_drafter = true; /// The smallest room among the other ranks and whose it is. pub const Peer = struct { bytes: u64, rank: usize }; diff --git a/zig/src/families/glm/dflash.zig b/zig/src/families/glm/dflash.zig index 8fb578f3d45..aa039ffc299 100644 --- a/zig/src/families/glm/dflash.zig +++ b/zig/src/families/glm/dflash.zig @@ -1,4 +1,4 @@ -//! DFlash2 drafts (--drafter): incoai/GLM-5.3-Flash-DFlash2 (CC BY-NC-ND 4.0, read unmodified), after z-lab (MIT). +//! DFlash2 drafts (--drafter): incoai/GLM-5.3-Flash-DFlash2 (Apache-2.0, read unmodified), after z-lab (MIT). const std = @import("std"); const cuda = @import("cuda"); const mm = @import("cuda_mm.zig"); diff --git a/zig/src/families/glm/dflash_load.zig b/zig/src/families/glm/dflash_load.zig index 5726ec05fae..a2096579776 100644 --- a/zig/src/families/glm/dflash_load.zig +++ b/zig/src/families/glm/dflash_load.zig @@ -206,7 +206,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: anytype, dir: []const u8, ran try @import("dflash_load_owner.zig").acquire(gpa, dr, &tail, Tail, .{ .book = df.VOCAB * df.RANK, .ctx = df.MAX_CTX, .cand = candidates.host, .hidden = df.MAX_M * df.RANK }); dr.log = try @import("sys").environment.read(u64, "TF_GLM_DFLASH2_LOG", dr.log); if (std.c.getenv("TF_GLM_DFLASH2_PROF") != null) dr.prof = io; - std.log.info("dflash2: drafter from {s} (CC BY-NC-ND 4.0: non-commercial use only), linears at {d} bits, {d:.2} GiB on the device", .{ dir, if (bits == 0) 16 else bits, @as(f64, @floatFromInt(l.at)) / (1 << 30) }); + std.log.info("dflash2: drafter from {s} (Apache-2.0), linears at {d} bits, {d:.2} GiB on the device", .{ dir, if (bits == 0) 16 else bits, @as(f64, @floatFromInt(l.at)) / (1 << 30) }); return dr; } From a3aecd8d0fe5f07b6b67b46b3a1cca94a759d835 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:06 +0000 Subject: [PATCH 089/186] changelog: #3084 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b66d1e88107..8975e9de0af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). + - CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. - DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. From 6667aeef406265fe626c30610bb130887903f1ff Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 090/186] metal: dispatch generic prompt lanes as independent row blocks Give tf_lane the reg path's 16-row grid axis and runtime sums stride. Each block keeps the same group products, FMAs and K reduction; cooperative and older-chip encodes retain 16 rows. Qwen27 selects a whole-prompt key for non-reg widths behind the existing Ultra gate. Decode and mixed stacks stay short. Compare each width with main's frozen short kernel, including ragged tails, producer sums, strided outputs and three controls. Pin real quant dispatch counts with main-red host regressions. Refs #2818 --- zig/kernels/metal/core/lane_projection.metal | 17 +- zig/qwen27_quant_metal.zig | 17 +- zig/qwen27_quant_test.zig | 86 +++++ zig/src/core/lane_projection.zig | 23 +- zig/src/families/qwen27/quant_gpu.zig | 8 +- zig/tests/core_lane_long.zig | 153 ++++++++ zig/tests/core_lane_reg.zig | 13 +- zig/tests/fixtures/lane_prompt/README.md | 7 + .../lane_prompt/main_fb8fdf3_lane.metal | 356 ++++++++++++++++++ 9 files changed, 653 insertions(+), 27 deletions(-) create mode 100644 zig/tests/core_lane_long.zig create mode 100644 zig/tests/fixtures/lane_prompt/README.md create mode 100644 zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 2fedb0c13c2..05b55977cd1 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -141,12 +141,8 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #endif #if TF_PRECOMPUTE_SUMS -#if TF_BITS == 8 #define TF_SUM_VALUE(row, g) ((row) < M ? XS[(g) * mdims[1] + (row)] : 0.0f) #else -#define TF_SUM_VALUE(row, g) XS[(g) * 16 + (row)] -#endif -#else #define TF_SUM_VALUE(row, g) tf_lane_xsum(X, row, g, M) #endif @@ -178,7 +174,14 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in const device float* XS [[buffer(5)]], #endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], - uint tgx [[threadgroup_position_in_grid]]) { + uint2 tg [[threadgroup_position_in_grid]]) { + const int row0 = int(tg.y) * 16, rows = mdims[0] - row0; + if (rows <= 0) return; + X += (size_t)row0 * TF_K; + Y += (size_t)row0 * TF_YS; +#if TF_PRECOMPUTE_SUMS + XS += row0; +#endif const int lane = int(lanei), sg = int(sgi); const short qid = short(lane) >> 2; const short fm = (qid & 4) | ((short(lane) >> 1) & 3); @@ -187,8 +190,8 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #else const short fn = ((qid & 2) | (short(lane) & 1)) * 4; #endif - const int M = mdims[0]; - const int tile = tf_tile(int(tgx)); + const int M = min(rows, 16); + const int tile = tf_tile(int(tg.x)); const int n0 = tile * NT; const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); diff --git a/zig/qwen27_quant_metal.zig b/zig/qwen27_quant_metal.zig index 02a97ba26d5..9684fb4bba7 100644 --- a/zig/qwen27_quant_metal.zig +++ b/zig/qwen27_quant_metal.zig @@ -32,9 +32,11 @@ pub const Record = struct { released: usize = 0, live_bytes: usize = 0, fail_pipeline: bool = false, + tensor_units: bool = true, current: usize = 0, rows: u32 = 0, - launches: [256]struct { pipeline: usize, rows: u32 } = undefined, + launches: [256]struct { pipeline: usize, rows: u32, grid: Size, stride: i32 } = undefined, + sums_stride: i32 = 0, count: usize = 0, pub fn device(r: *Record) Device { return .{ .r = r }; @@ -65,8 +67,8 @@ pub const Buffer = struct { }; pub const Device = struct { r: *Record, - pub fn tensorUnits(_: Device) bool { - return true; + pub fn tensorUnits(d: Device) bool { + return d.r.tensor_units; } pub fn generation(_: Device) usize { return 17; @@ -111,10 +113,13 @@ pub const ComputeEncoder = struct { } pub fn setBuffer(_: ComputeEncoder, _: Buffer, _: usize, _: usize) void {} pub fn setValue(e: ComputeEncoder, v: anytype, at: usize) void { - if (@TypeOf(v) == [2]i32 and at == 3) e.r.rows = @intCast(v[0]); + if (@TypeOf(v) == [2]i32 and at == 3) { + e.r.rows = @intCast(v[0]); + e.r.sums_stride = v[1]; + } } - pub fn dispatchGroups(e: ComputeEncoder, _: Size, _: Size) void { - e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows }; + pub fn dispatchGroups(e: ComputeEncoder, grid: Size, _: Size) void { + e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows, .grid = grid, .stride = e.r.sums_stride }; e.r.count += 1; } pub fn dispatchThreads(_: ComputeEncoder, _: Size, _: Size) void {} diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 73d432c8705..1c151d8a718 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -132,3 +132,89 @@ test "qwen27 quant: preparation skips projections that cannot take a decode lane try t.expectEqual(@as(u32, 0), k.lanes.count()); try t.expectEqual(@as(usize, 0), r.libraries); } + +const wide_formats = [_]lane_projection.Format{ + .{ .bits = 8, .group = 64 }, .{ .bits = 6, .group = 32 }, .{ .bits = 5, .group = 64 }, + .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, +}; + +test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing source" { + for (wide_formats) |format| { + const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; + try long.validate(); + var short = long; + short.rows = 16; + const text = try lane_projection.source(a, long); + defer a.free(text); + const old = try lane_projection.source(a, short); + defer a.free(old); + try t.expectEqualStrings(old, text); + for ([_]usize{ 0, 15, 17, 2064 }) |rows| { + short.rows = rows; + try t.expectError(error.InvalidSchedule, short.validate()); + } + } +} + +test "qwen27 quant: each non-reg whole prompt is one projection dispatch" { + for (wide_formats) |format| for ([_]u32{ 17, 1065, 2048 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const producer: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 2080 } else null; + try k.quant(r.encoder(), p, ref, producer, ref, ref, rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(rows, r.launches[0].rows); + try t.expectEqual(@as(usize, (rows + 15) / 16), r.launches[0].grid.height); + try t.expectEqual(@as(i32, if (producer != null) 2080 else @intCast((rows + 15) / 16 * 16)), r.launches[0].stride); + var entries = k.lanes.valueIterator(); + const entry = entries.next().?; + try t.expectEqual(@as(u32, 2048), entry.max_rows); + if (entry.sums_buffer) |b| try t.expectEqual(@as(usize, 2048 * (5120 / format.group) * 4), b.length()); + // A decode or verify row still uses the short key and its original 16-row steps. + k.prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 17); + try t.expectEqual(@as(usize, 2), r.count); + try t.expectEqual(@as(u32, 16), r.launches[0].rows); + try t.expectEqual(@as(u32, 1), r.launches[1].rows); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectEqual(@as(u32, 2), k.lanes.count()); + k.prompt = true; + k.long_prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 2048); + try t.expectEqual(@as(usize, 128), r.count); + for (r.launches[0..r.count]) |call| { + try t.expectEqual(@as(u32, 16), call.rows); + try t.expectEqual(@as(usize, 1), call.grid.height); + } + }; +} + +test "qwen27 quant: cooperative and older-chip projections keep their 16-row limits" { + var r: mtl.Record = .{}; + var layout = lane_projection.Layout{ .n = 64, .k = 512, .format = .{ .bits = 4, .group = 64 }, .cooperative = true }; + const coop = try lane_projection.Projection.init(a, r.device(), layout); + defer coop.deinit(); + try t.expectEqual(@as(u32, 16), coop.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try coop.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, coop.encode(r.encoder(), ref, ref, ref, ref, 17)); + layout.rows = 2048; + try t.expectError(error.InvalidSchedule, layout.validate()); + layout.cooperative = false; + layout.format = .{ .bits = 6, .group = 32 }; + r.tensor_units = false; + const old = try lane_projection.Projection.init(a, r.device(), layout); + defer old.deinit(); + try t.expectEqual(@as(u32, 16), old.max_rows); + try t.expectError(error.InvalidRows, old.encode(r.encoder(), ref, ref, ref, ref, 17)); + const over_budget = lane_projection.Layout{ .n = 2048, .k = 5120, .format = .{ .bits = 8, .group = 128 }, .sk = 8 }; + try t.expectError(error.ThreadgroupMemoryLimit, over_budget.validate()); +} diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index a12f9d52f2b..a27141818d9 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -1,4 +1,4 @@ -//! Affine lane projections for 1-16 bf16 rows on Metal. No family or model dependency. +//! Affine lane projections in independent 16-row blocks on Metal. No family or model dependency. const std = @import("std"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -36,7 +36,7 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, - /// Rows one reg encode may take, `reg_rows` to `max_reg_rows`: it sizes the sums scratch, never the kernel source. + /// Rows one encode may take: it sizes scratch, never the kernel source; cooperative encodes still take 16. rows: usize = reg_rows, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, @@ -50,7 +50,9 @@ pub const Layout = struct { if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; - if (l.rows != reg_rows and (!l.reg or l.rows < reg_rows or l.rows % 16 != 0 or l.rows > max_reg_rows)) return error.InvalidSchedule; + if (l.cooperative) { + if (l.rows != reg_rows) return error.InvalidSchedule; + } else if (l.rows < @as(usize, if (l.reg) reg_rows else 16) or l.rows % 16 != 0 or l.rows > max_reg_rows) return error.InvalidSchedule; const g = l.groupCut(); if (g[1] == 0 or g[0] >= l.k / l.format.group or g[1] > l.k / l.format.group - g[0]) return error.InvalidGroupCut; if (l.groups != null and l.output != .f32) return error.PartialOutputNeedsFloat; @@ -238,12 +240,15 @@ pub const Projection = struct { sums_buffer: ?mtl.Buffer = null, threads_per_slice: usize = 32, reg: bool = false, - /// Rows one encode takes at most: the layout's rows for reg, 16 otherwise. + cooperative: bool = false, + /// Rows one encode takes at most: the layout's rows on tensor units, 16 on cooperative or older-chip paths. max_rows: u32 = reg_rows, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); - if (!device.tensorUnits() and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; + const units = device.tensorUnits(); + if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; + const capacity: usize = if (l.cooperative or !units) 16 else l.rows; const text = try source(a, l); defer a.free(text); const adapted = try frags.source(device, a, text); @@ -258,9 +263,9 @@ pub const Projection = struct { if (l.precompute_sums) { sums_pipe = try mtl.Pipeline.init(device, lib, "tf_lane_sums", false); errdefer sums_pipe.?.deinit(); - sums_buffer = try device.buffer(@as(usize, if (l.reg) l.rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + sums_buffer = try device.buffer(capacity * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .max_rows = @intCast(if (l.reg) l.rows else 16) }; + return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .cooperative = l.cooperative, .max_rows = @intCast(capacity) }; } pub fn deinit(p: Projection) void { @@ -306,7 +311,7 @@ pub const Projection = struct { e.setValue([2]i32{ @intCast(rows), @intCast(stride) }, 3); e.setBuffer(y.buf, y.off, 4); if (given) |g| e.setBuffer(g.ref.buf, g.ref.off, 5) else if (p.sums_buffer) |buffer| e.setBuffer(buffer, 0, 5); - e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.reg) blocks else 1, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); + e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.cooperative) 1 else blocks, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); if (given == null and p.sums_buffer != null) e.barrier(); } }; @@ -341,7 +346,7 @@ test "a reg layout's rows size its sums scratch, never its kernel source" { } var plain = base; plain.reg = false; - plain.rows = 256; + plain.rows = max_reg_rows + 16; try std.testing.expectError(error.InvalidSchedule, plain.validate()); } diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 12ae470de2e..be2149254ae 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -13,7 +13,7 @@ pub const Sums = struct { ref: Ref, stride: u32 }; pub const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 }, bits: u8 = 4, group: u32 = 64, long: bool = false }; pub const Kind = enum { projection, xsum, norm, input_norm, mlp, ane_input, mlp_ane, add_ane }; -/// Whether prompt chunks take one reg dispatch a projection (`Kernels.long_prompt`): Ultras with tensor units only. +/// Whether prompt chunks take one dispatch a projection (`Kernels.long_prompt`): Ultras with tensor units only. pub fn longPrompt(name: []const u8, tensor_units: bool) bool { return tensor_units and std.mem.endsWith(u8, name, " Ultra"); } @@ -56,7 +56,7 @@ pub const Kernels = struct { pair_twins: std.AutoHashMapUnmanaged([2]usize, bool) = .empty, /// Set around prompt chunks: prompt rows take projection.promptSplit, decoded rows the decode split. prompt: bool = false, - /// Set on `longPrompt` chips: prompt chunks past 128 rows take one reg dispatch a projection, on their own key. + /// Set on `longPrompt` chips: chunks past a short entry's capacity take one dispatch on a separate long key. long_prompt: bool = false, lanes: std.AutoHashMapUnmanaged(Key, lane.Projection) = .empty, cache: std.AutoHashMapUnmanaged(Key, mtl.Pipeline) = .empty, @@ -174,7 +174,7 @@ pub const Kernels = struct { .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg or (p.bits == 8 and p.group == 64), .reg = reg, - .rows = if (long) lane.max_reg_rows else lane.reg_rows, + .rows = if (long) lane.max_reg_rows else if (reg) lane.reg_rows else 16, }); k.joined(entry.value_ptr); } @@ -246,7 +246,7 @@ pub const Kernels = struct { } const reg = p.bits == 4 and p.group == 64; if (!reg and (skip > 0 or cut[1] > 0)) return error.BadCoreProjectionRows; - const long = reg and longKey(k.long_prompt, k.prompt, rows); + const long = if (reg) longKey(k.long_prompt, k.prompt, rows) else k.long_prompt and k.prompt and rows > 16; const sums: ?Sums = if (reg or (p.bits == 8 and p.group == 64)) xs else null; const entry = try k.laneEntry(p, sk, skip, cut, long); const step = entry.value_ptr.max_rows; diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig new file mode 100644 index 00000000000..502aa57f3da --- /dev/null +++ b/zig/tests/core_lane_long.zig @@ -0,0 +1,153 @@ +//! Whole prompt lanes at each width against main's frozen 16-row kernel, including strided oQ member formats. +const std = @import("std"); +const mtl = @import("metal"); +const lane = @import("core_lane"); +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const original = @embedFile("fixtures/lane_prompt/main_fb8fdf3_lane.metal"); + +fn bf16(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn pipe(device: mtl.Device, text: []const u8) !mtl.Pipeline { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + return mtl.Pipeline.init(device, lib, "tf_lane", false); +} + +fn frozen(a: std.mem.Allocator, d: mtl.Device, l: lane.Layout) !mtl.Pipeline { + const current = try lane.source(a, l); + defer a.free(current); + const end = std.mem.indexOf(u8, current, "#include ") orelse return error.SourceHeader; + const old = try std.mem.concat(a, u8, &.{ current[0..end], original }); + defer a.free(old); + return pipe(d, old); +} + +fn chunks(q: mtl.Queue, p: lane.Projection, x: mtl.Buffer, w: mtl.Buffer, sb: mtl.Buffer, xs: ?lane.Sums, y: mtl.Buffer, rows: u32) !void { + const cb = q.commandBuffer(); + const e = cb.compute(.serial); + var first: u32 = 0; + while (first < rows) : (first += p.max_rows) { + const xr = lane.Ref{ .buf = x, .off = @as(usize, first) * p.k * 2 }; + const yr = lane.Ref{ .buf = y, .off = @as(usize, first) * p.y_stride * 2 }; + const n = @min(rows - first, p.max_rows); + if (xs) |s| { + try p.encodeSums(e, xr, .{ .buf = w }, .{ .buf = sb }, .{ .ref = .{ .buf = s.ref.buf, .off = first * 4 }, .stride = s.stride }, yr, n); + } else try p.encode(e, xr, .{ .buf = w }, .{ .buf = sb }, yr, n); + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.GpuFailure; +} + +fn clear(y: mtl.Buffer, rows: usize, n: usize, stride: usize, candidate: bool) void { + @memset(y.contents()[0..y.length()], 0x5a); + if (candidate) for (0..rows) |m| @memset(y.contents()[m * stride * 2 ..][0 .. n * 2], 0xa5); +} + +fn same(a: mtl.Buffer, b: mtl.Buffer) bool { + return std.mem.eql(u8, a.contents()[0..a.length()], b.contents()[0..b.length()]); +} + +fn controls(a: std.mem.Allocator, d: mtl.Device, q: mtl.Queue, l: lane.Layout, p: lane.Projection, x: mtl.Buffer, w: mtl.Buffer, sb: mtl.Buffer, sums: lane.Sums, want: mtl.Buffer, got: mtl.Buffer) !void { + const text = try lane.source(a, l); + defer a.free(text); + const from = [_][]const u8{ "int(tg.y) * 16", "XS += row0;", "XS[(g) * mdims[1] + (row)]" }; + const to = [_][]const u8{ "0", "", "XS[(g) * 16 + (row)]" }; + for (from, to, 0..) |needle, replacement, i| { + const bad = try std.mem.replaceOwned(u8, a, text, needle, replacement); + defer a.free(bad); + if (std.mem.eql(u8, bad, text)) return error.ControlNotPlanted; + const bad_pipe = try pipe(d, bad); + defer bad_pipe.deinit(); + var broken = p; + broken.pipe = bad_pipe; + clear(got, 17, l.n, l.y_stride.?, true); + try chunks(q, broken, x, w, sb, sums, got, 17); + if (same(want, got)) return error.ControlMissed; + std.debug.print("LANE LONG CONTROL {d} caught\n", .{i}); + } +} + +pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: std.Random) !usize { + var count: usize = 0; + const formats = [_]lane.Format{ + .{ .bits = 8, .group = 64 }, .{ .bits = 6, .group = 32 }, .{ .bits = 5, .group = 64 }, + .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, + .{ .bits = 5, .group = 64 }, // the generic precomputed-sums layout, beside the normal inline-sums one + }; + for (formats, 0..) |format, fi| { + // Non-reg entries retain splitFor's loader split; only reg entries use promptSplit's smaller values. + for ([_][3]usize{ .{ 34816, 5120, 2 }, .{ 5120, 17408, 8 }, .{ 10240, 5120, 4 }, .{ 12288, 5120, 4 }, .{ 6240, 5120, 8 }, .{ 5120, 6144, 8 }, .{ 2048, 5120, 8 }, .{ 48, 512, 2 } }, 0..) |shape, si| { + const n = shape[0]; + const k = shape[1]; + const stride = n + @as(usize, if (si == 7) 32 else 0); + const most = lane.max_reg_rows; + const sum_stride = most + 16; + const ng = k / format.group; + // Group128 SK8 exceeds the unchanged admission bound; compare its supported SK4 variant on both sides. + const sk = if (format.group == 128 and shape[2] == 8) 4 else shape[2]; + const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; + const p_long = try lane.Projection.init(a, device, layout); + defer p_long.deinit(); + var short = layout; + short.rows = 16; + var p_short = try lane.Projection.init(a, device, short); + defer p_short.deinit(); + const old_pipe = try frozen(a, device, short); + p_short.pipe.deinit(); + p_short.pipe = old_pipe; + const raw = try a.alloc(u32, layout.sourceWords()); + defer a.free(raw); + for (raw) |*v| v.* = rng.int(u32); + const scales = try a.alloc(u16, n * ng); + defer a.free(scales); + const biases = try a.alloc(u16, n * ng); + defer a.free(biases); + for (scales, biases) |*s, *b| { + s.* = bf16(0.002 + rng.float(f32) * 0.02); + b.* = bf16(-0.01 - rng.float(f32) * 0.15); + } + const w = try device.buffer(layout.weightWords() * 4, opts); + defer w.deinit(); + const sb = try device.buffer(layout.metadataElements() * 2, opts); + defer sb.deinit(); + try lane.pack(layout, raw, scales, biases, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); + const x = try device.buffer(most * k * 2, opts); + defer x.deinit(); + const input = x.slice(u16, most * k); + for (input) |*v| v.* = bf16((rng.float(f32) - 0.5) * 4); + const given = try device.buffer(ng * sum_stride * 4, opts); + defer given.deinit(); + const sums = given.slice(f32, ng * sum_stride); + @memset(sums, -999); + for (0..ng) |g| for (0..most) |m| { + var acc: f32 = 0; + for (0..format.group) |j| acc += @as(f32, @bitCast(@as(u32, input[m * k + g * format.group + j]) << 16)); + sums[g * sum_stride + m] = acc; + }; + const want = try device.buffer(most * stride * 2, opts); + defer want.deinit(); + const got = try device.buffer(most * stride * 2, opts); + defer got.deinit(); + for ([_]u32{ 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { + if (external and fi != 0) continue; + const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; + clear(want, rows, n, stride, false); + clear(got, rows, n, stride, true); + try chunks(queue, p_short, x, w, sb, xs, want, rows); + try chunks(queue, p_long, x, w, sb, xs, got, rows); + if (!same(want, got)) { + std.debug.print("LANE LONG DIFFER bits={d} group={d} n={d} k={d} sk={d} rows={d} external={}\n", .{ format.bits, format.group, n, k, sk, rows, external }); + return error.LongLaneBytesDiffer; + } + if (fi == 0 and si == 0 and rows == 17 and external) try controls(a, device, queue, layout, p_long, x, w, sb, xs.?, want, got); + count += 1; + }; + } + } + return count; +} diff --git a/zig/tests/core_lane_reg.zig b/zig/tests/core_lane_reg.zig index eb33896366f..b00151b9f71 100644 --- a/zig/tests/core_lane_reg.zig +++ b/zig/tests/core_lane_reg.zig @@ -1,4 +1,4 @@ -//! The reg kernel, every output byte: 1-16 rows against the cooperative one, 129-2048 rows against 128-row dispatches. +//! Each lane width's output bytes: reg against cooperative, whole prompts against short dispatches. const std = @import("std"); const mtl = @import("metal"); const lane = @import("core_lane"); @@ -117,6 +117,15 @@ pub fn main(init: std.process.Init) !void { if (!device.tensorUnits()) return error.TensorUnitsRequired; const queue = try device.queue(); defer queue.deinit(); + const args = try init.minimal.args.toSlice(a); + const only_long = args.len == 2 and std.mem.eql(u8, args[1], "--lane-long-only"); + if (args.len > 1 and !only_long) return error.BadArguments; + if (only_long) { + var random = std.Random.DefaultPrng.init(0x7e9); + const checked = try @import("core_lane_long.zig").check(a, device, queue, random.random()); + std.debug.print("LANE LONG CHECK PASS cases={d}: every byte equals main's frozen 16-row kernel\n", .{checked}); + return; + } const order = try lane.RegOrder.init(device); defer order.deinit(); var prng = std.Random.DefaultPrng.init(0x7e9); @@ -187,4 +196,6 @@ pub fn main(init: std.process.Init) !void { std.debug.print("REG CHECK PASS cases={d}: every output byte equals the cooperative kernel's\n", .{cases}); const long = try longCheck(a, device, queue, order, rng); std.debug.print("LONG CHECK PASS cases={d}: one dispatch of up to {d} rows equals {d}-row dispatches byte for byte\n", .{ long, lane.max_reg_rows, lane.reg_rows }); + const wide = try @import("core_lane_long.zig").check(a, device, queue, rng); + std.debug.print("LANE LONG CHECK PASS cases={d}: every byte equals main's frozen 16-row kernel\n", .{wide}); } diff --git a/zig/tests/fixtures/lane_prompt/README.md b/zig/tests/fixtures/lane_prompt/README.md new file mode 100644 index 00000000000..0ce2210fe66 --- /dev/null +++ b/zig/tests/fixtures/lane_prompt/README.md @@ -0,0 +1,7 @@ +The frozen Metal file is TensorFold's own lane kernel from main +fb8fdf3d8501255058a3221452164c56cb9a57ca. It is copied without modification from +zig/kernels/metal/core/lane_projection.metal and retains the repository's Apache-2.0 license. + +core_lane_long.zig compiles this kernel with each tested shape's defines and uses +16-row dispatches as the byte reference for the current whole-prompt kernel. +The fixture must stay frozen when the production kernel changes. diff --git a/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal b/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal new file mode 100644 index 00000000000..2fedb0c13c2 --- /dev/null +++ b/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal @@ -0,0 +1,356 @@ +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +// 1-16 bf16 rows share each group's weight read; every row's sums and split-K order hold at every width. + +#ifndef TF_SIMD_FRAGS +#define TF_SIMD_FRAGS 0 +#endif + +#ifndef TF_PRECOMPUTE_SUMS +#define TF_PRECOMPUTE_SUMS 0 +#endif +#ifndef TF_YS +#define TF_YS TF_N +#endif +// Output rows lie TF_YS elements apart, so a member of a fused stack writes its columns of the fused rows. +#define TF_ODD_WORDS (TF_GROUP * TF_BITS / 32 % 2) + +constant constexpr int NT = 32, GS = TF_GROUP, KG = TF_K / GS, NF = NT / 16; +constant constexpr int WP = GS * TF_BITS / 32; + +#if TF_ODD_WORDS +struct tf_raw { uint w[WP]; uint4 s[NF]; }; // 3- and 5-bit groups of 32 are an odd word count, loaded a word at a time +#else +struct tf_raw { uint2 w[WP / 2]; uint4 s[NF]; }; // a lane's packed code words and its scale/bias quads for one group +#endif + +inline tf_raw tf_lane_fetch(const device uint* W, const device uint4* SB, int tile, int g, int g_end, int n0, int lane, + short fn) { + tf_raw raw; + const bool ok = g < g_end; + const device uint2* src = (const device uint2*)(W + (((size_t)tile * KG + (ok ? g : 0)) * NT + lane) * WP); +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) raw.w[i] = ok && n0 + lane < TF_N ? ((const device uint*)src)[i] : 0u; +#elif TF_BITS != 4 && TF_BITS != 8 + for (int i = 0; i < WP / 2; i++) raw.w[i] = ok && n0 + lane < TF_N ? src[i] : uint2(0); +#else + for (int i = 0; i < WP / 2; i++) raw.w[i] = uint2(0); +#endif +#if TF_SIMD_FRAGS + const device uint2* sb2 = (const device uint2*)SB; + for (int f = 0; f < NF; f++) { + const int col = n0 + f * 16 + fn; + raw.s[f] = uint4(ok && col < TF_N ? sb2[((size_t)g * TF_N + col) / 2] : uint2(0), + ok && col + 8 < TF_N ? sb2[((size_t)g * TF_N + col + 8) / 2] : uint2(0)); + } +#else + // Four lanes load the metadata quads and broadcast them; all 32 lanes fetch, even for inactive rows. + const ushort leader = ushort((fn & 8) | ((fn >> 2) & 1)); + for (int f = 0; f < NF; f++) { + const uint4 sb = lane == int(leader) && ok && n0 + f * 16 + fn < TF_N + ? SB[((size_t)(ok ? g : 0) * TF_N + n0 + f * 16 + fn) / 4] : uint4(0); + raw.s[f] = simd_shuffle(sb, leader); + } +#endif + return raw; +} + +inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { + float acc = 0.0f; + if (m < M) for (int i = 0; i < GS; i++) acc += float(X[(size_t)m * TF_K + g * GS + i]); + return acc; +} + +// Codes: 4/8-bit device views; 2/3-bit widened to nibbles, 5/6-bit to bytes. +#if TF_BITS == 4 +#define TF_STAGE_CODES(raw) +#define TF_GROUP_PRODUCT(g) \ + tensor, tensor_inline> b( \ + (device uchar*)W + ((size_t)tile * KG + (g)) * (NT * GS / 2), dextents(GS, NT)); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); +#elif TF_BITS == 8 +#define TF_STAGE_CODES(raw) +#define TF_GROUP_PRODUCT(g) \ + tensor, tensor_inline> b( \ + (device uint8_t*)W + ((size_t)tile * KG + (g)) * (NT * GS), dextents(GS, NT)); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); +#else +#if TF_BITS < 4 +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; +#if TF_BITS == 3 + for (int c = 0; c < GS / 8; c++) { + const int bit = 24 * c, i = bit >> 5, sh = bit & 31; + uint pack = words[i] >> sh; + if (sh > 8) pack |= words[i + 1] << (32 - sh); + uint nib = 0u; + for (int j = 0; j < 8; j++) nib |= ((pack >> (3 * j)) & 7u) << (4 * j); + stage[lane * (GS / 8) + c] = nib; + } +#else + for (int c = 0; c < GS / 8; c++) { // word c/2's half c%2: 8 codes of 2 bits to 8 nibbles + uint v = (words[c >> 1] >> (16 * (c & 1))) & 0xFFFFu; + v = (v | (v << 8)) & 0x00FF00FFu; + v = (v | (v << 4)) & 0x0F0F0F0Fu; + v = (v | (v << 2)) & 0x33333333u; + stage[lane * (GS / 8) + c] = v; + } +#endif +} +#else +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; + for (int c = 0; c < GS / 4; c++) { + uint word = 0u; + for (int j = 0; j < 4; j++) { + const int bit = (4*c + j) * TF_BITS, i = bit >> 5, sh = bit & 31; + uint code = words[i] >> sh; + if (sh + TF_BITS > 32) code |= words[i+1] << (32-sh); + word |= (code & ((1u << TF_BITS)-1u)) << (8*j); + } + stage[lane * (GS / 4) + c] = word; + } +} +#endif +#define TF_STAGE_CODES(raw) tf_stage_codes(stage, raw, lane); +#define TF_GROUP_PRODUCT(g) \ + simdgroup_barrier(mem_flags::mem_threadgroup); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); \ + simdgroup_barrier(mem_flags::mem_threadgroup); +#endif + +#if TF_PRECOMPUTE_SUMS +#if TF_BITS == 8 +#define TF_SUM_VALUE(row, g) ((row) < M ? XS[(g) * mdims[1] + (row)] : 0.0f) +#else +#define TF_SUM_VALUE(row, g) XS[(g) * 16 + (row)] +#endif +#else +#define TF_SUM_VALUE(row, g) tf_lane_xsum(X, row, g, M) +#endif + +// Each group loads packed codes or stages widened codes, performs the tensor product, then applies the scale/bias FMAs. +#define TF_LANE_GROUP(g, raw) \ + { \ + TF_STAGE_CODES(raw) \ + float s[NF][4], bb[NF][4]; \ + for (int f = 0; f < NF; f++) { \ + const vec v = as_type>(raw.s[f]); \ + for (int j = 0; j < 4; j++) { s[f][j] = float(v[2 * j]); bb[f][j] = float(v[2 * j + 1]); } \ + } \ + raw = tf_lane_fetch(W, sbv, tile, (g) + TF_PF, g_end, n0, lane, fn); \ + TF_GROUP_PRODUCT(g) \ + const float xs0 = TF_SUM_VALUE(fm, g); \ + const float xs1 = TF_SUM_VALUE(fm + 8, g); \ + for (int f = 0; f < NF; f++) \ + for (int r = 0; r < 2; r++) \ + for (int j = 0; j < 4; j++) { \ + const int i = f * 8 + r * 4 + j; \ + C[i] = fma(s[f][j], P[TF_SIMD_FRAGS ? r * 8 + f * 4 + j : i], \ + fma(bb[f][j], r ? xs1 : xs0, C[i])); \ + } \ + } + +[[kernel]] void tf_lane(const device bfloat* X [[buffer(0)]], const device uint* W [[buffer(1)]], + const device bfloat* SBt [[buffer(2)]], const device int* mdims [[buffer(3)]], device TF_OUT* Y [[buffer(4)]], +#if TF_PRECOMPUTE_SUMS + const device float* XS [[buffer(5)]], +#endif + uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], + uint tgx [[threadgroup_position_in_grid]]) { + const int lane = int(lanei), sg = int(sgi); + const short qid = short(lane) >> 2; + const short fm = (qid & 4) | ((short(lane) >> 1) & 3); +#if TF_SIMD_FRAGS + const short fn = ((short(lane) & 8) >> 1) | ((short(lane) & 1) << 1); +#else + const short fn = ((qid & 2) | (short(lane) & 1)) * 4; +#endif + const int M = mdims[0]; + const int tile = tf_tile(int(tgx)); + const int n0 = tile * NT; + const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; + constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); + matmul2d op; + tensor, tensor_inline> tA((device bfloat*)X, dextents(TF_K, M)); +#if TF_BITS < 4 + threadgroup uint stage_all[TF_SK * NT * (GS / 8)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 8); + tensor, tensor_inline> b((threadgroup uchar*)stage, dextents(GS, NT)); +#elif TF_BITS != 4 && TF_BITS != 8 + threadgroup uint stage_all[TF_SK * NT * (GS / 4)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 4); + tensor, tensor_inline> b((threadgroup uint8_t*)stage, dextents(GS, NT)); +#endif + const device uint4* sbv = (const device uint4*)SBt; + float C[NF * 8]; + for (int i = 0; i < NF * 8; i++) C[i] = 0.0f; + tf_raw r0 = tf_lane_fetch(W, sbv, tile, g_begin, g_end, n0, lane, fn); +#if TF_PF == 2 + tf_raw r1 = tf_lane_fetch(W, sbv, tile, g_begin + 1, g_end, n0, lane, fn); + for (int g = g_begin; g < g_end; g += 2) { + TF_LANE_GROUP(g, r0) + if (g + 1 < g_end) TF_LANE_GROUP(g + 1, r1) + } +#else + for (int g = g_begin; g < g_end; g++) TF_LANE_GROUP(g, r0) +#endif + threadgroup float part[(TF_SK > 1 ? TF_SK - 1 : 1) * NF * 8 * 32]; + if (TF_SK > 1) { + if (sg > 0) for (int i = 0; i < NF * 8; i++) part[((sg - 1) * NF * 8 + i) * 32 + lane] = C[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) + for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < NF * 8; i++) C[i] += part[((s2 - 1) * NF * 8 + i) * 32 + lane]; + } + if (sg == 0) + for (int f = 0; f < NF; f++) + for (int r = 0; r < 2; r++) { + const int m = fm + 8 * r; + const int nn = n0 + f * 16 + fn; + if (m < M && nn < TF_N) + for (int j = 0; j < 4; j++) { + const int col = nn + (TF_SIMD_FRAGS ? (j & 1) + 8 * (j >> 1) : j); + if (col < TF_N) Y[(size_t)m * TF_YS + col] = static_cast(C[f * 8 + r * 4 + j]); + } + } +} + +#if TF_PRECOMPUTE_SUMS +kernel void tf_lane_sums(const device bfloat* X [[buffer(0)]],constant int& rows [[buffer(1)]],device float* XS [[buffer(2)]],constant int& stride [[buffer(3)]],uint index [[thread_position_in_grid]]) { + const int row=int(index)%stride,group=int(index)/stride; + if(group>=KG)return; + XS[index]=tf_lane_xsum(X,row,group,rows); +} +#endif + +#if TF_BITS == 4 +kernel void tf_lane_pair(const device bfloat* X [[buffer(0)]],const device uint* W [[buffer(1)]],const device bfloat* SB [[buffer(2)]],const device int* mdims [[buffer(3)]],device TF_OUT* Y [[buffer(4)]], +#if TF_PRECOMPUTE_SUMS +const device float* XS [[buffer(5)]], +#endif +uint group_id [[simdgroup_index_in_threadgroup]],uint lane [[thread_index_in_simdgroup]],uint thread_id [[thread_position_in_threadgroup]],uint tile_id [[threadgroup_position_in_grid]]) { + const uint sg=group_id/2,tip=thread_id-sg*64; + const int M=mdims[0],tile=tf_tile(int(tile_id)),n0=tile*NT; + const int first=TF_G0+(int(sg)*TF_GN)/TF_SK,last=TF_G0+((int(sg)+1)*TF_GN)/TF_SK; + constexpr auto descriptor=matmul2d_descriptor(16,NT,GS,false,true,false,matmul2d_descriptor::mode::multiply); + matmul2d> op; + tensor,tensor_inline> a((device bfloat*)X,dextents(TF_K,M)); + tensor,tensor_inline> b0((device uchar*)W,dextents(GS,NT)); + auto a0=a.slice(0,0); + auto p=op.template get_destination_cooperative_tensor(); + short column[8],row[8];float value[8]; + for(int i=0;i<8;i++){auto index=p.get_multidimensional_index(i);column[i]=index[0];row[i]=index[1];value[i]=0.0f;} + for(int group=first;group,tensor_inline> b((device uchar*)W+((size_t)tile*KG+group)*(NT*GS/2),dextents(GS,NT)); + auto x=a.slice(group*GS,0);op.run(x,b,p); + for(int i=0;i<8;i++) { + const int n=n0+column[i]; + const float scale=float(SB[(size_t(group)*TF_N+n)*2]),bias=float(SB[(size_t(group)*TF_N+n)*2+1]); +#if TF_PRECOMPUTE_SUMS + const float sum=XS[group*16+row[i]]; +#else + const float sum=tf_lane_xsum(X,row[i],group,M); +#endif + value[i]=fma(scale,p[i],fma(bias,sum,value[i])); + } + } + threadgroup float partial[(TF_SK>1?TF_SK-1:1)*8*64]; + if(TF_SK>1) { + if(sg>0) for(int i=0;i<8;i++) partial[((sg-1)*8+i)*64+tip]=value[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if(sg==0) for(int part=1;part> 3) & 1), mb = ((lane >> 1) & 3) + 4 * ((lane >> 4) & 1); + const int first = TF_G0 + (ks * TF_GN) / TF_SK, last = TF_G0 + ((ks + 1) * TF_GN) / TF_SK; + constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); + matmul2d op; + tensor, tensor_inline> a((device bfloat*)X, dextents(TF_K, M)); + auto a0 = a.slice(0, 0); + auto cb = op.template get_right_input_cooperative_tensor(); + auto P = op.template get_destination_cooperative_tensor, metal::remove_addrspace_t, float>(); + thread half2* b2 = (thread half2*)&cb[0]; + const device uint4* wl = W + (size_t)tile * KG * 64 + lane; + const device uint4* sbv = (const device uint4*)SB; + const half2 bias = half2(1024.0h); + float C[16]; + for (int i = 0; i < 16; i++) C[i] = 0.0f; + tf_reg_load next = {}; + if (first < last) next = tf_reg_fetch(wl, sbv, XS, first, n0 + nb, mb, stride); + for (int g = first; g < last; g++) { + const tf_reg_load cur = next; + for (int w = 0; w < 8; w++) { + // 1024 + code as fp16 bits, codes (0,4) (1,5) (2,6) (3,7) of a word; subtracting 1024 is exact + const uint word = w < 4 ? cur.w0[w] : cur.w1[w - 4], lo = word & 0x0F0F0F0Fu, hi = (word >> 4) & 0x0F0F0F0Fu; + const int e = 16 * (w & 1) + 2 * (w >> 1); + b2[e] = as_type((lo & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 1] = as_type((hi & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 8] = as_type(((lo >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 9] = as_type(((hi >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; + } + if (g + 1 < last) next = tf_reg_fetch(wl, sbv, XS, g + 1, n0 + nb, mb, stride); + auto x = a.slice(g * GS, 0); + op.run(x, cb, P); + for (int j = 0; j < 4; j++) { + const float s0 = as_type(cur.q0[j] << 16), b0 = as_type(cur.q0[j] & 0xFFFF0000u); + const float s1 = as_type(cur.q1[j] << 16), b1 = as_type(cur.q1[j] & 0xFFFF0000u); + C[j] = fma(s0, P[j], fma(b0, cur.x0, C[j])); + C[4 + j] = fma(s0, P[4 + j], fma(b0, cur.x1, C[4 + j])); + C[8 + j] = fma(s1, P[8 + j], fma(b1, cur.x0, C[8 + j])); + C[12 + j] = fma(s1, P[12 + j], fma(b1, cur.x1, C[12 + j])); + } + } + threadgroup float part[(TF_SK > 1 ? TF_SK - 1 : 1) * 16 * 32]; + if (TF_SK > 1) { + if (ks > 0) for (int i = 0; i < 16; i++) part[((ks - 1) * 16 + i) * 32 + lane] = C[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0) for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < 16; i++) C[i] += part[((s2 - 1) * 16 + i) * 32 + lane]; + } + if (ks == 0) + for (int i = 0; i < 16; i++) { + const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); + if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); + } +} +#endif From 8c4b12bf6f51a56f483cd8c10f4f4d05df28f365 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 091/186] metal: record measured offline limits for standalone byte-lane sums The full lane contract compiles all four tf_lane_sums sources for M1 and M2 at a 1024-thread pipeline limit. Record those measured values and keep the actual 256-thread launch. Tensor-op sources remain explicitly unmeasured on these chips. Pin the distinction in a host regression; no shader changes. Refs #2818 --- zig/src/core/contract/lane_q8_records.zig | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/zig/src/core/contract/lane_q8_records.zig b/zig/src/core/contract/lane_q8_records.zig index 16bdd8bd2c1..0e734861227 100644 --- a/zig/src/core/contract/lane_q8_records.zig +++ b/zig/src/core/contract/lane_q8_records.zig @@ -1,4 +1,4 @@ -//! M5 byte-lane sources retain their launch envelopes before and after removing code staging. +//! Byte-lane launch envelopes; the standalone sums compiled for M1/M2 even when tensor ops cannot. const M = @import("limits.zig").MetalLaunch; pub const Bad = enum { none, group, stride }; pub const Source = struct { name: []const u8, n: usize, k: usize, sk: usize, group: usize = 64, sums: bool = false, pf: usize = 1, f32_out: bool = false, bad: Bad = .none }; @@ -35,7 +35,24 @@ pub const records = blk: { const groups: usize = @max(@as(usize, s.sk) -| 1, @as(usize, 1)); const threads: usize = if (sum) 256 else 32 * s.sk; const memory: usize = if (sum) 0 else groups * 16 * 32 * 4; - out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .offline = false }; + out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .m1 = if (sum) 1024 else 0, .m2 = if (sum) 1024 else 0, .offline = sum }; } break :blk out; }; + +test "byte lane records: standalone sums have measured offline limits; tensor ops stay unmeasured" { + const std = @import("std"); + var count: usize = 0; + for (records) |r| { + const sum = std.mem.endsWith(u8, r.name, "tf_lane_sums"); + try std.testing.expectEqual(sum, r.offline); + try std.testing.expectEqual(@as(u32, if (sum) 1024 else 0), r.m1); + try std.testing.expectEqual(r.m1, r.m2); + if (sum) { + count += 1; + try std.testing.expectEqual(@as(u32, 256), r.threads); + try std.testing.expectEqual(@as(u32, 0), r.tg_memory); + } + } + try std.testing.expectEqual(@as(usize, 4), count); +} From 79fc3d4b3abbbf4e48c0fbae128505eb81fcb472 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 092/186] metal: keep short byte-lane projections on a scalar row grid Only whole-prompt layouts compile the independent row-block header. Short layouts retain a scalar group coordinate and their original row extent; guarded sums remain bounds-safe. Keep older-chip sources on that short variant as well. Compare both short and long programs against frozen main at every maintained byte-check shape and row count. Refs #2818 --- zig/kernels/metal/core/lane_projection.metal | 17 ++++++++++++++++- zig/qwen27_quant_test.zig | 5 +++-- zig/src/core/lane_projection.zig | 5 ++++- zig/tests/core_lane_long.zig | 13 +++++++++---- 4 files changed, 32 insertions(+), 8 deletions(-) diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 05b55977cd1..7e65ddb08b6 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -12,6 +12,9 @@ using namespace mpp::tensor_ops; #ifndef TF_PRECOMPUTE_SUMS #define TF_PRECOMPUTE_SUMS 0 #endif +#ifndef TF_ROW_BLOCKS +#define TF_ROW_BLOCKS 0 +#endif #ifndef TF_YS #define TF_YS TF_N #endif @@ -174,13 +177,20 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in const device float* XS [[buffer(5)]], #endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], - uint2 tg [[threadgroup_position_in_grid]]) { +#if TF_ROW_BLOCKS + uint2 tg [[threadgroup_position_in_grid]] +#else + uint tgx [[threadgroup_position_in_grid]] +#endif +) { +#if TF_ROW_BLOCKS const int row0 = int(tg.y) * 16, rows = mdims[0] - row0; if (rows <= 0) return; X += (size_t)row0 * TF_K; Y += (size_t)row0 * TF_YS; #if TF_PRECOMPUTE_SUMS XS += row0; +#endif #endif const int lane = int(lanei), sg = int(sgi); const short qid = short(lane) >> 2; @@ -190,8 +200,13 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #else const short fn = ((qid & 2) | (short(lane) & 1)) * 4; #endif +#if TF_ROW_BLOCKS const int M = min(rows, 16); const int tile = tf_tile(int(tg.x)); +#else + const int M = mdims[0]; + const int tile = tf_tile(int(tgx)); +#endif const int n0 = tile * NT; const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 1c151d8a718..783ec7c2eef 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -138,7 +138,7 @@ const wide_formats = [_]lane_projection.Format{ .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, }; -test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing source" { +test "qwen27 quant: non-reg layouts admit whole prompt chunks with row-block source variants" { for (wide_formats) |format| { const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; try long.validate(); @@ -148,7 +148,8 @@ test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing s defer a.free(text); const old = try lane_projection.source(a, short); defer a.free(old); - try t.expectEqualStrings(old, text); + try t.expect(std.mem.indexOf(u8, old, "#define TF_ROW_BLOCKS 0\n") != null); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 1\n") != null); for ([_]usize{ 0, 15, 17, 2064 }) |rows| { short.rows = rows; try t.expectError(error.InvalidSchedule, short.validate()); diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index a27141818d9..ca7adc4121c 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -120,6 +120,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { } try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); + try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(!l.reg and !l.cooperative and l.rows > 16)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); @@ -249,7 +250,9 @@ pub const Projection = struct { const units = device.tensorUnits(); if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; const capacity: usize = if (l.cooperative or !units) 16 else l.rows; - const text = try source(a, l); + var kernel_layout = l; + if (!units and !l.cooperative and !l.reg) kernel_layout.rows = capacity; + const text = try source(a, kernel_layout); defer a.free(text); const adapted = try frags.source(device, a, text); defer a.free(adapted); diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig index 502aa57f3da..6add4d36ffe 100644 --- a/zig/tests/core_lane_long.zig +++ b/zig/tests/core_lane_long.zig @@ -95,11 +95,12 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st defer p_long.deinit(); var short = layout; short.rows = 16; - var p_short = try lane.Projection.init(a, device, short); + const p_short = try lane.Projection.init(a, device, short); defer p_short.deinit(); const old_pipe = try frozen(a, device, short); - p_short.pipe.deinit(); - p_short.pipe = old_pipe; + defer old_pipe.deinit(); + var p_old = p_short; + p_old.pipe = old_pipe; const raw = try a.alloc(u32, layout.sourceWords()); defer a.free(raw); for (raw) |*v| v.* = rng.int(u32); @@ -138,7 +139,10 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; clear(want, rows, n, stride, false); clear(got, rows, n, stride, true); - try chunks(queue, p_short, x, w, sb, xs, want, rows); + try chunks(queue, p_old, x, w, sb, xs, want, rows); + try chunks(queue, p_short, x, w, sb, xs, got, rows); + if (!same(want, got)) return error.ShortLaneBytesDiffer; + clear(got, rows, n, stride, true); try chunks(queue, p_long, x, w, sb, xs, got, rows); if (!same(want, got)) { std.debug.print("LANE LONG DIFFER bits={d} group={d} n={d} k={d} sk={d} rows={d} external={}\n", .{ format.bits, format.group, n, k, sk, rows, external }); @@ -149,5 +153,6 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st }; } } + std.debug.print("LANE SHORT CHECK PASS cases={d}: current short variant equals frozen main\n", .{count}); return count; } From 9359fe73c18971005cef40e2cb35244282b3b4b8 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Tue, 10 Nov 2026 23:13:35 -0800 Subject: [PATCH 093/186] metal: opt into whole-prompt row blocks explicitly Leave other affine lane callers on their scalar 16-row source and capacity. Only the Qwen prompt key and byte fixture request independent row blocks. Pin the shared shader bytes, describe the row layout accurately and cover the 16/17-row boundary. Refs #2818 --- zig/gemma4_shared_affine_test.zig | 4 +-- zig/qwen27_quant_test.zig | 39 ++++++++++++++++++++++++++- zig/src/core/lane_projection.zig | 10 ++++--- zig/src/families/qwen27/quant_gpu.zig | 1 + zig/tests/core_lane_long.zig | 6 ++--- 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/zig/gemma4_shared_affine_test.zig b/zig/gemma4_shared_affine_test.zig index 33e99b67ad7..4ab73b90968 100644 --- a/zig/gemma4_shared_affine_test.zig +++ b/zig/gemma4_shared_affine_test.zig @@ -109,8 +109,8 @@ test "landed Metal source bytes are pinned without shader compilation" { { var hash: [32]u8 = undefined; std.crypto.hash.sha2.Sha256.hash(ks.core_lane_projection, &hash, .{}); - // 66155d5ee's bytes (8-bit lane codes and producer sums, Refs #2502, under the metal_lane_q8 contract) - try std.testing.expectEqualStrings("159f39224b5bb60fd2ced301496d91662d2958056bd291d13d8a5db5186dd41a", &std.fmt.bytesToHex(hash, .lower)); + // Packed affine lanes with producer sums and independent prompt row blocks. + try std.testing.expectEqualStrings("7befe2028dae5036541f311f230a02d5652e634450ef16dc66dce6e15e935277", &std.fmt.bytesToHex(hash, .lower)); } } test "released lane parameters and conservative scratch size" { diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 783ec7c2eef..8fcfa9a5bcf 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -140,7 +140,7 @@ const wide_formats = [_]lane_projection.Format{ test "qwen27 quant: non-reg layouts admit whole prompt chunks with row-block source variants" { for (wide_formats) |format| { - const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; + const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048, .row_blocks = true }; try long.validate(); var short = long; short.rows = 16; @@ -219,3 +219,40 @@ test "qwen27 quant: cooperative and older-chip projections keep their 16-row lim const over_budget = lane_projection.Layout{ .n = 2048, .k = 5120, .format = .{ .bits = 8, .group = 128 }, .sk = 8 }; try t.expectError(error.ThreadgroupMemoryLimit, over_budget.validate()); } + +test "qwen27 quant: default non-reg callers retain scalar source and 16-row capacity" { + for (wide_formats) |format| { + var r: mtl.Record = .{}; + const layout = lane_projection.Layout{ .n = 64, .k = 5120, .format = format }; + const text = try lane_projection.source(a, layout); + defer a.free(text); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 0\n") != null); + const p = try lane_projection.Projection.init(a, r.device(), layout); + defer p.deinit(); + try t.expectEqual(@as(u32, 16), p.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try p.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, p.encode(r.encoder(), ref, ref, ref, ref, 17)); + } +} + +test "qwen27 quant: non-reg prompt key opts in only above 16 rows" { + for (wide_formats) |format| for ([_]u32{ 16, 17 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const sums: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 32 } else null; + try k.quant(r.encoder(), p, ref, sums, ref, ref, rows); + var entries = k.lanes.iterator(); + const entry = entries.next().?; + try t.expectEqual(rows > 16, entry.key_ptr.long); + try t.expectEqual(@as(u32, if (rows > 16) 2048 else 16), entry.value_ptr.max_rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(@as(usize, if (rows > 16) 2 else 1), r.launches[0].grid.height); + }; +} diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index ca7adc4121c..14e6029d370 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -36,8 +36,10 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, - /// Rows one encode may take: it sizes scratch, never the kernel source; cooperative encodes still take 16. + /// Row capacity for reg layouts and opted-in row blocks; cooperative and scalar layouts take 16. rows: usize = reg_rows, + /// Non-reg prompt callers opt into independent 16-row blocks on the grid's y. + row_blocks: bool = false, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, @@ -120,7 +122,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { } try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); - try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(!l.reg and !l.cooperative and l.rows > 16)}); + try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(l.row_blocks and !l.reg and !l.cooperative and l.rows > 16)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); @@ -242,14 +244,14 @@ pub const Projection = struct { threads_per_slice: usize = 32, reg: bool = false, cooperative: bool = false, - /// Rows one encode takes at most: the layout's rows on tensor units, 16 on cooperative or older-chip paths. + /// Capacity is the reg or opted-in row count on tensor units, otherwise 16. max_rows: u32 = reg_rows, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); const units = device.tensorUnits(); if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; - const capacity: usize = if (l.cooperative or !units) 16 else l.rows; + const capacity: usize = if (l.cooperative or !units or (!l.reg and !l.row_blocks)) 16 else l.rows; var kernel_layout = l; if (!units and !l.cooperative and !l.reg) kernel_layout.rows = capacity; const text = try source(a, kernel_layout); diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index be2149254ae..19fcd480139 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -174,6 +174,7 @@ pub const Kernels = struct { .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg or (p.bits == 8 and p.group == 64), .reg = reg, + .row_blocks = long and !reg, .rows = if (long) lane.max_reg_rows else if (reg) lane.reg_rows else 16, }); k.joined(entry.value_ptr); diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig index 6add4d36ffe..c2b00f105ad 100644 --- a/zig/tests/core_lane_long.zig +++ b/zig/tests/core_lane_long.zig @@ -1,4 +1,4 @@ -//! Whole prompt lanes at each width against main's frozen 16-row kernel, including strided oQ member formats. +//! Whole prompt lanes at each width against a frozen 16-row kernel, including strided oQ member formats. const std = @import("std"); const mtl = @import("metal"); const lane = @import("core_lane"); @@ -90,7 +90,7 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st const ng = k / format.group; // Group128 SK8 exceeds the unchanged admission bound; compare its supported SK4 variant on both sides. const sk = if (format.group == 128 and shape[2] == 8) 4 else shape[2]; - const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; + const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .row_blocks = true, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; const p_long = try lane.Projection.init(a, device, layout); defer p_long.deinit(); var short = layout; @@ -134,7 +134,7 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st defer want.deinit(); const got = try device.buffer(most * stride * 2, opts); defer got.deinit(); - for ([_]u32{ 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { + for ([_]u32{ 16, 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { if (external and fi != 0) continue; const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; clear(want, rows, n, stride, false); From ae705fe50b54e95c6101222c7a71545408ca1c71 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:41 +0000 Subject: [PATCH 094/186] changelog: #2851 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8975e9de0af..535cd9c58e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). + - GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). - CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. From 36aa8300a96c7ef7921191b3f64f44370159acb9 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:25:40 +0900 Subject: [PATCH 095/186] glm cuda: TF_GLM_CONTEXT_CEILING serves windows past 1,048,576 tokens, read at GLM's open GLM-5.3-Flash's attention keeps no rotary positions, so only its trained length bounds the window. TF_GLM_CONTEXT_CEILING, a .glm row, sets the window --context may ask for, up to 2^31 - 32, where the slots keep i32 positions. GLM reads it at its open once its .glm rows are checked; unset, the ceiling and the refusal are main's. A window past the trained length says so at startup. --- docs/environment.md | 1 + docs/recipes/glm-5.3-flash-cuda.md | 9 ++++++ zig/src/families/glm/cuda_env_test.zig | 33 +++++++++++++++++++++- zig/src/families/glm/cuda_native.zig | 17 +++++++++++- zig/src/native/cuda.zig | 38 ++++++++++++++++++++++++-- zig/src/native/cuda_window.zig | 21 ++++++++++++++ zig/src/platform/environment.zig | 2 ++ 7 files changed, 117 insertions(+), 4 deletions(-) diff --git a/docs/environment.md b/docs/environment.md index 0af40a6f12c..b1818515681 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -43,6 +43,7 @@ Diagnostics quote and escape the supplied text so it cannot create another termi | `TF_GLM_SHARED` | Exactly `0` or `1`; default on (GLM windows share a forward); read on rank 0 | | `TF_GLM_DFLASH2_LOG` | Whole unsigned 64-bit count of DFlash2 rounds whose drafts are logged; default 0 | | `TF_GLM_DRAFT_CONFIDENCE` | Finite number 0 through 1 where a DFlash2 chain's chance product stops it; default 0.7 | +| `TF_GLM_CONTEXT_CEILING` | GLM-5.3-Flash on CUDA: whole token count 1 through 2147483616, the window `--context` may ask for and `--context 0` serves, in place of config.json's 1048576; read at GLM's open. Past that length the model's quality is the operator's to check. Unset keeps config.json's window | | `TF_PROMPT_ROWS` | Native serving: a multiple of 16 from 128 through 2048; unset keeps the chip/family default | | `TF_ANE_COLUMNS` | Whole u32 count; zero disables the experimental share; default zero | | `TF_FLASHNEXT_SCALE_INDEX` | Flash Next on Metal before tensor units: exactly `0` or `1`; `1` has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs, the same values with fewer weight bytes; default 0 | diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 93f671bd8fe..e8ed3350863 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -250,6 +250,15 @@ Long prompts and long runs of replies heat a GB10. The runbook's [heat bands](.. 1,035,295-token prompt at TP=3), [CPU pinning](../../RUNBOOK.md#gb10-cpu-pinning-for-long-cuda-runs) and [GPU clock cap](../../RUNBOOK.md#gb10-gpu-clock-cap-for-long-cuda-runs) each carry a measured GLM-5.3-Flash run. +## Windows past 1,048,576 tokens + +GLM-5.3-Flash's attention keeps no rotary positions (`qk_rope_head_dim` is 0, and KDA keeps none), so only its trained +length, config.json's 1,048,576 tokens, bounds the window. `TF_GLM_CONTEXT_CEILING` sets the window `--context` may ask +for and `--context 0` serves. GLM reads it at its open, and the startup log says when the window passes the trained +length, past which the quality is the operator's to check. On three GB10s with `Bizuayeu/GLM-5.3-Flash-XSFP4`, +`TF_GLM_CONTEXT_CEILING=2097152`, `--context 2097152 --parallel 1` and no prompt cache, rank 0 took 4,097 KV pages, and +a 1,996,430-token prompt answered all three of its passphrases. + ## Shared KDA operator `cuda.kda` owns the fused and three-kernel chains and single-layer and batched saved-row replay. diff --git a/zig/src/families/glm/cuda_env_test.zig b/zig/src/families/glm/cuda_env_test.zig index f7a6e7de348..023b4cebde7 100644 --- a/zig/src/families/glm/cuda_env_test.zig +++ b/zig/src/families/glm/cuda_env_test.zig @@ -18,6 +18,8 @@ const switches = [_]Name{ }; const counts = [_][:0]const u8{"TF_GLM_DFLASH2_LOG"}; const reals = [_][:0]const u8{"TF_GLM_DRAFT_CONFIDENCE"}; +/// The window GLM may serve past config.json's: GLM's raisedCeiling reads it at the open, after these rows' check. +const windows = [_][:0]const u8{"TF_GLM_CONTEXT_CEILING"}; /// The row by name without `rule`, so a missing row fails an assertion instead of panicking. fn row(name: []const u8) !e.Rule { @@ -54,12 +56,13 @@ test "GLM strict env: every GLM reader's name is a glm row, and the glm scope ho for (switches) |s| try t.expectEqual(e.Kind.flag, (try row(s.name)).kind); for (counts) |name| try t.expectEqual(e.Kind.u64, (try row(name)).kind); for (reals) |name| try t.expectEqual(e.Kind.real, (try row(name)).kind); + for (windows) |name| try t.expectEqual(e.Kind.u64, (try row(name)).kind); var n: usize = 0; for (e.rules) |r| if (r.scope == .glm) { try t.expect(std.mem.startsWith(u8, r.name, "TF_GLM_")); n += 1; }; - try t.expectEqual(switches.len + counts.len + reals.len, n); + try t.expectEqual(switches.len + counts.len + reals.len + windows.len, n); } test "GLM strict env: switches keep unset, blank, 0 and 1, and refuse words and padded digits by name" { @@ -137,3 +140,31 @@ test "GLM strict env: the DFlash2 cut is a finite chance from 0 to 1, read as a try t.expectError(error.InvalidEnvironment, confidence()); try t.expectError(error.InvalidEnvironment, @import("cuda_env.zig").startup()); } + +test "GLM strict env: the context ceiling is a whole token count from 1 to 2147483616; unset or blank keeps config.json's" { + for (windows) |name| { + _ = try row(name); + for ([_]?[]const u8{ null, "", " \t" }) |raw| { + try t.expectEqual(@as(?u64, null), try e.value(?u64, name, raw, null)); + if (raw) |v| try check(name, v, true); + } + for ([_][]const u8{ "1", "2097152", "2147483616" }) |raw| { + try t.expectEqual(@as(?u64, try std.fmt.parseInt(u64, raw, 10)), try e.value(?u64, name, raw, null)); + try check(name, raw, true); + } + for ([_][]const u8{ "0", "2147483617", "+2097152", " 2097152", "2097152 ", "2M", "0x200000", "2097152.0" }) |raw| try check(name, raw, false); + } +} + +test "GLM strict env: the open's check refuses a malformed ceiling before ranks, NCCL or weights" { + const startup = @import("cuda_env.zig").startup; + const name = "TF_GLM_CONTEXT_CEILING"; + _ = try row(name); + _ = env.unsetenv(name); + defer _ = env.unsetenv(name); + try startup(); + _ = env.setenv(name, "2097152", 1); + try startup(); + _ = env.setenv(name, "2M", 1); + try t.expectError(error.InvalidEnvironment, startup()); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 904d7cb72aa..01323464187 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -22,6 +22,15 @@ pub const model_type = "glm5_next"; pub const formats: []const []const u8 = &.{ "modelopt-nvfp4", "modelopt-mixed", "exl3" }; // cc-defer: a fixed window; the Python family sizes it from free memory (capacity.admit), port that when serving needs pub const default_context: i64 = 8192; +/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. +pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; + +/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. +pub fn raisedCeiling() !?u64 { + try @import("cuda_env.zig").startup(); + const environment = @import("sys").environment; + return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); +} pub const max_segments: u32 = 1; pub const prompt_rows: u32 = fw.prefill_rows; /// Chunk rows change prompt arithmetic: states keep and resume only where a cold pass cuts (cuda_prompt_cuts). @@ -581,11 +590,17 @@ test "capacity keeps a verify window past the context and a multiple of 64" { test "the plan's window, the context and a verify window, takes the pages of a stream's table at that capacity" { const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = @import("cuda_state.zig").page_rows, .snap_bytes = 0 }; - for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576 }) |context| { + for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576, 2_097_152, 2_097_152 + 1 }) |context| { try std.testing.expectEqual(pages_mod.tableLen(capacity(context)), l.pages(context + window_rows)); } } +test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { + const most = @import("sys").environment.rule(ceiling_switch).max; + try std.testing.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); + try std.testing.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); +} + // Enough words for every offered frame pair and image, even with a small text context. pub fn linkCapacity(cap: usize, media: ?images_mod.Limits) !usize { const lim = media orelse return cap; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 57c34a6bac1..c9b1be77b68 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -181,18 +181,26 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; } // a family whose rope can stretch (Qwen's YaRN, from its config or --rope-yarn) names its own ceiling - const native = if (comptime @hasDecl(F, "contextCeiling")) F.contextCeiling(a, io, o.dir, o.rope_yarn) catch |e| { + const own = if (comptime @hasDecl(F, "contextCeiling")) F.contextCeiling(a, io, o.dir, o.rope_yarn) catch |e| { problem.* = try std.fmt.allocPrint(a, "{s} ({s}): YaRN takes rope_type yarn, factor, original_max_position_embeddings and their product as max_position_embeddings, and --rope-yarn a factor from 1 matching it", .{ if (o.rope_yarn != null) "--rope-yarn" else "the checkpoint's rope scaling", @errorName(e) }); return null; } else modelContext(a, io, o.dir); + // a family with no positions to stretch (GLM-5.3-Flash) raises it with a row of its own, read at its open + const raised: ?u64 = if (comptime @hasDecl(F, "raisedCeiling")) F.raisedCeiling() catch |e| { + problem.* = F.explain(null, e) orelse @errorName(e); + return null; + } else null; + const switch_name = if (comptime @hasDecl(F, "raisedCeiling")) F.ceiling_switch else ""; + const native: i64 = if (raised) |n| @intCast(n) else own; const window = contextWindow(o.context, native, F.default_context) catch |e| { problem.* = switch (e) { error.Negative => try std.fmt.allocPrint(a, "--context {d}: a token count, or 0 for the model's window", .{o.context.?}), error.NoNative => "--context 0 asks for the model's window, and its config.json names none: give a token count", - error.PastNative => try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), + error.PastNative => if (raised != null) try windows.pastCeiling(a, o.context.?, native, switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), }; return null; }; + if (raised != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, switch_name)}); if (o.slide) { // a family learns once its module declares `slide` (its learner), `learnModel` and `linked` if (!@hasDecl(F, "slide")) { @@ -554,6 +562,32 @@ test "Flash Next takes kept prompt states and drafts together on any rank count" test "Flash Next takes --vision at family admission" { try @import("cuda_flashnext_refusal_test.zig").vision(); } +test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { + const env = platform.env; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + var named: usize = 0; + inline for (registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { + defer _ = env.unsetenv(F.ceiling_switch); + var why: []const u8 = ""; + _ = env.unsetenv(F.ceiling_switch); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); + _ = env.setenv(F.ceiling_switch, "2M", 1); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); + _ = env.setenv(F.ceiling_switch, "2097152", 1); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); + named += 1; + }; + try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING +} test { _ = @import("cuda_health_test.zig"); _ = @import("cuda_vision_offer.zig"); diff --git a/zig/src/native/cuda_window.zig b/zig/src/native/cuda_window.zig index 355dc4d1b4b..7bd01019232 100644 --- a/zig/src/native/cuda_window.zig +++ b/zig/src/native/cuda_window.zig @@ -13,6 +13,17 @@ pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { return if (limit == .integer and limit.integer > 0) limit.integer else 0; } +/// A --context past the window a family's switch raises its ceiling to: the switch named. +pub fn pastCeiling(a: Allocator, requested: i64, ceiling: i64, switch_name: []const u8) ![]const u8 { + return std.fmt.allocPrint(a, "--context {d} exceeds the {d}-token window {s} sets", .{ requested, ceiling, switch_name }); +} + +/// The startup line of a window past the checkpoint's trained one, which only the family's switch serves. +pub fn pastTrained(a: Allocator, window: i64, trained: i64, switch_name: []const u8) ![]const u8 { + const words = "{s}: a {d}-token window, past the checkpoint's trained {d} tokens; past them the quality is the operator's to check"; + return std.fmt.allocPrint(a, words, .{ switch_name, window, trained }); +} + /// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. pub fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { const r = requested orelse return if (native > 0) @min(default, native) else default; @@ -31,3 +42,13 @@ test "context windows: the family default inside the model's, 0 for the model's, try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); } + +test "context windows: past a family's raised ceiling the refusal names its switch, and the startup line the trained window" { + const a = std.testing.allocator; + const past = try pastCeiling(a, 2097153, 2097152, "TF_GLM_CONTEXT_CEILING"); + defer a.free(past); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window TF_GLM_CONTEXT_CEILING sets", past); + const line = try pastTrained(a, 2097152, 1048576, "TF_GLM_CONTEXT_CEILING"); + defer a.free(line); + try std.testing.expectEqualStrings("TF_GLM_CONTEXT_CEILING: a 2097152-token window, past the checkpoint's trained 1048576 tokens; past them the quality is the operator's to check", line); +} diff --git a/zig/src/platform/environment.zig b/zig/src/platform/environment.zig index 3ad7ecb4941..c73ab2066f4 100644 --- a/zig/src/platform/environment.zig +++ b/zig/src/platform/environment.zig @@ -82,6 +82,8 @@ pub const rules = [_]Rule{ .{ .name = "TF_GLM_SHARED", .scope = .glm, .kind = .flag, .forms = "0 or 1" }, .{ .name = "TF_GLM_DFLASH2_LOG", .scope = .glm, .kind = .u64, .forms = "a whole unsigned 64-bit round count (0 logs no drafts)" }, .{ .name = "TF_GLM_DRAFT_CONFIDENCE", .scope = .glm, .kind = .real, .max = 1, .forms = "a finite chain-chance cut from 0 to 1 (default 0.7)" }, + // 2^31 - 32: GLM's slots, the window and its 32 verify rows (glm/cuda_native.zig's capacity), keep i32 positions + .{ .name = "TF_GLM_CONTEXT_CEILING", .scope = .glm, .kind = .u64, .first_digit = true, .min = 1, .max = 2147483616, .forms = "a whole token count from 1 to 2147483616 (the window GLM-5.3-Flash may serve in place of config.json's)" }, .{ .name = "TF_DSV41_MHC_DEFER", .scope = .dsv41, .kind = .toggle, .trim = " \t", .no_word = false, .forms = "0, 1, true, false, on, off or yes (any case)" }, .{ .name = "TF_DSV41_COEF_LATE", .scope = .dsv41, .kind = .toggle, .trim = " \t", .no_word = false, .forms = "0, 1, true, false, on, off or yes (any case)" }, .{ .name = "TF_DSV41_LOOKUP", .scope = .dsv41, .kind = .toggle, .trim = " \t", .forms = "0, 1, true, false, on, off, yes or no (any case); blank turns lookup drafts off" }, From bbb60d85200f5fad4fc3e12fe03545fd7049e202 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 02:17:33 +0900 Subject: [PATCH 096/186] glm cuda: TF_GLM_CONTEXT_CEILING's reader joins the .glm check in cuda_env.zig cuda_native.zig had passed the 600-line limit. ceiling_switch and raisedCeiling move to cuda_env.zig, which owns the .glm check, and the i32 test to cuda_env_test.zig; cuda_native.zig keeps the two names the CUDA open looks up. --- zig/src/families/glm/cuda_env.zig | 9 +++++++++ zig/src/families/glm/cuda_env_test.zig | 7 +++++++ zig/src/families/glm/cuda_native.zig | 17 ++--------------- 3 files changed, 18 insertions(+), 15 deletions(-) diff --git a/zig/src/families/glm/cuda_env.zig b/zig/src/families/glm/cuda_env.zig index c929b466b70..c5f37582c11 100644 --- a/zig/src/families/glm/cuda_env.zig +++ b/zig/src/families/glm/cuda_env.zig @@ -15,6 +15,15 @@ pub fn requestStartsAlone(learn_conversations: bool) !void { if (on and try kp.sharing() > 0) return error.GlmRequestStartsWithSharedKv; } +/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. +pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; + +/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. +pub fn raisedCeiling() !?u64 { + try startup(); + return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); +} + test { _ = @import("cuda_env_test.zig"); } diff --git a/zig/src/families/glm/cuda_env_test.zig b/zig/src/families/glm/cuda_env_test.zig index 023b4cebde7..6f678d94dcb 100644 --- a/zig/src/families/glm/cuda_env_test.zig +++ b/zig/src/families/glm/cuda_env_test.zig @@ -168,3 +168,10 @@ test "GLM strict env: the open's check refuses a malformed ceiling before ranks, _ = env.setenv(name, "2M", 1); try t.expectError(error.InvalidEnvironment, startup()); } + +test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { + const capacity = @import("cuda_native.zig").capacity; + const most = e.rule(@import("cuda_env.zig").ceiling_switch).max; + try t.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); + try t.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 01323464187..099e47af6b2 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -22,15 +22,8 @@ pub const model_type = "glm5_next"; pub const formats: []const []const u8 = &.{ "modelopt-nvfp4", "modelopt-mixed", "exl3" }; // cc-defer: a fixed window; the Python family sizes it from free memory (capacity.admit), port that when serving needs pub const default_context: i64 = 8192; -/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. -pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; - -/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. -pub fn raisedCeiling() !?u64 { - try @import("cuda_env.zig").startup(); - const environment = @import("sys").environment; - return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); -} +pub const ceiling_switch = @import("cuda_env.zig").ceiling_switch; +pub const raisedCeiling = @import("cuda_env.zig").raisedCeiling; pub const max_segments: u32 = 1; pub const prompt_rows: u32 = fw.prefill_rows; /// Chunk rows change prompt arithmetic: states keep and resume only where a cold pass cuts (cuda_prompt_cuts). @@ -595,12 +588,6 @@ test "the plan's window, the context and a verify window, takes the pages of a s } } -test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { - const most = @import("sys").environment.rule(ceiling_switch).max; - try std.testing.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); - try std.testing.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); -} - // Enough words for every offered frame pair and image, even with a small text context. pub fn linkCapacity(cap: usize, media: ?images_mod.Limits) !usize { const lim = media orelse return cap; From fcfa424fba368279e51e740d5ed52947b5fea381 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 07:00:52 +0900 Subject: [PATCH 097/186] native cuda: the ceiling test and lookup leave cuda.zig for its own test file and cuda_window.zig Main grew cuda.zig to 587 lines, so the branch's additions passed the 600-line limit again. The open-time test of raised context ceilings moves to cuda_ceiling_test.zig, imported beside cuda_health_test.zig, and the lookup of a family's raised ceiling and its switch to cuda_window.zig's familyCeiling; cuda.zig keeps one call. --- zig/src/native/cuda.zig | 39 ++++------------------------ zig/src/native/cuda_ceiling_test.zig | 30 +++++++++++++++++++++ zig/src/native/cuda_window.zig | 13 ++++++++++ 3 files changed, 48 insertions(+), 34 deletions(-) create mode 100644 zig/src/native/cuda_ceiling_test.zig diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index c9b1be77b68..8036a4b1c82 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -186,21 +186,17 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; } else modelContext(a, io, o.dir); // a family with no positions to stretch (GLM-5.3-Flash) raises it with a row of its own, read at its open - const raised: ?u64 = if (comptime @hasDecl(F, "raisedCeiling")) F.raisedCeiling() catch |e| { - problem.* = F.explain(null, e) orelse @errorName(e); - return null; - } else null; - const switch_name = if (comptime @hasDecl(F, "raisedCeiling")) F.ceiling_switch else ""; - const native: i64 = if (raised) |n| @intCast(n) else own; + const raised = windows.familyCeiling(F, problem) catch return null; + const native: i64 = if (raised.window) |n| @intCast(n) else own; const window = contextWindow(o.context, native, F.default_context) catch |e| { problem.* = switch (e) { error.Negative => try std.fmt.allocPrint(a, "--context {d}: a token count, or 0 for the model's window", .{o.context.?}), error.NoNative => "--context 0 asks for the model's window, and its config.json names none: give a token count", - error.PastNative => if (raised != null) try windows.pastCeiling(a, o.context.?, native, switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), + error.PastNative => if (raised.window != null) try windows.pastCeiling(a, o.context.?, native, raised.switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), }; return null; }; - if (raised != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, switch_name)}); + if (raised.window != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, raised.switch_name)}); if (o.slide) { // a family learns once its module declares `slide` (its learner), `learnModel` and `linked` if (!@hasDecl(F, "slide")) { @@ -562,34 +558,9 @@ test "Flash Next takes kept prompt states and drafts together on any rank count" test "Flash Next takes --vision at family admission" { try @import("cuda_flashnext_refusal_test.zig").vision(); } -test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { - const env = platform.env; - var tmp = std.testing.tmpDir(.{}); - defer tmp.cleanup(); - try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); - var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); - defer arena.deinit(); - const a = arena.allocator(); - const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); - var named: usize = 0; - inline for (registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { - defer _ = env.unsetenv(F.ceiling_switch); - var why: []const u8 = ""; - _ = env.unsetenv(F.ceiling_switch); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); - try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); - _ = env.setenv(F.ceiling_switch, "2M", 1); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); - try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); - _ = env.setenv(F.ceiling_switch, "2097152", 1); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); - try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); - named += 1; - }; - try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING -} test { _ = @import("cuda_health_test.zig"); + _ = @import("cuda_ceiling_test.zig"); _ = @import("cuda_vision_offer.zig"); _ = @import("cuda_ledger.zig"); _ = kv_pages; diff --git a/zig/src/native/cuda_ceiling_test.zig b/zig/src/native/cuda_ceiling_test.zig new file mode 100644 index 00000000000..f9d1eacc0cd --- /dev/null +++ b/zig/src/native/cuda_ceiling_test.zig @@ -0,0 +1,30 @@ +//! A family's raised context ceiling (GLM's TF_GLM_CONTEXT_CEILING) through the real CUDA open, before any weights. +const std = @import("std"); +const env = @import("sys").env; +const cuda = @import("cuda.zig"); + +test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + var named: usize = 0; + inline for (cuda.registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { + defer _ = env.unsetenv(F.ceiling_switch); + var why: []const u8 = ""; + _ = env.unsetenv(F.ceiling_switch); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); + _ = env.setenv(F.ceiling_switch, "2M", 1); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); + _ = env.setenv(F.ceiling_switch, "2097152", 1); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); + named += 1; + }; + try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING +} diff --git a/zig/src/native/cuda_window.zig b/zig/src/native/cuda_window.zig index 7bd01019232..2e94e48fc39 100644 --- a/zig/src/native/cuda_window.zig +++ b/zig/src/native/cuda_window.zig @@ -13,6 +13,19 @@ pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { return if (limit == .integer and limit.integer > 0) limit.integer else 0; } +/// The window a family's own row raises --context's ceiling to, and the switch its refusal names; none without one. +pub const Ceiling = struct { window: ?u64 = null, switch_name: []const u8 = "" }; + +/// A family's raised ceiling, read at its open (GLM's TF_GLM_CONTEXT_CEILING); a malformed row is the family's refusal. +pub fn familyCeiling(comptime F: type, problem: *[]const u8) error{Refused}!Ceiling { + if (comptime !@hasDecl(F, "raisedCeiling")) return .{}; + const window = F.raisedCeiling() catch |e| { + problem.* = F.explain(null, e) orelse @errorName(e); + return error.Refused; + }; + return .{ .window = window, .switch_name = F.ceiling_switch }; +} + /// A --context past the window a family's switch raises its ceiling to: the switch named. pub fn pastCeiling(a: Allocator, requested: i64, ceiling: i64, switch_name: []const u8) ![]const u8 { return std.fmt.allocPrint(a, "--context {d} exceeds the {d}-token window {s} sets", .{ requested, ceiling, switch_name }); From 90eeb3d689826f4780460abb529769c744806d26 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Thu, 12 Nov 2026 05:54:24 +0900 Subject: [PATCH 098/186] glm cuda: the capacity test leaves cuda_native.zig for cuda_pages_test.zig Main's eb29a68a5f added a line to cuda_native.zig, so merged with main the file comes to 601 lines, past the 600-line limit. The test that capacity keeps a verify window past the context, a multiple of 64, moves whole to cuda_pages_test.zig, which tests GLM's pages at that capacity: 594 lines in cuda_native.zig here, 595 merged with main. --- zig/src/families/glm/cuda_native.zig | 6 ------ zig/src/families/glm/cuda_pages_test.zig | 6 ++++++ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 099e47af6b2..4a1f5cd1a5c 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -575,12 +575,6 @@ test "TF_KV_PAGES with --kv-cache fp8: auto and off keep windows at every rank c try std.testing.expect(explain(null, error.GlmKvHostNeedsPages) != null); } -test "capacity keeps a verify window past the context and a multiple of 64" { - try std.testing.expectEqual(@as(usize, 8256), capacity(8192)); - try std.testing.expectEqual(@as(usize, 330624), capacity(330529)); - try std.testing.expect(capacity(330529) % 64 == 0); -} - test "the plan's window, the context and a verify window, takes the pages of a stream's table at that capacity" { const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = @import("cuda_state.zig").page_rows, .snap_bytes = 0 }; for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576, 2_097_152, 2_097_152 + 1 }) |context| { diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index 7fdfc60b657..76de1632cfa 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -530,3 +530,9 @@ test "every host allocation failure while making the pages leaves no device buff } }.run, .{}); } + +test "capacity keeps a verify window past the context and a multiple of 64" { + try expectEqual(@as(usize, 8256), native.capacity(8192)); + try expectEqual(@as(usize, 330624), native.capacity(330529)); + try expect(native.capacity(330529) % 64 == 0); +} From 5ef0b1fa2a34b694a793e5074d064880d914919d Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:44 +0000 Subject: [PATCH 099/186] changelog: #2765 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 535cd9c58e2..144b294f96e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. + - Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). - GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). From 100e4a0848d0be4be96b57cc04ef6427dff6ba5e Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Wed, 11 Nov 2026 08:39:08 +0000 Subject: [PATCH 100/186] fix(server): let the last --compact-at replace the earlier one A fraction after --compact-at auto left auto mode on, because the flag set compact_auto and a later fraction never cleared it, so the server kept the first value where every other flag keeps the last. Each occurrence now replaces both fields, which a config file's value under a command-line override depends on, and a filtered host step runs the test that fails without the change. --- zig/build/cuda_server_host.zig | 1 + zig/src/server/cli.zig | 4 +++- zig/src/server/compact_at_cli_test.zig | 23 +++++++++++++++++++++++ zig/src/server/root.zig | 1 + 4 files changed, 28 insertions(+), 1 deletion(-) create mode 100644 zig/src/server/compact_at_cli_test.zig diff --git a/zig/build/cuda_server_host.zig b/zig/build/cuda_server_host.zig index c77531a9b00..9da10918454 100644 --- a/zig/build/cuda_server_host.zig +++ b/zig/build/cuda_server_host.zig @@ -29,6 +29,7 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = mods.tokenizer, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), .filters = &.{"EOS policy:"} })).step); + b.step("test-serve-flags-host", "Serve flag parsing a later flag replaces, without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--compact-at:"} })).step); b.step("test-dory-cli-host", "Dory CUDA loop option forwarding and invalid-value refusals").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--device and --segments:"} })).step); b.step("test-dory-chat-host", "Dory chat template, EOS, reasoning split and qwen3_coder call on the CPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Dory chat:"} })).step); b.step("test-embeddings-protocol-host", "Embedding request and response contracts without an admitted model or GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"embeddings protocol:"} })).step); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 387dff1c8c6..82975c5f7e5 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -472,7 +472,9 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.decode_share = try float(u, a, name, v); if (!(out.decode_share.? >= 0 and out.decode_share.? < 1)) return fail(u, a, "argument --decode-share: expected a share from 0 to under 1: '{s}'", .{v}); } else if (is(name, "--prompt-seed")) out.prompt_seed = true else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--instruct-sampling")) out.instruct_sampling = if (try @import("family_sampling.zig").pairsProblem(a, v)) |why| return fail(u, a, "argument {s}: {s}: '{s}'", .{ name, why, v }) else v else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--loop-guard-stop")) out.loop_guard_stop = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { - if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { + out.compact_auto = std.mem.eql(u8, v, "auto"); // each occurrence replaces both fields, so the last one wins + out.compact_fraction = null; + if (!out.compact_auto) { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); out.compact_fraction = f; diff --git a/zig/src/server/compact_at_cli_test.zig b/zig/src/server/compact_at_cli_test.zig new file mode 100644 index 00000000000..31cd835c5af --- /dev/null +++ b/zig/src/server/compact_at_cli_test.zig @@ -0,0 +1,23 @@ +//! ``--compact-at`` given twice: each occurrence replaces both of its fields, so the last one wins. +const std = @import("std"); +const cli = @import("cli.zig"); + +test "--compact-at: a later occurrence replaces an earlier one in full" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var u: cli.Usage = .{}; + const fraction_last = try cli.parse(a, &.{ "m", "--compact-at", "auto", "--compact-at=0.5" }, &u); + try std.testing.expect(!fraction_last.compact_auto); + try std.testing.expectEqual(@as(?f64, 0.5), fraction_last.compact_fraction); + const auto_last = try cli.parse(a, &.{ "m", "--compact-at=0.5", "--compact-at", "auto" }, &u); + try std.testing.expect(auto_last.compact_auto); + try std.testing.expectEqual(@as(?f64, null), auto_last.compact_fraction); + const auto = try cli.parse(a, &.{ "m", "--compact-at", "auto" }, &u); + try std.testing.expect(auto.compact_auto and auto.compact_fraction == null); + const fraction = try cli.parse(a, &.{ "m", "--compact-at", "0.25" }, &u); + try std.testing.expect(!fraction.compact_auto); + try std.testing.expectEqual(@as(?f64, 0.25), fraction.compact_fraction); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--compact-at", "1.5" }, &u)); + try std.testing.expectEqualStrings("argument --compact-at: expected auto or a fraction in (0, 1]: '1.5'", u.message); +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 3bfbc4ad96c..66d0da36073 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -76,6 +76,7 @@ test { _ = @import("cli.zig"); _ = @import("cli_test.zig"); _ = @import("cli_profile_test.zig"); + _ = @import("compact_at_cli_test.zig"); _ = @import("unix_socket_cli_test.zig"); _ = @import("effort_cli_audit_test.zig"); _ = @import("live.zig"); From d8a2da42fcd0d63048cc07d6b604a3d8e97aa839 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:45 +0000 Subject: [PATCH 101/186] changelog: #2883 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 144b294f96e..a62f1873c78 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. + - GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. - Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). From 60a6656accf4e44fa248ececd0d3b8f383b82b91 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 01:23:44 -0800 Subject: [PATCH 102/186] test: align Qwen27 reply history with newest-state retention Pin eviction of redundant own history before another conversation's newest state under pressure. Retain prompt-born provenance and prove a same-prompt regenerate restores history when all three states fit. Refs #2881 --- zig/src/native/qwen27_reply_cache_test.zig | 37 ++++++++++++++++++++++ zig/src/native/qwen27_served.zig | 24 ++------------ 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/zig/src/native/qwen27_reply_cache_test.zig b/zig/src/native/qwen27_reply_cache_test.zig index 91b8d078dbc..50d2be0d67c 100644 --- a/zig/src/native/qwen27_reply_cache_test.zig +++ b/zig/src/native/qwen27_reply_cache_test.zig @@ -9,6 +9,43 @@ const Metal = @import("qwen27_lanes.zig").Metal; const t = std.testing; const a = t.allocator; +/// Reply ends keep their request's provenance; spare budget also retains the pre-reply regenerate state. +pub fn replyBorn(keepEnd: anytype, snapshots: pc.Snapshots, live: *usize) !void { + const at: u32 = 5000; + for ([_]bool{ false, true }) |roomy| { + const rows: u64 = 3 * at + @as(u64, if (roomy) 107 else 100); + var store = pc.Store.init(a, snapshots, .{ .warm = .on }, rows << 16); + defer store.deinit(); + var other: [at + 7]u32 = undefined; + for (&other, 0..) |*word, i| word.* = @intCast(i); + try t.expect(store.keep(&other, at, null, &.{}, &.{})); + var tokens: [at + 107]u32 = undefined; + for (&tokens, 0..) |*word, i| word.* = @intCast(1_000_000 + i); + const prompt = tokens[0 .. at + 7]; + try t.expect(store.keep(prompt, at, null, &.{}, &.{})); + var s = try lanes.Stream.init(a, .{ .id = "reply", .prompt = prompt, .max_new = 100 }); + defer s.deinit(a); + try s.context.appendSlice(a, tokens[at + 7 ..]); + keepEnd(&store, a, &s, tokens.len); + const newest_other = store.find(&other, &.{}, &.{}) orelse return error.NewestConversationEvicted; + try t.expectEqual(at, newest_other.at); + const reply = store.entries.items[store.entries.items.len - 1]; + try t.expectEqual(@as(u32, tokens.len), reply.at); + try t.expectEqual(@as(u32, at + 7), reply.born); + try t.expectEqual(@as(usize, if (roomy) 3 else 2), store.entries.items.len); + try t.expectEqual(store.entries.items.len, live.*); + if (roomy) { + const history = store.find(prompt, &.{}, &.{}) orelse return error.HistoryEvicted; + try t.expectEqual(at, history.at); + const plan = try store.begin(a, prompt, at, &.{}, &.{}, null, &.{}); + defer a.free(plan.marks); + try t.expectEqual(at, plan.from); + try t.expectEqual(@as(u64, 1), store.counts.hits); + } else try t.expect(store.find(prompt, &.{}, &.{}) == null); + } + try t.expectEqual(@as(usize, 0), live.*); +} + const Family = struct { at: u32 = 0, sum: u64 = 0, diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 05c40ddba7f..2e24131464e 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -265,28 +265,10 @@ test "--prompt-cache-gib 40: a later 192k-token slot takes its room back from ke try std.testing.expectEqual(@as(usize, 1), fake.live); } -test "a reply's end is born with its request's prompt: the history mark that request kept outlives the reply's keep" { - const gpa = std.testing.allocator; - const at: u32 = 5000; // a history mark at 5000 tokens, then a 7-token generation prompt and a 100-token reply +test "a reply's end keeps its prompt's born: newest conversations precede redundant history under pressure" { var fake: FakeStates = .{}; - var store = pc.Store.init(gpa, .{ .ptr = &fake, .vtable = &.{ .bytes = FakeStates.bytes, .save = FakeStates.save, .restore = FakeStates.restore, .drop = FakeStates.drop } }, .{ .warm = .on }, (3 * at + 100) << 16); - defer store.deinit(); - var other: [at + 7]u32 = undefined; // another conversation's prompt, its history kept first - for (&other, 0..) |*t, i| t.* = @intCast(i); - try std.testing.expect(store.keep(&other, at, null, &.{}, &.{})); - var tokens: [at + 107]u32 = undefined; - for (&tokens, 0..) |*t, i| t.* = @intCast(1_000_000 + i); - const prompt = tokens[0 .. at + 7]; - try std.testing.expect(store.keep(prompt, at, null, &.{}, &.{})); - var s = try lanes.Stream.init(gpa, .{ .id = "reply", .prompt = prompt, .max_new = 100 }); - defer s.deinit(gpa); - try s.context.appendSlice(gpa, tokens[at + 7 ..]); - keepEnd(&store, gpa, &s, tokens.len); // three states pass the budget: the least recently used goes - try std.testing.expect(store.find(&other, &.{}, &.{}) == null); - const history = store.find(prompt, &.{}, &.{}) orelse return error.HistoryEvicted; // a regenerate resumes it - try std.testing.expectEqual(at, history.at); - try std.testing.expectEqual(@as(u32, at + 7), store.entries.items[1].born); - try std.testing.expectEqual(@as(usize, 2), fake.live); + const snapshots: pc.Snapshots = .{ .ptr = &fake, .vtable = &.{ .bytes = FakeStates.bytes, .save = FakeStates.save, .restore = FakeStates.restore, .drop = FakeStates.drop } }; + try @import("qwen27_reply_cache_test.zig").replyBorn(keepEnd, snapshots, &fake.live); } /// Memory the slots past the ready ones take on first use: their K/V rows and taps rings. From f13d3de486b9299677c852466569894dde5ea0e2 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:47 +0000 Subject: [PATCH 103/186] changelog: #2889 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a62f1873c78..7c6e70d4bfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). + - Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. - GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. From 72b0d7035eb3215276d1fb0ccfb8e9841e77d8a4 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 01:12:55 -0800 Subject: [PATCH 104/186] ci: run Metal engine host tests in the macOS nightly Give test-native a timed shard with the Metal toolchain. Keep it nightly until the CI runner establishes its wall time. Refs #2880 --- .github/workflows/no-gpu.yml | 20 +++++++++++++++++++- zig/build/host_ci.zig | 3 +++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 26b34ab3ef3..6960dc59073 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -335,7 +335,7 @@ jobs: strategy: fail-fast: false matrix: - suite: [fast, adapters, deepseek, lanes, deferred] + suite: [fast, adapters, native, deepseek, lanes, deferred] steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 with: @@ -344,11 +344,29 @@ jobs: with: version: 0.17.0 use-cache: false + - name: Metal toolchain for native engine host checks + if: matrix.suite == 'native' + run: xcrun -sdk macosx metal -v || { xcodebuild -downloadComponent MetalToolchain && xcrun -sdk macosx metal -v; } - name: Host checks and aggregate wall time shell: bash env: SUITE: ${{ matrix.suite }} run: | printf 'Runner class: %s %s\n' "$RUNNER_OS" "$RUNNER_ARCH" + if [ "$SUITE" = native ]; then + printf 'Native source: %s\n' "$GITHUB_SHA" + status=0 + for phase in initial warm; do + TIMEFORMAT="test-native $phase wall seconds: %R" + if { time zig build test-host-ci-macos-native --test-timeout 30m --summary all; }; then + printf 'test-native %s exit: 0\n' "$phase" + else + code=$? + printf 'test-native %s exit: %s\n' "$phase" "$code" + status=$code + fi + done + exit "$status" + fi TIMEFORMAT='Host checks wall seconds: %R' time zig build "test-host-ci-macos-$SUITE" --test-timeout 30m diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index 85206a8dd76..feaa56d56f0 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -14,6 +14,7 @@ pub fn steps(b: *std.Build) void { .macos => { wire(b, "test-host-ci-macos-fast", "Nightly host checks without a GPU or shader compiler", &macos_fast); wire(b, "test-host-ci-macos-adapters", "Nightly host checks without a GPU or shader compiler", &macos_adapters); + wire(b, "test-host-ci-macos-native", "Nightly Metal engine host checks with the Metal toolchain", &macos_native); wire(b, "test-host-ci-macos-deepseek", "Nightly host checks without a GPU or shader compiler", &macos_deepseek); wire(b, "test-host-ci-macos-lanes", "Nightly host checks without a GPU or shader compiler", &macos_lanes); wire(b, "test-host-ci-macos-deferred", "Nightly host checks without a GPU or shader compiler", &macos_deferred); @@ -289,6 +290,8 @@ const macos_adapters = [_][]const u8{ "test-verbs", }; +const macos_native = [_][]const u8{"test-native"}; + const macos_deepseek = [_][]const u8{ "test-dsv41-device-host", "test-dsv41-draft", From 77c2c1a0833e6fd85c10142db7e906ea5ae890c9 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:49 +0000 Subject: [PATCH 105/186] changelog: #2891 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c6e70d4bfd..4b01c602295 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). + - Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). - Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. From 1bc80c5f78bf66072dd77994392047840d1d0813 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 11 Nov 2026 07:59:25 +0100 Subject: [PATCH 106/186] core: one block-FP8 reader and exact scale widening; Flash-Next MTP experts, Drama 3 and Kolibri-1 read through it --- zig/build/checkpoint_io_host.zig | 2 + zig/src/core/quant/fp8.zig | 54 +++++++ zig/src/core/quant/fp8_block_test.zig | 140 ++++++++++++++++++ zig/src/families/drama3/cuda/weights.zig | 106 +++++++++++-- zig/src/families/drama3/drama3.zig | 1 + .../flashnext/cuda_fp8_experts_load.zig | 14 +- zig/src/families/kolibri1/names.zig | 13 +- zig/src/families/kolibri1/pack.zig | 13 +- 8 files changed, 300 insertions(+), 43 deletions(-) create mode 100644 zig/src/core/quant/fp8_block_test.zig diff --git a/zig/build/checkpoint_io_host.zig b/zig/build/checkpoint_io_host.zig index 77767ffa710..583d42f9b1e 100644 --- a/zig/build/checkpoint_io_host.zig +++ b/zig/build/checkpoint_io_host.zig @@ -11,6 +11,8 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget) void { step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"checkpoint admission:"} })).step); if (mode == .debug) b.step("test-modelopt-part-host", "The shared ModelOpt part reader (core/modelopt_part.zig) without a GPU") .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"part:"} })).step); + if (mode == .debug) b.step("test-fp8-block-host", "The shared block-FP8 reader and exact scale widening (core/quant/fp8.zig)") + .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"fp8 block:"} })).step); const qwen = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/loader.zig"), .target = target, .optimize = mode, .link_libc = true, .imports = &.{.{ .name = "core", .module = core }} }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = qwen, .filters = &.{"Qwen shard integrity:"} })).step); } diff --git a/zig/src/core/quant/fp8.zig b/zig/src/core/quant/fp8.zig index e95e38fbb43..3d400788b39 100644 --- a/zig/src/core/quant/fp8.zig +++ b/zig/src/core/quant/fp8.zig @@ -4,6 +4,7 @@ const types = @import("types.zig"); const convert = @import("convert.zig"); const Tensor = types.Tensor; +const Why = @import("../modelopt.zig").Why; /// Rows and columns one scale_inv entry covers. pub const block = 128; @@ -66,6 +67,59 @@ pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void { }; } +/// The tensor beside `{module}.weight` that holds a block-FP8 projection's scales. +pub const scale_tail = "weight_scale_inv"; + +/// Scale blocks along a dimension of `n`: a last partial block takes a scale of its own. +pub fn blocks(n: usize) usize { + return (n + block - 1) / block; +} + +/// `{module}.weight`, e4m3 [n, k], from any store with `get(name) !Tensor`; a refusal named in `why`, else logged. +pub fn weight(store: anytype, module: []const u8, n: usize, k: usize, why: ?*Why) !Tensor { + return checked(store, module, "weight", .f8_e4m3, &.{ n, k }, why); +} + +/// `{module}.weight_scale_inv` of `dtype` [ceil(n/128), ceil(k/128)]; each family names the dtype it stores. +pub fn scales(store: anytype, module: []const u8, n: usize, k: usize, dtype: types.DType, why: ?*Why) !Tensor { + if (!convert.tableDtype(dtype)) return error.UnexpectedTensor; + return checked(store, module, scale_tail, dtype, &.{ blocks(n), blocks(k) }, why); +} + +/// A block-FP8 projection [n, k] as stored: views into the store's bytes, nothing copied. +pub fn read(store: anytype, module: []const u8, n: usize, k: usize, dtype: types.DType, why: ?*Why) !Host { + return .{ .weight = try weight(store, module, n, k, why), .scale_inv = try scales(store, module, n, k, dtype, why) }; +} + +fn checked(store: anytype, module: []const u8, tail: []const u8, dtype: types.DType, shape: []const usize, why: ?*Why) !Tensor { + var buf: [320]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "{s}.{s}", .{ module, tail }); + const t = try store.get(name); + const size = std.math.mul(usize, shape[0], shape[1]) catch return error.UnexpectedTensor; + if (t.is(dtype, shape) and t.bytes.len == size * dtype.size()) return t; + var own: Why = .{}; + const w = why orelse &own; + w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); + if (why == null) std.log.err("{s}", .{w.text()}); + return error.UnexpectedTensor; +} + +/// `scale_inv` as f32 into `out`, exactly (f32 copied bit for bit); `finite` refuses NaN and infinity. +pub fn widen(scale_inv: Tensor, out: []f32, finite: bool) error{ UnexpectedTensor, UnsupportedQuantization }!void { + const dtype = scale_inv.dtype; + if (!convert.tableDtype(dtype) or out.len != scale_inv.numel() or scale_inv.bytes.len != out.len * dtype.size()) + return error.UnexpectedTensor; + if (dtype == .f32) @memcpy(std.mem.sliceAsBytes(out), scale_inv.bytes); // NaN payloads too + for (out, 0..) |*v, i| { + if (dtype != .f32) v.* = convert.load(dtype, scale_inv.bytes, i); + if (finite and !std.math.isFinite(v.*)) return error.UnsupportedQuantization; + } +} + +test { + _ = @import("fp8_block_test.zig"); +} + test "e4m3 codes and their rounding" { try std.testing.expectEqual(@as(f64, 1), e4m3(0x38)); try std.testing.expectEqual(@as(f64, -2), e4m3(0xc0)); diff --git a/zig/src/core/quant/fp8_block_test.zig b/zig/src/core/quant/fp8_block_test.zig new file mode 100644 index 00000000000..6cd7ba82ff1 --- /dev/null +++ b/zig/src/core/quant/fp8_block_test.zig @@ -0,0 +1,140 @@ +//! fp8 block: the shared block-FP8 reader (fp8.read) and its exact scale widening (fp8.widen), on the host. +const std = @import("std"); +const fp8 = @import("fp8.zig"); +const types = @import("types.zig"); +const Why = @import("../modelopt.zig").Why; + +const Tensor = types.Tensor; +const testing = std.testing; + +fn tensor(dtype: types.DType, shape: []const usize, bytes: []const u8) Tensor { + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +/// A store of named tensors, as safetensors maps, Drama 3's shards and the Flash-Next weight store answer `get`. +const Store = struct { + names: []const []const u8, + tensors: []const Tensor, + gets: usize = 0, + + pub fn get(s: *Store, name: []const u8) !Tensor { + s.gets += 1; + for (s.names, s.tensors) |n, t| if (std.mem.eql(u8, n, name)) return t; + return error.MissingTensor; + } +}; + +test "fp8 block: every bf16 bit pattern widens to its f32 value exactly, NaN and infinity kept without `finite`" { + var bits: [1 << 16]u16 = undefined; + for (&bits, 0..) |*b, i| b.* = @intCast(i); + var out: [1 << 16]f32 = undefined; + try fp8.widen(tensor(.bf16, &.{ 256, 256 }, std.mem.sliceAsBytes(&bits)), &out, false); + for (bits, out) |b, v| try testing.expectEqual(@as(u32, b) << 16, @as(u32, @bitCast(v))); +} + +test "fp8 block: f32 scales copy bit for bit, NaN payloads included" { + var prng = std.Random.DefaultPrng.init(0xf8b1); + var words: [4096]u32 = undefined; + prng.random().bytes(std.mem.sliceAsBytes(&words)); + words[0] = 0x7f800001; // a signalling NaN with a payload + words[1] = 0xffc12345; + var out: [4096]f32 = undefined; + try fp8.widen(tensor(.f32, &.{ 64, 64 }, std.mem.sliceAsBytes(&words)), &out, false); + try testing.expectEqualSlices(u32, &words, @ptrCast(&out)); +} + +test "fp8 block: every finite f16 bit pattern widens to its value exactly" { + var bits: [1 << 16]u16 = undefined; + for (&bits, 0..) |*b, i| b.* = @intCast(i); + var out: [1 << 16]f32 = undefined; + try fp8.widen(tensor(.f16, &.{1 << 16}, std.mem.sliceAsBytes(&bits)), &out, false); + for (bits, out) |b, v| { + const h: f16 = @bitCast(b); + if (std.math.isFinite(h)) try testing.expectEqual(@as(f64, @floatCast(h)), @as(f64, @floatCast(v))); + } +} + +test "fp8 block: `finite` refuses a NaN or infinite scale, and only then" { + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80, 0x7f81 }) |bad| { + var bits = [_]u16{ 0x3f80, bad, 0x4000 }; + var out: [3]f32 = undefined; + try testing.expectError(error.UnsupportedQuantization, fp8.widen(tensor(.bf16, &.{ 1, 3 }, std.mem.sliceAsBytes(&bits)), &out, true)); + try fp8.widen(tensor(.bf16, &.{ 1, 3 }, std.mem.sliceAsBytes(&bits)), &out, false); + } + const ok = [_]f32{ 0.5, 1e-30, -2 }; + var out: [3]f32 = undefined; + try fp8.widen(tensor(.f32, &.{ 1, 3 }, std.mem.sliceAsBytes(&ok)), &out, true); + try testing.expectEqualSlices(f32, &ok, &out); +} + +test "fp8 block: widen refuses a wrong output length and a dtype the kernels do not read" { + const two = [_]f32{ 1, 2 }; + var out: [3]f32 = undefined; + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f32, &.{ 1, 2 }, std.mem.sliceAsBytes(&two)), &out, false)); + const wide = [_]f64{1}; + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f64, &.{ 1, 1 }, std.mem.sliceAsBytes(&wide)), out[0..1], false)); + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f32, &.{ 1, 2 }, std.mem.sliceAsBytes(&two)[0..4]), out[0..2], false)); +} + +const codes: [130 * 260]u8 = @splat(0x38); +const inv: [2 * 3]f32 = .{ 1, 2, 3, 4, 5, 6 }; +const inv16: [2 * 3]u16 = .{ 1, 2, 3, 4, 5, 6 }; + +test "fp8 block: read returns the stored views of weight and weight_scale_inv, a partial last block counted" { + var s: Store = .{ + .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, + .tensors = &.{ tensor(.f8_e4m3, &.{ 130, 260 }, &codes), tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)) }, + }; + var why: Why = .{}; + const h = try fp8.read(&s, "m.q", 130, 260, .f32, &why); + try testing.expectEqual(@intFromPtr(&codes), @intFromPtr(h.weight.bytes.ptr)); + try testing.expectEqual(@intFromPtr(&inv), @intFromPtr(h.scale_inv.bytes.ptr)); + try testing.expectEqual(@as(usize, 2), s.gets); + try testing.expectEqual(@as(usize, 0), why.len); + try testing.expectEqual(@as(usize, 2), fp8.blocks(130)); + try testing.expectEqual(@as(usize, 1), fp8.blocks(128)); +} + +test "fp8 block: read refuses by name a wrong dtype, a wrong shape and short bytes; a missing tensor is the store's" { + const Case = struct { w: Tensor, s: Tensor, n: usize, k: usize, dtype: types.DType, text: []const u8 }; + const w = tensor(.f8_e4m3, &.{ 130, 260 }, &codes); + const s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)); + const cases = [_]Case{ + .{ .w = tensor(.u8, &.{ 130, 260 }, &codes), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: u8 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + .{ .w = w, .s = s, .n = 128, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 128, 260 }" }, + .{ .w = w, .s = s, .n = 130, .k = 260, .dtype = .bf16, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected bf16 { 2, 3 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 3, 2 }, std.mem.sliceAsBytes(&inv)), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 3, 2 }, expected f32 { 2, 3 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected f32 { 2, 3 }" }, + .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + }; + for (cases) |c| { + var store: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ c.w, c.s } }; + var why: Why = .{}; + try testing.expectError(error.UnexpectedTensor, fp8.read(&store, "m.q", c.n, c.k, c.dtype, &why)); + try testing.expectEqualStrings(c.text, why.text()); + } + var none: Store = .{ .names = &.{"m.q.weight"}, .tensors = &.{w} }; + var why: Why = .{}; + try testing.expectError(error.MissingTensor, fp8.read(&none, "m.q", 130, 260, .f32, &why)); + // bf16 scales of the same blocks read as stored + var b: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ w, tensor(.bf16, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv16)) } }; + _ = try fp8.read(&b, "m.q", 130, 260, .bf16, &why); +} + +test "fp8 block: read then widen gives reference()'s scales, block by block" { + var cs: [130 * 260]u8 = @splat(0x38); // 1.0 + cs[129 * 260 + 259] = 0x40; // 2.0 in the last, partial block + var s: Store = .{ .names = &.{ "e.weight", "e.weight_scale_inv" }, .tensors = &.{ tensor(.f8_e4m3, &.{ 130, 260 }, &cs), tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)) } }; + var why: Why = .{}; + const h = try fp8.read(&s, "e", 130, 260, .f32, &why); + var widened: [6]f32 = undefined; + try fp8.widen(h.scale_inv, &widened, true); + var ref: [130 * 260]f64 = undefined; + try fp8.reference(h, &ref); + for (0..130) |r| for (0..260) |c| { + const want = fp8.e4m3(cs[r * 260 + c]) * widened[(r / fp8.block) * 3 + c / fp8.block]; + try testing.expectEqual(want, ref[r * 260 + c]); + }; +} diff --git a/zig/src/families/drama3/cuda/weights.zig b/zig/src/families/drama3/cuda/weights.zig index 7a0a2d941c3..0f7182bb991 100644 --- a/zig/src/families/drama3/cuda/weights.zig +++ b/zig/src/families/drama3/cuda/weights.zig @@ -7,6 +7,7 @@ const cfg = @import("../config.zig"); const shards = @import("../shards.zig"); const Shards = shards.Shards; +const fp8 = core.quant.fp8; pub const Gdn = struct { proj: u64, conv: u64, a_log: u64, dt_bias: u64, norm: u64, out: u64 }; pub const Attn = struct { qkv: u64, q_norm: u64, k_norm: u64, o: u64 }; @@ -114,6 +115,18 @@ fn fmt(buf: []u8, comptime f: []const u8, args: anytype) []const u8 { return std.fmt.bufPrint(buf, f, args) catch unreachable; } +/// An expert projection's e4m3 codes by the shared block-FP8 reader; a refused one stays Drama's BadExpert. +fn codes(sh: *const Shards, module: []const u8, n: usize, k: usize, why: ?*core.modelopt.Why) ![]const u8 { + const t = fp8.weight(sh, module, n, k, why) catch |e| return if (e == error.UnexpectedTensor) error.BadExpert else e; + return t.bytes; +} + +/// Its f32 weight_scale_inv blocks copied into `out` bit for bit; a refused one stays Drama's BadExpertScale. +fn scalesInto(sh: *const Shards, module: []const u8, n: usize, k: usize, out: []f32, why: ?*core.modelopt.Why) !void { + const t = fp8.scales(sh, module, n, k, .f32, why) catch |e| return if (e == error.UnexpectedTensor) error.BadExpertScale else e; + fp8.widen(t, out, false) catch return error.BadExpertScale; +} + const PackJob = struct { sh: *const Shards, layer: usize, @@ -135,28 +148,24 @@ const PackJob = struct { const gpa = std.heap.c_allocator; const n = cfg.expert_width; const k = cfg.hidden; - var nb: [160]u8 = undefined; var e = j.first; while (e < cfg.experts) : (e += j.stride) { const base = "model.language_model.layers.{d}.mlp.experts.{d}."; - const g = try j.sh.get(fmt(&nb, base ++ "gate_proj.weight", .{ j.layer, e })); - const u = try j.sh.get(fmt(&nb, base ++ "up_proj.weight", .{ j.layer, e })); - const dn = try j.sh.get(fmt(&nb, base ++ "down_proj.weight", .{ j.layer, e })); - if (!g.is(.f8_e4m3, &.{ n, k }) or !u.is(.f8_e4m3, &.{ n, k }) or !dn.is(.f8_e4m3, &.{ k, n })) return error.BadExpert; + var mb: [3][160]u8 = undefined; + const mods = [3][]const u8{ fmt(&mb[0], base ++ "gate_proj", .{ j.layer, e }), fmt(&mb[1], base ++ "up_proj", .{ j.layer, e }), fmt(&mb[2], base ++ "down_proj", .{ j.layer, e }) }; + const g = try codes(j.sh, mods[0], n, k, null); + const u = try codes(j.sh, mods[1], n, k, null); + const dn = try codes(j.sh, mods[2], k, n, null); const per_up = n * k * 2 / 4; - try cuda.experts.Fp8Experts.packGateUp(gpa, j.up[e * per_up ..][0..per_up], g.bytes, u.bytes, n, k); + try cuda.experts.Fp8Experts.packGateUp(gpa, j.up[e * per_up ..][0..per_up], g, u, n, k); const per_down = k * n / 4; - cuda.experts.Fp8Experts.packOne(j.down[e * per_down ..][0..per_down], dn.bytes, k, n); - const gs = try j.sh.get(fmt(&nb, base ++ "gate_proj.weight_scale_inv", .{ j.layer, e })); - const us = try j.sh.get(fmt(&nb, base ++ "up_proj.weight_scale_inv", .{ j.layer, e })); - const ds = try j.sh.get(fmt(&nb, base ++ "down_proj.weight_scale_inv", .{ j.layer, e })); - if (!gs.is(.f32, &.{ n / 128, k / 128 }) or !us.is(.f32, &.{ n / 128, k / 128 }) or !ds.is(.f32, &.{ k / 128, n / 128 })) return error.BadExpertScale; + cuda.experts.Fp8Experts.packOne(j.down[e * per_down ..][0..per_down], dn, k, n); const per_us = 2 * (n / 128) * (k / 128); - const usd = std.mem.sliceAsBytes(j.up_scale[e * per_us ..][0..per_us]); - @memcpy(usd[0..gs.bytes.len], gs.bytes); - @memcpy(usd[gs.bytes.len..][0..us.bytes.len], us.bytes); + const us = j.up_scale[e * per_us ..][0..per_us]; + try scalesInto(j.sh, mods[0], n, k, us[0 .. per_us / 2], null); + try scalesInto(j.sh, mods[1], n, k, us[per_us / 2 ..], null); const per_ds = (k / 128) * (n / 128); - @memcpy(std.mem.sliceAsBytes(j.down_scale[e * per_ds ..][0..per_ds]), ds.bytes); + try scalesInto(j.sh, mods[2], k, n, j.down_scale[e * per_ds ..][0..per_ds], null); for ([_][]const u8{ "gate_proj.weight", "up_proj.weight", "down_proj.weight" }) |suffix| { var nb2: [160]u8 = undefined; const name = std.fmt.bufPrint(&nb2, "model.language_model.layers.{d}.mlp.experts.{d}.{s}", .{ j.layer, e, suffix }) catch unreachable; @@ -303,3 +312,70 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []co if (opt.log) std.log.info("drama3: weights on the device: {d:.2} GB", .{@as(f64, @floatFromInt(w.bytes)) / 1e9}); return w; } + +test "drama experts: block-FP8 experts pack through the shared reader, f32 scales bit for bit, Drama's refusals kept" { + const a = std.testing.allocator; + const io = std.testing.io; + const n = cfg.expert_width; + const k = cfg.hidden; + const sb = (n / 128) * (k / 128); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + var prng = std.Random.DefaultPrng.init(0xd3a3); + const pre = "model.language_model.layers.0.mlp.experts.1."; + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var codes_: [3][]u8 = undefined; + var inv: [3][sb]u32 = undefined; + var entries: [6]core.checkpoint_fixture.Entry = undefined; + var names: [6][96]u8 = undefined; + for (0..3) |p| { + codes_[p] = try a.alloc(u8, n * k); + prng.random().bytes(codes_[p]); + prng.random().bytes(std.mem.sliceAsBytes(&inv[p])); + const shape: []const usize = if (p == 2) &.{ k, n } else &.{ n, k }; + const blocks: []const usize = if (p == 2) &.{ k / 128, n / 128 } else &.{ n / 128, k / 128 }; + entries[2 * p] = .{ .name = fmt(&names[2 * p], pre ++ "{s}.weight", .{projs[p]}), .dtype = "F8_E4M3", .shape = shape, .bytes = codes_[p] }; + entries[2 * p + 1] = .{ .name = fmt(&names[2 * p + 1], pre ++ "{s}.weight_scale_inv", .{projs[p]}), .dtype = "F32", .shape = blocks, .bytes = std.mem.sliceAsBytes(&inv[p]) }; + } + defer for (codes_) |c| a.free(c); + inv[1][5] = 0x7f800123; // a NaN with a payload: Drama copies scales as stored, no finiteness rule + try core.checkpoint_fixture.write(a, io, tmp.dir, "model.safetensors", &entries); + var index: std.ArrayList(u8) = .empty; + defer index.deinit(a); + try index.appendSlice(a, "{\"weight_map\":{"); + for (entries, 0..) |e, i| try index.print(a, "{s}\"{s}\":\"model.safetensors\"", .{ if (i > 0) "," else "", e.name }); + try index.appendSlice(a, "}}"); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.items }); + const cwd = try std.process.currentPathAlloc(io, a); + defer a.free(cwd); + const dir = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); + defer a.free(dir); + var sh = try Shards.open(a, io, dir); + defer sh.close(io); + const per_up = n * k * 2 / 4; + const per_down = k * n / 4; + const up = try a.alloc(u32, 2 * per_up); + defer a.free(up); + const down = try a.alloc(u32, 2 * per_down); + defer a.free(down); + var up_scale: [2 * 2 * sb]f32 = undefined; + var down_scale: [2 * sb]f32 = undefined; + var job: PackJob = .{ .sh = &sh, .layer = 0, .first = 1, .stride = cfg.experts, .up = up, .down = down, .up_scale = &up_scale, .down_scale = &down_scale }; + try job.work(); + const want_up = try a.alloc(u32, per_up); + defer a.free(want_up); + try cuda.experts.Fp8Experts.packGateUp(a, want_up, codes_[0], codes_[1], n, k); + try std.testing.expectEqualSlices(u32, want_up, up[per_up..]); + const want_down = try a.alloc(u32, per_down); + defer a.free(want_down); + cuda.experts.Fp8Experts.packOne(want_down, codes_[2], k, n); + try std.testing.expectEqualSlices(u32, want_down, down[per_down..]); + try std.testing.expectEqualSlices(u32, &(inv[0] ++ inv[1]), @ptrCast(up_scale[2 * sb ..])); + try std.testing.expectEqualSlices(u32, &inv[2], @ptrCast(down_scale[sb..])); + // a weight or scale of the wrong shape keeps Drama's own error, the reason named by tensor + var why: core.modelopt.Why = .{}; + try std.testing.expectError(error.BadExpert, codes(&sh, pre ++ "down_proj", n, k, &why)); + try std.testing.expectEqualStrings(pre ++ "down_proj.weight: f8_e4m3 { 2048, 512 }, expected f8_e4m3 { 512, 2048 }", why.text()); + try std.testing.expectError(error.BadExpertScale, scalesInto(&sh, pre ++ "down_proj", n, k, down_scale[0..sb], &why)); + try std.testing.expectEqualStrings(pre ++ "down_proj.weight_scale_inv: f32 { 16, 4 }, expected f32 { 4, 16 }", why.text()); +} diff --git a/zig/src/families/drama3/drama3.zig b/zig/src/families/drama3/drama3.zig index 87a955bf15e..d929cb5b4af 100644 --- a/zig/src/families/drama3/drama3.zig +++ b/zig/src/families/drama3/drama3.zig @@ -43,4 +43,5 @@ test { test { _ = @import("cuda/owner_test.zig"); + _ = @import("cuda/weights.zig"); } diff --git a/zig/src/families/flashnext/cuda_fp8_experts_load.zig b/zig/src/families/flashnext/cuda_fp8_experts_load.zig index 3c79c3e52fe..0c4a0cf944c 100644 --- a/zig/src/families/flashnext/cuda_fp8_experts_load.zig +++ b/zig/src/families/flashnext/cuda_fp8_experts_load.zig @@ -1,8 +1,8 @@ //! The MTP head's block-FP8 routed experts: E4M3 weights, BF16 128x128 weight_scale_inv, packed for Fp8Experts (W8A16). const std = @import("std"); +const core = @import("core"); const cuda = @import("cuda"); const experts = @import("cuda_modelopt_experts.zig"); -const wt = @import("cuda_prompt_wt.zig"); const fx = cuda.experts.Fp8Experts; const blk = fx.block; @@ -78,15 +78,9 @@ pub fn packExpert(a: std.mem.Allocator, store: anytype, prefix: []const u8, i: u /// One projection's E4M3 weight [n, k]; its BF16 weight_scale_inv [n/128, k/128] widened (exactly) into `scale`. fn part(store: anytype, prefix: []const u8, i: usize, name: []const u8, n: usize, k: usize, scale: []f32) ![]const u8 { var buf: [512]u8 = undefined; - const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight", .{ prefix, i, name })); - if (!w.is(.f8_e4m3, &.{ n, k }) or w.bytes.len != n * k) return error.UnexpectedTensor; - const s = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight_scale_inv", .{ prefix, i, name })); - if (!s.is(.bf16, &.{ n / blk, k / blk }) or s.bytes.len != 2 * scale.len) return error.UnexpectedTensor; - for (scale, 0..) |*v, j| { - v.* = wt.promote(std.mem.readInt(u16, s.bytes[2 * j ..][0..2], .little)); - if (!std.math.isFinite(v.*)) return error.UnsupportedQuantization; - } - return w.bytes; + const h = try core.quant.fp8.read(store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}", .{ prefix, i, name }), n, k, .bf16, null); + try core.quant.fp8.widen(h.scale_inv, scale, true); + return h.weight.bytes; } test { diff --git a/zig/src/families/kolibri1/names.zig b/zig/src/families/kolibri1/names.zig index 361864d86fb..2f179a07895 100644 --- a/zig/src/families/kolibri1/names.zig +++ b/zig/src/families/kolibri1/names.zig @@ -6,7 +6,7 @@ const Config = @import("config.zig").Config; const Tensor = core.checkpoint.Tensor; const Checkpoint = core.checkpoint.Checkpoint; -pub const block = 128; // the checkpoint's scale block, both ways +pub const block = core.quant.fp8.block; // the checkpoint's scale block, both ways /// A block-FP8 projection as stored: e4m3 codes [n, k] and fp32 scale_inv [ceil(n/128), ceil(k/128)]. pub const Fp8 = struct { w: Tensor, s: Tensor }; @@ -42,10 +42,6 @@ pub const Map = struct { } }; -fn blocks(n: usize) usize { - return (n + block - 1) / block; -} - const Names = struct { ck: *Checkpoint, buf: [160]u8 = undefined, @@ -59,9 +55,8 @@ const Names = struct { } fn fp8(self: *Names, n: usize, k: usize, comptime fmt: []const u8, args: anytype) !Fp8 { - const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .f8_e4m3, &.{ n, k }); - const s = try self.ck.expect(try self.name(fmt ++ ".weight_scale_inv", args), .f32, &.{ blocks(n), blocks(k) }); - return .{ .w = w, .s = s }; + const h = try core.quant.fp8.read(self.ck, try self.name(fmt, args), n, k, .f32, null); + return .{ .w = h.weight, .s = h.scale_inv }; } fn expert(self: *Names, c: Config, comptime pre: []const u8, args: anytype) !Expert { @@ -130,7 +125,7 @@ fn writeTiny(a: std.mem.Allocator, tmp: testing.TmpDir) !void { const Add = struct { fn fp8(list: anytype, al: std.mem.Allocator, pre: []const u8, n: usize, k: usize) !void { try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k }) }); - try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_inv", .{pre}), .dtype = "F32", .shape = try al.dupe(usize, &.{ blocks(n), blocks(k) }) }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_inv", .{pre}), .dtype = "F32", .shape = try al.dupe(usize, &.{ core.quant.fp8.blocks(n), core.quant.fp8.blocks(k) }) }); } fn bf(list: anytype, al: std.mem.Allocator, name: []const u8, shape: []const usize) !void { try list.append(al, .{ .name = name, .dtype = "BF16", .shape = try al.dupe(usize, shape) }); diff --git a/zig/src/families/kolibri1/pack.zig b/zig/src/families/kolibri1/pack.zig index 8c93570cf65..a51c0975fc5 100644 --- a/zig/src/families/kolibri1/pack.zig +++ b/zig/src/families/kolibri1/pack.zig @@ -49,11 +49,6 @@ pub fn widen(gpa: std.mem.Allocator, t: Tensor) ![]f32 { return out; } -/// fp32 values of a tensor whose bytes may sit at any offset of the mapped file. -fn floats(dst: []f32, t: Tensor) void { - for (dst, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); -} - /// Row-stacked block-FP8 parts as one projection: codes and scales concatenated by rows, then packed. fn linear(gpa: std.mem.Allocator, parts: []const names.Fp8) !Linear { const k = parts[0].w.dim(1); @@ -73,7 +68,7 @@ fn linear(gpa: std.mem.Allocator, parts: []const names.Fp8) !Linear { for (parts) |p| { @memcpy(codes[at..][0..p.w.bytes.len], p.w.bytes); at += p.w.bytes.len; - floats(inv[sat..][0..p.s.numel()], p.s); + try core.quant.fp8.widen(p.s, inv[sat..][0..p.s.numel()], false); sat += p.s.numel(); } const npad = fp8.padded(n); @@ -98,9 +93,9 @@ fn experts(gpa: std.mem.Allocator, l: *Layer, ex: []const names.Expert, c: Confi for (ex, 0..) |e, i| { try fx8.packGateUp(gpa, l.up[i * one / 2 ..][0 .. one / 2], e.gate.w.bytes, e.up.w.bytes, ni, d); fx8.packOne(l.down[i * one / 4 ..][0 .. one / 4], e.down.w.bytes, d, ni); - floats(l.up_scale[2 * i * sb ..][0..sb], e.gate.s); - floats(l.up_scale[(2 * i + 1) * sb ..][0..sb], e.up.s); - floats(l.down_scale[i * sb ..][0..sb], e.down.s); + try core.quant.fp8.widen(e.gate.s, l.up_scale[2 * i * sb ..][0..sb], false); + try core.quant.fp8.widen(e.up.s, l.up_scale[(2 * i + 1) * sb ..][0..sb], false); + try core.quant.fp8.widen(e.down.s, l.down_scale[i * sb ..][0..sb], false); } } From 40a68f42e613f11711e34fc3062b0dcafd3eeff6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:33:51 +0100 Subject: [PATCH 107/186] core fp8: a byte-length refusal names the byte counts; widen's comment says an f16 NaN stays a NaN --- zig/src/core/quant/fp8.zig | 6 ++++-- zig/src/core/quant/fp8_block_test.zig | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/zig/src/core/quant/fp8.zig b/zig/src/core/quant/fp8.zig index 3d400788b39..f28169ddae9 100644 --- a/zig/src/core/quant/fp8.zig +++ b/zig/src/core/quant/fp8.zig @@ -99,12 +99,14 @@ fn checked(store: anytype, module: []const u8, tail: []const u8, dtype: types.DT if (t.is(dtype, shape) and t.bytes.len == size * dtype.size()) return t; var own: Why = .{}; const w = why orelse &own; - w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); + if (t.is(dtype, shape)) { + w.set("{s}: {d} bytes, {t} {any} needs {d}", .{ name, t.bytes.len, dtype, shape, size * dtype.size() }); + } else w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); if (why == null) std.log.err("{s}", .{w.text()}); return error.UnexpectedTensor; } -/// `scale_inv` as f32 into `out`, exactly (f32 copied bit for bit); `finite` refuses NaN and infinity. +/// `scale_inv` as f32 into `out`: exact (f32 bit for bit), an f16 NaN stays a NaN; `finite` refuses NaN and infinity. pub fn widen(scale_inv: Tensor, out: []f32, finite: bool) error{ UnexpectedTensor, UnsupportedQuantization }!void { const dtype = scale_inv.dtype; if (!convert.tableDtype(dtype) or out.len != scale_inv.numel() or scale_inv.bytes.len != out.len * dtype.size()) diff --git a/zig/src/core/quant/fp8_block_test.zig b/zig/src/core/quant/fp8_block_test.zig index 6cd7ba82ff1..9ea6e30f906 100644 --- a/zig/src/core/quant/fp8_block_test.zig +++ b/zig/src/core/quant/fp8_block_test.zig @@ -106,8 +106,8 @@ test "fp8 block: read refuses by name a wrong dtype, a wrong shape and short byt .{ .w = w, .s = s, .n = 128, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 128, 260 }" }, .{ .w = w, .s = s, .n = 130, .k = 260, .dtype = .bf16, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected bf16 { 2, 3 }" }, .{ .w = w, .s = tensor(.f32, &.{ 3, 2 }, std.mem.sliceAsBytes(&inv)), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 3, 2 }, expected f32 { 2, 3 }" }, - .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected f32 { 2, 3 }" }, - .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: 20 bytes, f32 { 2, 3 } needs 24" }, + .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: 33799 bytes, f8_e4m3 { 130, 260 } needs 33800" }, }; for (cases) |c| { var store: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ c.w, c.s } }; From b42b92f39f769d18b1d53143ebbcf7fc96bcc18c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:50 +0000 Subject: [PATCH 108/186] changelog: #2892 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b01c602295..e55d1906beb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. + - CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). - Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). From 8721d0341f8308b4a2625fd109548cfcb05d7471 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:26:07 +0100 Subject: [PATCH 109/186] perf(hip): the routed prompt tile takes K-parallel passes for an item of few rows On gfx1151 a routed prompt holds about 80 experts of 8 rows or fewer at any length. On the 64-row matrix tile each of their blocks reads 128 columns of codes a stage, a page of weights every four columns, and the cost of those reads follows the pages a block touches. wmma_item_tile now runs gemm_kp_wave for an item of at most the launch's cut: a lane a column's group, a column's whole row of codes read at once, no barrier. gemm_kp_wave is a wave's passes of gemm_kp_tile, factored out so both tiles run them; the K-parallel tiles keep their statement order and their register and scratch use. A pass of the item tile takes 2 rows and 2 columns a lane, which stays in registers, and keeps its partials in the block's code stages, so the tile's LDS is unchanged. The cut is a card's: `few_rows` in the tuning row (8 on gfx1151's item64 row, none elsewhere), carried to the kernel in Routing.few by prefillLaunch. A kernel test runs one plan through both paths at every width and compares bytes; the contract runs the tile with the cut and a control whose passes stop after a wave's first. The tile's record holds ROCm 10.1's hipcc. The build now tracks tiles/matrix_item.hpp. Refs #2795 --- zig/build/hip.zig | 18 ++--- zig/kernels/hip/contract_controls.hip | 13 ++++ zig/kernels/hip/quant/mlx_tiles.hpp | 11 ++- zig/kernels/hip/tiles/gemm_kp.hpp | 58 ++++++++++------ zig/kernels/hip/tiles/matrix_item.hpp | 30 +++++++- zig/kernels/hip/tiles/plan.hpp | 1 + zig/src/core/contract/hip_records_item.zig | 4 +- zig/src/core/registry.zig | 6 ++ zig/src/core/tuning.zig | 2 + zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/affine/cases.zig | 2 + zig/src/hip/kernel_tests/affine/item_few.zig | 72 ++++++++++++++++++++ zig/src/hip/kernel_tests/affine/product.zig | 9 ++- zig/src/hip/launches.zig | 1 + zig/src/hip/launches/affine.zig | 9 ++- zig/src/hip/launches/matrix_item.zig | 9 ++- zig/src/hip/tuning/gfx1151.zon | 4 +- zig/tests/hip/contract_main.zig | 10 ++- zig/tests/hip/contract_wmma_item.zig | 38 ++++++++--- 19 files changed, 245 insertions(+), 53 deletions(-) create mode 100644 zig/src/hip/kernel_tests/affine/item_few.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 70c3f26e360..88c0f6fb363 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -518,15 +518,15 @@ const groups = [_]Source{ }; /// What the group sources include, so an edit to one rebuilds them. const headers = [_][]const u8{ - "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", - "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", - "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", - "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", - "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", - "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", "flashnext/streams.hpp", - "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", "quant/bf16_decoder.hpp", - "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", "columns_controls.hpp", - "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", + "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", + "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", + "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", + "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", + "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", + "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/matrix_item.hpp", "tiles/plan.hpp", "tiles/stream.hpp", + "flashnext/streams.hpp", "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", + "quant/bf16_decoder.hpp", "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", + "columns_controls.hpp", "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. const flags = [_][]const u8{ "-D__HIP_PLATFORM_AMD__=1", "-DUSE_ROCM=1", "-DHIPBLAS_V2", "-fPIC", "-DCUDA_HAS_FP16=1", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", "-std=c++20", "-fno-gpu-rdc", "-mno-wavefrontsize64", "-ffp-contract=off" }; diff --git a/zig/kernels/hip/contract_controls.hip b/zig/kernels/hip/contract_controls.hip index 231f769c61c..327da668856 100644 --- a/zig/kernels/hip/contract_controls.hip +++ b/zig/kernels/hip/contract_controls.hip @@ -7,6 +7,7 @@ #include "attention/prefill.hpp" #include "quant/dense_tiles.hpp" +#include "quant/mlx_tiles.hpp" #include "tiles/wmma12_gemm_k32_ref.hpp" __device__ __forceinline__ float tf_ctl_x(const uint16_t* x, int64_t i) { @@ -196,6 +197,18 @@ tf_ctl_fa_paged_unmasked_bf16(const float* q, const hip_bfloat16* k, const hip_b 0, 0, 0, 0, table, count); } +// Gate 3: the routed item tile's few-row passes stop after a wave's first, leaving the block's other columns unwritten. +struct TfCtlItemFewOnePass { + using Dot = tf::rocm::DotBF16; + using Epi = tf::rocm::GemmOut; + static constexpr bool every_pass = false; +}; + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_item_few_one_pass(tf::rocm::Affine a) { + using namespace tf::rocm; + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, 4, TfCtlItemFewOnePass>(a); +} + // Gate 3: the big block loading its operands without the stores' swizzle, and without its K loop's barriers. struct TfCtlUnswizzled { static constexpr bool swizzled_loads = false, ordered = true; diff --git a/zig/kernels/hip/quant/mlx_tiles.hpp b/zig/kernels/hip/quant/mlx_tiles.hpp index a3f6ca68940..70bdd72d1e5 100644 --- a/zig/kernels/hip/quant/mlx_tiles.hpp +++ b/zig/kernels/hip/quant/mlx_tiles.hpp @@ -45,11 +45,18 @@ __global__ void __launch_bounds__(256) affine_wmma_skew(Affine a) { wmma_gemm_tile, BF16Act, WmmaBF16, GemmOut, true>(a); } -// ---- the same on RM x 16 rows a block, for a routed item of few rows ---- +// ---- the same on RM x 16 rows a block, for a routed item of few rows; K-parallel passes for one of very few ---- + +// The item tile's few-row passes as served: BF16 dot2 into the fp32 sums, every pass of a wave's columns. +struct ItemFew { + using Dot = DotBF16; + using Epi = GemmOut; + static constexpr bool every_pass = true; +}; template __global__ void __launch_bounds__(256) affine_wmma_item(Affine a) { - wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM>(a); + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM, ItemFew>(a); } // ---- the decode stream tile ---- diff --git a/zig/kernels/hip/tiles/gemm_kp.hpp b/zig/kernels/hip/tiles/gemm_kp.hpp index d759c7da59f..f4b7da96a48 100644 --- a/zig/kernels/hip/tiles/gemm_kp.hpp +++ b/zig/kernels/hip/tiles/gemm_kp.hpp @@ -12,26 +12,23 @@ namespace rocm { // The lanes a column takes, a group each: 32, or 16 or 8 when it has few groups (a wave then takes more columns). __host__ __device__ inline int kp_lanes(int groups) { return groups > 16 ? 32 : groups > 8 ? 16 : 8; } -// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. -template -__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { - const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows - if (!Dec::take_item(a, blockIdx.z)) return; +// A lane's row of each of a wave's LDS tables, what `lds(lane)` names: R * CB dots, R sums of x, CB scales, CB biases. +struct KpRows { + float* part; + float* spart; + float* scl; + float* bil; +}; + +// A wave's passes over its columns from `colw`, CB a lane and R rows a pass: rows from `first`, every `step` when LOOP. +template +__device__ __forceinline__ void gemm_kp_wave(typename Dec::Args& a, int colw, int first, int step, Lds lds) { constexpr int NO = R * CB; // outputs of a lane set - constexpr int PS = NO + 1; constexpr int OPL = (4 * NO + 31) / 32; // outputs a lane folds, at most const int groups = a.k / a.group; const int gp = kp_lanes(groups); // the lanes of a set; a wave has 32 / gp sets const int sets = 32 / gp; - const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; - if (colw >= a.n) return; - __shared__ float part[WAVES][32][PS]; - __shared__ float spart[WAVES][32][R + 1]; - __shared__ float scl[WAVES][32][CB + 1]; - __shared__ float bil[WAVES][32][CB + 1]; - const int lane = threadIdx.x & 31; - const int wave = threadIdx.x >> 5; const int per = a.group / 32; const int rounds = (groups + gp - 1) / gp; const int gi = lane & (gp - 1); @@ -128,14 +125,14 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { // this group's results to LDS, then the round's groups in order #pragma unroll for (int r = 0; r < R; ++r) { - spart[wave][lane][r] = sx[r]; + lds(lane).spart[r] = sx[r]; #pragma unroll - for (int c = 0; c < CB; ++c) part[wave][lane][r * CB + c] = d[r][c]; + for (int c = 0; c < CB; ++c) lds(lane).part[r * CB + c] = d[r][c]; } #pragma unroll for (int c = 0; c < CB; ++c) { - scl[wave][lane][c] = Dec::scale_value(a, term[c]); - bil[wave][lane][c] = Dec::bias_value(a, term[c]); + lds(lane).scl[c] = Dec::scale_value(a, term[c]); + lds(lane).bil[c] = Dec::bias_value(a, term[c]); } if (round + 1 < rounds) tables(g + gp); __syncwarp(); @@ -150,8 +147,8 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { const int c = o % CB; float v = acc[t]; for (int gg = 0; gg < count; ++gg) { - v = Dec::fold_scale(v, part[wave][base + gg][o], scl[wave][base + gg][c]); - v = Dec::fold_bias(v, spart[wave][base + gg][r], bil[wave][base + gg][c]); + v = Dec::fold_scale(v, lds(base + gg).part[o], lds(base + gg).scl[c]); + v = Dec::fold_bias(v, lds(base + gg).spart[r], lds(base + gg).bil[c]); } acc[t] = v; } @@ -168,11 +165,28 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { } }; if constexpr (LOOP) { - for (int row0 = (blockIdx.x % blocks) * R; row0 < a.m; row0 += blocks * R) pass(row0); + for (int row0 = first; row0 < a.m; row0 += step) pass(row0); } else { - pass((blockIdx.x % blocks) * R); + pass(first); } } +// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. +template +__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { + const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows + if (!Dec::take_item(a, blockIdx.z)) return; + const int sets = 32 / kp_lanes(a.k / a.group); + const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; + if (colw >= a.n) return; + __shared__ float part[WAVES][32][R * CB + 1]; + __shared__ float spart[WAVES][32][R + 1]; + __shared__ float scl[WAVES][32][CB + 1]; + __shared__ float bil[WAVES][32][CB + 1]; + const int wave = threadIdx.x >> 5; + const auto lds = [wave](int lane) { return KpRows{part[wave][lane], spart[wave][lane], scl[wave][lane], bil[wave][lane]}; }; + gemm_kp_wave(a, colw, (blockIdx.x % blocks) * R, blocks * R, lds); +} + } // namespace rocm } // namespace tf diff --git a/zig/kernels/hip/tiles/matrix_item.hpp b/zig/kernels/hip/tiles/matrix_item.hpp index 06195bb373e..734ccf5a056 100644 --- a/zig/kernels/hip/tiles/matrix_item.hpp +++ b/zig/kernels/hip/tiles/matrix_item.hpp @@ -4,13 +4,19 @@ #include "common/wmma.hpp" #include "tiles/gemm.hpp" +#include "tiles/gemm_kp.hpp" #include "tiles/plan.hpp" namespace tf { namespace rocm { +// The most rows a launch's cut sends to K-parallel passes, the rows a pass takes and the columns a lane takes in it. +constexpr int kItemFewRows = 8; +constexpr int kItemFewPass = 2; +constexpr int kItemFewCols = 2; + // Same staging, stage order and per-element sums as wmma_gemm_tile, so each output equals that tile's byte for byte. -template +template __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { #if TF_DEVICE_WMMA_GFX11 if (!Dec::take_item(a, blockIdx.z)) return; @@ -27,6 +33,28 @@ __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { const int tid = threadIdx.x; const int n0 = bx * kGemmN; + + // An item within the launch's cut takes K-parallel passes with Few's Dot and epilogue: same bytes, no matrix core. + if (a.m <= (a.route.few < kItemFewRows ? a.route.few : kItemFewRows)) { + constexpr int CB = kItemFewCols; + constexpr int R = kItemFewPass; + // a pass's partials, a row a lane: each wave's own slice of the two code stages, which this path leaves unused + constexpr int kLane = (R * CB + 1) + (R + 1) + 2 * (CB + 1); + static_assert(8 * 32 * kLane <= 2 * kGemmN * kLdw, "the few-row passes' partials fit the tile's code stages"); + const int wave = __builtin_amdgcn_readfirstlane(tid >> 5); + const auto lds = [wave](int lane) { + float* const row = reinterpret_cast(&ws[0][0]) + (wave * 32 + lane) * kLane; + return KpRows{row, row + R * CB + 1, row + R * CB + 1 + R + 1, row + R * CB + 1 + R + 1 + CB + 1}; + }; + // The 8 waves walk the block's columns side by side, CB a lane set a pass, R rows of the item a pass. + const int span = CB * (32 / kp_lanes(a.k / a.group)); + const int end = n0 + kGemmN < a.n ? n0 + kGemmN : a.n; + for (int colw = n0 + wave * span; colw < end; colw += 8 * span) { + gemm_kp_wave(a, colw, 0, R, lds); + if constexpr (!Few::every_pass) break; // the contract's control stops after a wave's first + } + return; + } const int words_row = Dec::template row_words(a); const int groups = a.k / a.group; const int stages = a.k / kGemmK; diff --git a/zig/kernels/hip/tiles/plan.hpp b/zig/kernels/hip/tiles/plan.hpp index 6cd355ba754..260d14b3b1e 100644 --- a/zig/kernels/hip/tiles/plan.hpp +++ b/zig/kernels/hip/tiles/plan.hpp @@ -10,6 +10,7 @@ struct Routing { const int* members = nullptr; int x_div = 1; int first = 0; // set by the kernel from its item + int few = 0; // an item of at most this many rows takes its tile's K-parallel passes, where the tile has them }; // Row r of this block's x and out: the plain matrix's r, or the routed item's r-th pair. diff --git a/zig/src/core/contract/hip_records_item.zig b/zig/src/core/contract/hip_records_item.zig index e7d154bb843..265209c83c7 100644 --- a/zig/src/core/contract/hip_records_item.zig +++ b/zig/src/core/contract/hip_records_item.zig @@ -18,8 +18,8 @@ pub const records = [1]Record{ .launch = itemTile(4), .usage = .{ .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, }, }, diff --git a/zig/src/core/registry.zig b/zig/src/core/registry.zig index 7d9895e8bed..1f07408a611 100644 --- a/zig/src/core/registry.zig +++ b/zig/src/core/registry.zig @@ -110,6 +110,12 @@ pub fn Registry(comptime Launch: type) type { return r; } + /// The tuned row of one of the registry's entries, or null for an entry this GPU's table does not name. + pub fn rowOf(r: *const Self, e: *const E) ?tuning.Row { + const i = (@intFromPtr(e) - @intFromPtr(r.entries.ptr)) / @sizeOf(E); + return if (i < r.entries.len) r.tuned[i] else null; + } + /// The cost of entry `i` at `shape`, or null where its row has no room for the shape. fn cost(r: *const Self, i: usize, shape: Shape, env: Env) ?f32 { const e = &r.entries[i]; diff --git a/zig/src/core/tuning.zig b/zig/src/core/tuning.zig index a24e2a346a5..45a8a169483 100644 --- a/zig/src/core/tuning.zig +++ b/zig/src/core/tuning.zig @@ -12,6 +12,8 @@ pub const Row = struct { items: u32 = 0, /// The most pairs a routed plan holds per item: more than `pairs_per_item * items` is refused (0: any). pairs_per_item: f32 = 0, + /// A routed prompt item of at most this many rows takes its kernel's K-parallel passes (0: none does). + few_rows: u32 = 0, plain: bool = true, routed: bool = true, }; diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index f5e44b272e0..14415a76989 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -67,6 +67,7 @@ test { _ = @import("kernel_tests/affine/affine.zig"); _ = @import("kernel_tests/affine/rows.zig"); _ = @import("kernel_tests/affine/order.zig"); + _ = @import("kernel_tests/affine/item_few.zig"); _ = @import("kernel_tests/affine/round.zig"); _ = @import("kernel_tests/affine/layout.zig"); _ = @import("kernel_tests/affine/golden.zig"); diff --git a/zig/src/hip/kernel_tests/affine/cases.zig b/zig/src/hip/kernel_tests/affine/cases.zig index f367933e7fa..cf68a09ced8 100644 --- a/zig/src/hip/kernel_tests/affine/cases.zig +++ b/zig/src/hip/kernel_tests/affine/cases.zig @@ -19,6 +19,8 @@ pub const Case = struct { item_rows: usize = 8, /// Every expert takes the same number of pairs. even: bool = false, + /// The pairs each expert takes, in expert order (they add up to rows x slots); empty draws an expert a pair. + loads: []const u16 = &.{}, /// Each token's last slot is the shared expert (index `experts`), the others distinct routed experts. shared: bool = false, /// The plan serving passes (`hip/moe.zig`): items cut at its tile, padded with empty items to its capacity. diff --git a/zig/src/hip/kernel_tests/affine/item_few.zig b/zig/src/hip/kernel_tests/affine/item_few.zig new file mode 100644 index 00000000000..86cdba36f3d --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/item_few.zig @@ -0,0 +1,72 @@ +//! The 64-row routed prompt tile's K-parallel passes write the matrix path's bytes for an item of few rows, at any cut. + +const std = @import("std"); +const base = @import("../rig.zig"); +const Buffer = base.Buffer; +const gpa = base.gpa; +const cases = @import("cases.zig"); +const prod = @import("product.zig"); +const Product = prod.Product; +const Rig = prod.Rig; +const Entry = @import("../../launches/affine.zig").Entry; +const item = @import("../../launches/matrix_item.zig"); + +/// Experts on both sides of every cut, of a row tile, of a block and of an item, then three of few rows at the end. +const loads = [_]u16{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17, 63, 64, 65, 128, 129, 200, 5, 8, 3 }; +const pairs = blk: { + var n: usize = 0; + for (loads) |l| n += l; + break :blk n; +}; + +/// The 35B's gate and up, its down, and a ragged product: a 72-column last block and a last round of one group. +const shapes = [_]cases.Case{ + .{ .name = "35b gate_up", .path = .prefill, .rows = pairs, .n = 1024, .k = 2048, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads }, + .{ .name = "35b down", .path = .prefill, .rows = pairs, .n = 2048, .k = 512, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads, .down = true }, + .{ .name = "ragged", .path = .prefill, .rows = pairs, .n = 200, .k = 2080, .group = 32, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads }, +}; + +/// What a run of outputs holds when no launch wrote it. +const unwritten: [16]u8 = @splat(0xa5); + +/// The tile's output on `p` with items of at most `few` rows on the K-parallel passes. +fn run(p: *const Product, e: *const Entry, out: Buffer, few: c_int) ![]u8 { + var a = p.arg(out.ptr, 0); + a.route.few = few; + try out.fill8(unwritten[0], null); + try e.launch(p.rig.kernels, .{ .r = &p.rig.base.r, .s = p.rig.base.stream.handle, .arg = a, .items = @intCast(p.items) }); + try p.rig.base.stream.synchronize(); + const bytes = try gpa.alloc(u8, p.outBytes(false)); + errdefer gpa.free(bytes); + try out.download(0, bytes); + return bytes; +} + +test "the routed item tile writes the matrix path's bytes through its K-parallel passes, at every cut and width" { + const b = try base.Rig.open(1 << 20); + defer b.close(); + var rig = Rig.open(b); + if (!rig.kernels.env().matrix) return error.SkipZigTest; + for (cases.widths) |bits| for (shapes) |shape| { + var c = shape; + c.bits = bits; + var p = try Product.init(&rig, c); + defer p.deinit(); + const e = for (rig.kernels.reg.entries) |*x| { + if (std.mem.endsWith(u8, x.id, ".routed.prefill.item64") and prod.takes(&p, x)) break x; + } else return error.NoItemTile; + var out = try Buffer.alloc(rig.base, p.outBytes(false)); + defer out.free(); + const matrix = try run(&p, e, out, 0); + defer gpa.free(matrix); + // every pair's output is written: no byte of the fill is left where an item of few rows sits + try base.expect(std.mem.indexOf(u8, matrix, &unwritten) == null, "{s} b{d}: the matrix path left outputs unwritten", .{ c.name, bits }); + // a cut below, at and past the rows a pass takes (the tile clamps the last) + for ([_]c_int{ 1, 3, item.few_rows, 64 }) |few| { + const passes = try run(&p, e, out, few); + defer gpa.free(passes); + const at = std.mem.indexOfDiff(u8, matrix, passes); + try base.expect(at == null, "{s} b{d} cut {d}: byte {?d} differs from the matrix path", .{ c.name, bits, few, at }); + } + }; +} diff --git a/zig/src/hip/kernel_tests/affine/product.zig b/zig/src/hip/kernel_tests/affine/product.zig index 24069c656be..2baeeaa27a3 100644 --- a/zig/src/hip/kernel_tests/affine/product.zig +++ b/zig/src/hip/kernel_tests/affine/product.zig @@ -101,7 +101,14 @@ pub const Product = struct { fn makePlan(p: *Product, pairs: usize) !void { const c = p.c; p.expert_of = try gpa.alloc(u32, pairs); - if (c.shared or c.served) { + if (c.loads.len != 0) { + var at: usize = 0; + for (c.loads, 0..) |load, e| { + @memset(p.expert_of[at..][0..load], @intCast(e)); + at += load; + } + std.debug.assert(at == pairs and c.loads.len <= c.experts); + } else if (c.shared or c.served) { drawServed(p.expert_of, c, &p.rig.rng); } else for (p.expert_of, 0..) |*e, i| { e.* = if (c.even) @intCast(i % c.experts) else @intCast(p.rig.rng.next() % c.experts); diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index 5d3c644de76..24b5aa308a2 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -10,6 +10,7 @@ const Stream = @import("stream.zig").Stream; const caps = @import("caps.zig"); const affine = @import("launches/affine.zig"); pub const dense = @import("launches/dense.zig"); +pub const item = @import("launches/matrix_item.zig"); const util = @import("launches/util.zig"); diff --git a/zig/src/hip/launches/affine.zig b/zig/src/hip/launches/affine.zig index 38d3d646731..a1b317a24b1 100644 --- a/zig/src/hip/launches/affine.zig +++ b/zig/src/hip/launches/affine.zig @@ -26,7 +26,7 @@ pub const Choices = struct { matrix: Matrix = .auto, stream: bool = true, gemm: /// quant/mlx.hpp's GroupTable, Routing and Affine as the kernels take them by value. pub const GroupTable = extern struct { p: u64, kind: c_int }; -pub const Routing = extern struct { items: u64 = 0, members: u64 = 0, x_div: c_int = 1, first: c_int = 0 }; +pub const Routing = extern struct { items: u64 = 0, members: u64 = 0, x_div: c_int = 1, first: c_int = 0, few: c_int = 0 }; pub const Arg = extern struct { x: u64, words: u64, @@ -49,7 +49,7 @@ pub const StreamSides = stream_layout.StreamSides; pub const Side = struct { words: u64, scale: u64, bias: u64, n: c_int, out: u64 }; comptime { - std.debug.assert(@sizeOf(GroupTable) == 16 and @sizeOf(Routing) == 24 and @sizeOf(Arg) == 112 and @sizeOf(StreamSides) == 184); + std.debug.assert(@sizeOf(GroupTable) == 16 and @sizeOf(Routing) == 32 and @sizeOf(Arg) == 120 and @sizeOf(StreamSides) == 184); } /// Everything a launch needs beside the kernels. @@ -457,7 +457,10 @@ pub const Kernels = struct { pub fn prefillLaunch(k: *const Kernels, r: *const runtime.Runtime, a: Arg, s: abi.Stream, items: c_int, pairs: c_int, expert_rows: c_int) Error!void { const op: registry.Op = if (a.route.items != 0) .routed else .project; const e = k.choose(op, .prefill, prefillShape(a, items, pairs, expert_rows)) orelse return refuse("prefill shape"); - try e.launch(k, .{ .r = r, .s = s, .arg = a, .items = items }); + var arg = a; + // this GPU's cut for a tile with K-parallel passes for items of few rows + if (k.reg.rowOf(e)) |row| arg.route.few = @intCast(row.few_rows); + try e.launch(k, .{ .r = r, .s = s, .arg = arg, .items = items }); } /// The product a prompt launch asks the registry about; a routed plan's pairs, and its average rows per expert. diff --git a/zig/src/hip/launches/matrix_item.zig b/zig/src/hip/launches/matrix_item.zig index b9be20377b4..3f9618d66a0 100644 --- a/zig/src/hip/launches/matrix_item.zig +++ b/zig/src/hip/launches/matrix_item.zig @@ -9,6 +9,8 @@ const util = @import("util.zig"); /// The 16-row tiles a block holds, by tier: 64 rows. pub const row_tiles = [_]usize{4}; +/// The most rows a tuning row's `few_rows` can cut at: what the tile sends to its K-parallel passes (kItemFewRows). +pub const few_rows = 8; pub const Functions = [affine.bit_widths.len][row_tiles.len]Function; /// The tier's functions at a width, as the module names them. @@ -28,9 +30,14 @@ pub fn rows(tier: usize) usize { /// Tier `tier` over every item of a plan: 128 columns by the tier's rows a block, all waves across the columns. pub fn go(k: *const affine.Kernels, c: affine.Call, tier: usize) util.Error!void { const b = std.mem.indexOfScalar(c_int, &affine.bit_widths, c.arg.bits) orelse return error.Invalid; + try goWith(k.item[b][tier], c, tier); +} + +/// `go` on a tile function of the tier's geometry: the tier's own, or a contract control in its place. +pub fn goWith(f: Function, c: affine.Call, tier: usize) util.Error!void { if (c.s == util.counting) return; var args: hl.Args = .{}; try args.add(c.arg); const grid = util.dim(util.cdiv(c.arg.n, 128), util.cdiv(c.arg.m, rows(tier)), c.items); - try hl.launch(k.item[b][tier], .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); + try hl.launch(f, .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); } diff --git a/zig/src/hip/tuning/gfx1151.zon b/zig/src/hip/tuning/gfx1151.zon index 175548d9394..cdecfe82346 100644 --- a/zig/src/hip/tuning/gfx1151.zon +++ b/zig/src/hip/tuning/gfx1151.zon @@ -7,9 +7,11 @@ // min_rows the fewest rows // items the most items of a routed plan (0: any) // pairs_per_item the most pairs a routed plan holds per item (0: any) +// few_rows a routed prompt item of at most this many rows takes the kernel's K-parallel passes (0: none) // plain, routed whether it takes plain products, routed plans (both by default) // gfx1100's rows, plus the 64-row routed matrix-core tile below 12 pairs an item (measured on gfx1151), // and the skewed-row matrix tile (three blocks a WGP) ahead of the padded one where it takes a product. +// The 64-row routed tile's items of 8 rows or fewer take its K-parallel passes (measured on gfx1151). .{ .gfx = "gfx1151", .entries = .{ @@ -43,7 +45,7 @@ .{ .id = "mlx.routed.prefill.kp3", .cost = 2, .items = 300, .plain = false }, .{ .id = "mlx.routed.prefill.kp1", .cost = 3, .rows = 4, .items = 300 }, .{ .id = "mlx.routed.prefill.kp2", .cost = 4, .rows = 32, .routed = false }, - .{ .id = "mlx.routed.prefill.item64", .cost = 9, .pairs_per_item = 12, .plain = false }, + .{ .id = "mlx.routed.prefill.item64", .cost = 9, .pairs_per_item = 12, .few_rows = 8, .plain = false }, .{ .id = "mlx.routed.prefill.skew", .cost = 9.5 }, .{ .id = "mlx.routed.prefill.matrix", .cost = 10 }, .{ .id = "mlx.routed.prefill.gemm", .cost = 11 }, diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 2f656e6d6c7..59416a7b740 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -298,10 +298,14 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&routers) |*x| try kernels.append(a, x.kernel()); // the few-row matrix-core tiles run on gfx11's matrix cores only const matrix = if (hip.Caps.of(arch)) |v| v.matrix == .wmma11 else false; - var items: [item_k.cases.len]item_k.Item = undefined; - for (&items, 0..) |*x, i| x.* = .init(a, &r, s, &l, i); + // the last one is the few-row passes' control, on its case's geometry + var items: [item_k.cases.len + 1]item_k.Item = undefined; + for (&items, 0..) |*x, i| x.* = .init(a, &r, s, &l, @min(i, item_k.control.index)); defer for (&items) |*x| x.deinit(); - if (matrix) for (&items) |*x| try kernels.append(a, x.kernel()); + if (matrix) { + items[item_k.cases.len].planted = try ctl.function(item_k.control.symbol); + for (&items) |*x| try kernels.append(a, x.kernel()); + } var embeds: [4]embed_k.Embed = undefined; var nembeds: usize = 0; defer for (embeds[0..nembeds]) |*x| x.deinit(); diff --git a/zig/tests/hip/contract_wmma_item.zig b/zig/tests/hip/contract_wmma_item.zig index f8e0b62bba1..33dbc3af8a7 100644 --- a/zig/tests/hip/contract_wmma_item.zig +++ b/zig/tests/hip/contract_wmma_item.zig @@ -13,18 +13,31 @@ const Kind = @typeInfo(@TypeOf(Kernels.blockGo)).@"fn".param_types[5].?; const Entry = @typeInfo(@FieldType(@FieldType(Kernels, "reg"), "entries")).pointer.child; const Call = @typeInfo(@typeInfo(@FieldType(Entry, "launch")).pointer.child).@"fn".param_types[1].?; -const Case = struct { tile: []const u8, rows: usize, group: usize, name: []const u8 }; +const Case = struct { tile: []const u8, rows: usize, group: usize, few: c_int, name: []const u8 }; const tiles = [_]struct { []const u8, usize }{.{ "item64", 64 }}; +/// The few-row cut a case's launches pass: none (the matrix path at every row count), then the tile's most. +const cuts = [_]c_int{ 0, hip.launches.item.few_rows }; pub const cases = blk: { - var out: [2 * tiles.len]Case = undefined; - for (tiles, 0..) |t, i| for ([_]usize{ 32, 64 }, 0..) |group, j| { - const name = std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out)", .{ t[1], group }); - out[2 * i + j] = .{ .tile = t[0], .rows = t[1], .group = group, .name = name }; + var out: [2 * cuts.len * tiles.len]Case = undefined; + for (tiles, 0..) |t, i| for (cuts, 0..) |few, f| for ([_]usize{ 32, 64 }, 0..) |group, j| { + const name: []const u8 = if (few == 0) + std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out)", .{ t[1], group }) + else + std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out, K-parallel passes to {d} rows)", .{ t[1], group, few }); + out[2 * (cuts.len * i + f) + j] = .{ .tile = t[0], .rows = t[1], .group = group, .few = few, .name = name }; }; break :blk out; }; +/// The planted control on the last case's geometry: the few-row passes stop after a wave's first. Red at gate 3. +pub const control = struct { + pub const index = cases.len - 1; + pub const symbol = "tf_ctl_item_few_one_pass"; + const name = "control: routed item tile whose few-row passes stop after a wave's first"; + const gate: g.Gate = .shapes; +}; + /// The 128-row tile's chain: a group's products by k, its scale and bias, then the groups' sum. fn depth(s: c.Shape) usize { return s.k / 2 + s.k / 32 + 8; @@ -39,6 +52,8 @@ pub const Item = struct { f: ?contract.affine.Affine = null, arg: Arg = undefined, base: u64 = 0, + /// The planted control (its kernel from the controls' object, on the tile's geometry) in place of the served one. + planted: ?hip.Function = null, pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, i: usize) Item { return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .arena = .init(a), .index = i }; @@ -54,7 +69,7 @@ pub const Item = struct { const flash: c.Shape = .{ .rows = 22, .n = 1280, .k = 2560 }; return .{ .spec = .{ - .name = cfg.name, + .name = if (self.planted != null) control.name else cfg.name, .accum = .f32, .depth = depth, .out = .f32, @@ -62,10 +77,12 @@ pub const Item = struct { .nan = .{ .stated = "the 128-row tile's bytes" }, // a prompt's routed items: about 22 rows at 1,024 tokens a step, 88 at 4,096 .served = &.{ flash, .{ .rows = 88, .n = 1280, .k = 2560 }, .{ .rows = 5, .n = 1024, .k = 2048 } }, - .tiles = .{ .{ .size = cfg.rows }, .{ .size = 128 }, .{ .size = 256, .step = cfg.group, .min = cfg.group } }, + // rows: the block's edge, or the cut's when the launch passes one + .tiles = .{ .{ .size = if (cfg.few > 0) @intCast(cfg.few) else cfg.rows }, .{ .size = 128 }, .{ .size = 256, .step = cfg.group, .min = cfg.group } }, }, .ctx = self, .vt = &vt, + .control = if (self.planted != null) control.gate else null, }; } @@ -96,13 +113,18 @@ pub const Item = struct { .group = @intCast(cfg.group), .fp16 = 0, }; + self.arg.route.few = cfg.few; return true; } pub fn run(self: *Item) !void { const tile = cases[self.index].tile; + const call: Call = .{ .r = self.l.r, .s = self.rig.s.handle, .arg = self.arg, .items = 1 }; + if (self.planted) |f| { + try hip.launches.item.goWith(f, call, 0); + return self.rig.s.synchronize(); + } for (self.l.affine.reg.entries) |*e| if (std.mem.startsWith(u8, e.id, "mlx.routed.prefill.") and std.mem.endsWith(u8, e.id, tile)) { - const call: Call = .{ .r = self.l.r, .s = self.rig.s.handle, .arg = self.arg, .items = 1 }; try e.launch(&self.l.affine, call); try self.rig.s.synchronize(); return; From fb85fe76dfccc48a4bcdde8fde7ed08ab2258af4 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:08:48 +0100 Subject: [PATCH 110/186] test(hip): pin the routed tile's cut from table to launch; rowOf checks bounds (#2878 review) The review's asks on the tile's cut for items of few rows: - registry: `rowOf` returns null for an entry before the registry's own instead of subtracting first, which trapped. A toy-registry test covers a named entry, an unnamed one and entries before and after it. - tuning: only `mlx.routed.prefill.item64` rows set `few_rows`, at no more than the tile's K-parallel passes take, and gfx1151's is set. - launches: `prefillLaunch` over a registry of toy entries carries its row's cut in `route.few`, whatever the caller's route held, and none for a row without one. It fails when the copy is lost. Refs #2795 --- zig/src/core/registry.zig | 24 ++++++++++++++-- zig/src/hip/launches/matrix_item.zig | 43 ++++++++++++++++++++++++++++ zig/src/hip/tuning/tuning.zig | 10 +++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/zig/src/core/registry.zig b/zig/src/core/registry.zig index 1f07408a611..e0ea439046e 100644 --- a/zig/src/core/registry.zig +++ b/zig/src/core/registry.zig @@ -110,9 +110,12 @@ pub fn Registry(comptime Launch: type) type { return r; } - /// The tuned row of one of the registry's entries, or null for an entry this GPU's table does not name. + /// The tuned row of entry `e`, or null for one this GPU's table does not name or the registry does not hold. pub fn rowOf(r: *const Self, e: *const E) ?tuning.Row { - const i = (@intFromPtr(e) - @intFromPtr(r.entries.ptr)) / @sizeOf(E); + const at = @intFromPtr(e); + const first = @intFromPtr(r.entries.ptr); + if (at < first) return null; + const i = (at - first) / @sizeOf(E); return if (i < r.entries.len) r.tuned[i] else null; } @@ -279,6 +282,23 @@ test "the cheapest entry that has room for the rows runs" { try std.testing.expect(std.mem.endsWith(u8, partial.select(toy_env, .mlx, .project, .decode, shape).?.id, "any")); } +test "an entry's tuned row is the table's, and null for an entry the table or the registry does not hold" { + const entries = [_]ToyEntry{ + toy("a.decode.before", .decode, .lanes, false), + toy("a.decode.few", .decode, .stream, true), + toy("a.decode.any", .decode, .lanes, false), + toy("a.decode.after", .decode, .lanes, false), + }; + const rows = [_]tuning.Row{.{ .id = "a.decode.few", .cost = 1, .few_rows = 8 }}; + // the registry holds the two entries in the middle, and the table names the first of them + const r = Toy.init(&.{ .gfx = "toy", .rows = &rows }, entries[1..3]); + try std.testing.expectEqual(@as(u32, 8), r.rowOf(&entries[1]).?.few_rows); + try std.testing.expect(r.rowOf(&entries[2]) == null); + // an entry before or after the registry's own has no row + try std.testing.expect(r.rowOf(&entries[0]) == null); + try std.testing.expect(r.rowOf(&entries[3]) == null); +} + test "a plan is chosen by its items, a plain product by its rows" { const entries = [_]ToyEntry{ toy("a.prefill.short", .prefill, .gemm, false), diff --git a/zig/src/hip/launches/matrix_item.zig b/zig/src/hip/launches/matrix_item.zig index 3f9618d66a0..0a6bd913830 100644 --- a/zig/src/hip/launches/matrix_item.zig +++ b/zig/src/hip/launches/matrix_item.zig @@ -6,6 +6,8 @@ const Module = @import("../module.zig").Module; const Function = @import("../module.zig").Function; const affine = @import("affine.zig"); const util = @import("util.zig"); +const registry = @import("core").registry; +const tuning = @import("../tuning/tuning.zig"); /// The 16-row tiles a block holds, by tier: 64 rows. pub const row_tiles = [_]usize{4}; @@ -41,3 +43,44 @@ pub fn goWith(f: Function, c: affine.Call, tier: usize) util.Error!void { const grid = util.dim(util.cdiv(c.arg.n, 128), util.cdiv(c.arg.m, rows(tier)), c.items); try hl.launch(f, .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); } + +// ---- the cut, from the tuning table to the launch + +var cut_launched: c_int = -1; + +fn keepCut(_: *const affine.Kernels, c: affine.Call) util.Error!void { + cut_launched = c.arg.route.few; +} + +fn anyEnv(_: registry.Env) bool { + return true; +} + +fn anyShape(_: registry.Shape) bool { + return true; +} + +fn toy(id: []const u8, op: registry.Op) affine.Entry { + return .{ .id = id, .format = .mlx, .op = op, .path = .prefill, .family = .gemm, .caps = &anyEnv, .policy = &anyEnv, .fits = &anyShape, .launch = &keepCut }; +} + +test "a prompt launch carries its entry's cut from the tuning table, and none for a row without one" { + const entries = [_]affine.Entry{ toy("toy.routed.prefill", .routed), toy("toy.project.prefill", .project) }; + const tuned = [_]tuning.Row{ .{ .id = "toy.routed.prefill", .cost = 1, .few_rows = few_rows }, .{ .id = "toy.project.prefill", .cost = 1 } }; + var k: affine.Kernels = undefined; + k.wmma = true; + k.matrix = true; + k.mode = .auto; + k.tile = .gemm; + k.stream_on = true; + k.reg = affine.Registry.init(&.{ .gfx = "toy", .rows = &tuned }, &entries); + // the caller's route holds another cut: the table's replaces it, for a plan and for a plain product + var a = std.mem.zeroes(affine.Arg); + a.m, a.n, a.k, a.bits, a.group = .{ 4, 128, 128, 4, 32 }; + a.route = .{ .items = 0x1000, .few = 3 }; + try k.prefillLaunch(undefined, a, util.counting, 2, 4, 2); + try std.testing.expectEqual(@as(c_int, few_rows), cut_launched); + a.route = .{ .few = 3 }; + try k.prefillLaunch(undefined, a, util.counting, 1, 0, 0); + try std.testing.expectEqual(@as(c_int, 0), cut_launched); +} diff --git a/zig/src/hip/tuning/tuning.zig b/zig/src/hip/tuning/tuning.zig index 4a636c328e6..b294625a2ba 100644 --- a/zig/src/hip/tuning/tuning.zig +++ b/zig/src/hip/tuning/tuning.zig @@ -19,6 +19,16 @@ test "only a routed prompt row bounds its pairs per item, by a finite number fro }; } +test "only the routed prompt's item tile cuts at few rows, at no more than its K-parallel passes take" { + const most = @import("../launches/matrix_item.zig").few_rows; + for ([_]*const Table{ &gfx1030, &gfx1100, &gfx1151, &gfx1201 }) |t| for (t.rows) |r| { + try std.testing.expect(r.few_rows <= most); + if (r.few_rows != 0) try std.testing.expectEqualStrings("mlx.routed.prefill.item64", r.id); + }; + // the cut is live where the receipts were taken: a table that lost it would pass the loop above + try std.testing.expect(find(&gfx1151, "mlx.routed.prefill.item64").?.few_rows != 0); +} + test "the tables name each entry once" { for ([_]*const Table{ &gfx1030, &gfx1100, &gfx1151, &gfx1201 }) |t| { for (t.rows, 0..) |r, i| { From ab78aef90456a93a45480d9bdc9713b6de6701ba Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:53 +0000 Subject: [PATCH 111/186] changelog: #2878 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e55d1906beb..3ee3b8becd0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. + - Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. - CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). From df79510351088285d45e183fb654d368bf806ac6 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:07:46 +0900 Subject: [PATCH 112/186] glm cuda: a split warm-up that NCCL fails refuses split by name before the rank stops Split's start-up warm-up sends and receives 4 MiB with every other rank, and runs before Engine.init's split probe, whose refusal names the fix. On a fabric where some ranks reach only their neighbours (a ring of four without a switch) the warm-up's group ends with an NCCL error first, and check stops every rank at once, so the probe's refusal never prints. From four ranks, where a ring has ranks that are not neighbours, the warm-up now prints that refusal, naming gather, when its group end fails under split, and hands check a name that starts with it, so a rank that stops on another rank's failure names the refusal as the cause too. Fewer ranks, ring and gather warm-ups, and a warm-up that finishes, print nothing new. --- zig/src/families/glm/cuda_warm_peers.zig | 17 ++++++++++++++++- zig/src/families/glm/cuda_warm_peers_test.zig | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/zig/src/families/glm/cuda_warm_peers.zig b/zig/src/families/glm/cuda_warm_peers.zig index 0be9a6e49d0..d95f0054531 100644 --- a/zig/src/families/glm/cuda_warm_peers.zig +++ b/zig/src/families/glm/cuda_warm_peers.zig @@ -111,7 +111,22 @@ pub fn post(g: *const cuda.tp_group.Group, stream: cuda.Stream, buffer: u64, pee } break :posting {}; }; - const ended = g.lib.check(g.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + const closed = g.lib.api.ncclGroupEnd(); + // split's warm-up is its first exchange with every rank; from four ranks a ring has some that are not neighbours + const refuse = peers == .every and g.world >= 4; + if (closed != 0 and refuse) refuseSplit(g.rank); // before check, which stops every rank and hands the others its name + const ended = g.lib.check(closed, if (refuse) "REFUSE TF_GLM_PREFILL_REDUCE=split at ncclGroupEnd" else "ncclGroupEnd"); try posted; try ended; } + +/// The latest rank that refused split, which a test build keeps instead of printing the refusal. +pub var refused_rank: ?usize = null; + +fn refuseSplit(rank: usize) void { + if (@import("builtin").is_test) { + refused_rank = rank; + return; + } + std.debug.print("REFUSE TF_GLM_PREFILL_REDUCE=split: rank {d} could not finish its start-up exchange with every other rank; split sends every rank's rows to every other, and a fabric where some ranks reach only their neighbours (a ring of four without a switch) needs TF_GLM_PREFILL_REDUCE=gather\n", .{rank}); +} diff --git a/zig/src/families/glm/cuda_warm_peers_test.zig b/zig/src/families/glm/cuda_warm_peers_test.zig index 6cea11161a9..50bd713ca99 100644 --- a/zig/src/families/glm/cuda_warm_peers_test.zig +++ b/zig/src/families/glm/cuda_warm_peers_test.zig @@ -139,6 +139,21 @@ test "GLM warm peers: failure at every posting closes only an opened group and s }; } +test "GLM warm peers: a split warm-up of four ranks or more that NCCL fails at its group's end refuses split by name first; fewer ranks, ring and gather do not" { + for (2..9) |world| for (0..world) |rank| for ([_]warm.Peers{ .every, .neighbours, .gather }) |peers| { + const g = Fake.group(rank, world); + Fake.reset(); + warm.refused_rank = null; + try Fake.run(&g, peers); + try t.expectEqual(@as(?usize, null), warm.refused_rank); + Fake.reset(); + if (peers == .gather) Fake.fail = 0 else Fake.end_fails = true; + try t.expectError(error.NcclFailed, Fake.run(&g, peers)); + try t.expectEqual(if (peers == .every and world >= 4) @as(?usize, rank) else null, warm.refused_rank); + }; + warm.refused_rank = null; +} + test "GLM warm peers: one rank does nothing and malformed placement refuses before NCCL" { Fake.reset(); var g = Fake.group(0, 1); From 05984eca0edf3a4f4dcc454fdeae1745efbd1813 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 16:10:36 +0900 Subject: [PATCH 113/186] changelog: #2868 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ee3b8becd0..807b24df04f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -190,6 +190,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: core gains one walk for copying kept prompt states, and GLM's copies on CUDA and Metal go through it, with the same copies in the same order; other families keep their own loops. Nothing a user sees changes (#2215, part of #1293). Thanks to @jschmied. +- GLM-5.3-Flash on CUDA at four ranks or more with the default split prefill: when the start-up warm-up's exchange with every rank fails, as on a switchless ring of four, the failing ranks print split's refusal naming `TF_GLM_PREFILL_REDUCE=gather` and the ranks that stop with them name it too; before, the ranks stopped without naming the setting (#2868). Thanks to @Bizuayeu. + - CUDA, opt-in build (`-Dnative-fp-decode`): qmmf's FP4 decode rows use the shared native e2m1 pair decode, `tf_common::e2m1_pairs`, which equals the existing decode on all 2^33 inputs; the default build's 113 fatbins are byte-identical and GB10 decode is level (#2721). Thanks to @jschmied. - Qwen3.8-27B on Metal: while a request runs, the keep-warm ticker now commits the residency sets' probe each interval, so the first command buffer after a wide prompt pass no longer waits in the kernel (likely for macOS to wire the model again; not traced). On an M5 Ultra with a 14 GiB prompt cache filling memory, five cold 24K- to 64K-token prompts in a row reach their first tokens 3.5% to 5.9% sooner in total; with the cache off, cold prompts are level. Replies are unchanged (#2725, fixes #2724). Thanks to @mrv777. From 1e8b2c3b8ea0dd29dd5fbd4107b2761dec3cea66 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:09:52 +0000 Subject: [PATCH 114/186] changelog: #2868's line from lines/2868.txt in place of its own --- CHANGELOG.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 807b24df04f..24ad3b85410 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. + - HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. - Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. @@ -99,7 +101,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. - - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. - Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. @@ -190,8 +191,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: core gains one walk for copying kept prompt states, and GLM's copies on CUDA and Metal go through it, with the same copies in the same order; other families keep their own loops. Nothing a user sees changes (#2215, part of #1293). Thanks to @jschmied. -- GLM-5.3-Flash on CUDA at four ranks or more with the default split prefill: when the start-up warm-up's exchange with every rank fails, as on a switchless ring of four, the failing ranks print split's refusal naming `TF_GLM_PREFILL_REDUCE=gather` and the ranks that stop with them name it too; before, the ranks stopped without naming the setting (#2868). Thanks to @Bizuayeu. - - CUDA, opt-in build (`-Dnative-fp-decode`): qmmf's FP4 decode rows use the shared native e2m1 pair decode, `tf_common::e2m1_pairs`, which equals the existing decode on all 2^33 inputs; the default build's 113 fatbins are byte-identical and GB10 decode is level (#2721). Thanks to @jschmied. - Qwen3.8-27B on Metal: while a request runs, the keep-warm ticker now commits the residency sets' probe each interval, so the first command buffer after a wide prompt pass no longer waits in the kernel (likely for macOS to wire the model again; not traced). On an M5 Ultra with a 14 GiB prompt cache filling memory, five cold 24K- to 64K-token prompts in a row reach their first tokens 3.5% to 5.9% sooner in total; with the cache off, cold prompts are level. Replies are unchanged (#2725, fixes #2724). Thanks to @mrv777. From 8d60e1ff00af3a01d9a27797ed4f40ef2cf2da46 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:53:09 +0000 Subject: [PATCH 115/186] refactor(radix): take a Query in the prompt store's begin and keep P2 of #2444 adds reuse identity (C6) beside the tree's words, so the store's begin and keep now take a Query instead of a bare prompt. The engine_api side carries every identity field with defaults and reads only the key for now; the core side carries the key and prompt length, which the CUDA adapter converts as it already does for Plan. Every caller passes the same words as before, so no store user changes behaviour. --- zig/src/core/kv_pool_store_test.zig | 4 +-- zig/src/core/kv_prompts.zig | 41 ++++++++++++++++-------- zig/src/core/prompt_radix.zig | 29 +++++++++++++++-- zig/src/core/prompt_radix_test.zig | 12 +++---- zig/src/families/dory/cuda_backend.zig | 4 +-- zig/src/families/gemma4/cuda_shared.zig | 4 +-- zig/src/families/glm/cuda_pages_test.zig | 16 ++++----- zig/src/families/glm/cuda_share.zig | 4 +-- zig/src/families/kolibri1/backend.zig | 4 +-- zig/src/families/qwen3_5/cuda_prefix.zig | 6 ++-- zig/src/families/spark25/backend.zig | 4 +-- zig/src/hip_serve/prefix.zig | 4 +-- zig/src/native/cuda_prompt_pages.zig | 17 ++++++---- 13 files changed, 96 insertions(+), 53 deletions(-) diff --git a/zig/src/core/kv_pool_store_test.zig b/zig/src/core/kv_pool_store_test.zig index 7260b1836a7..fa3f5b18123 100644 --- a/zig/src/core/kv_pool_store_test.zig +++ b/zig/src/core/kv_pool_store_test.zig @@ -79,7 +79,7 @@ const World = struct { fn admit(w: *World, a: Allocator, prompt: []const u32, history: u32, max_new: u32) !void { var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(w.gpa); - const plan = try w.store.begin(a, prompt, history, &.{}, null, &adopt); + const plan = try w.store.begin(a, .of(prompt), history, &.{}, null, &adopt); var l: Lease = .{}; errdefer w.pool.finish(&l); try w.pool.adopt(&l, adopt.items); @@ -89,7 +89,7 @@ const World = struct { for (plan.marks) |m| { var path: [64]u32 = undefined; const n = m / page; - const kept = w.store.keep(prompt, m, null, l.ids.items[0..n], path[0..n]); + const kept = w.store.keep(.of(prompt), m, null, l.ids.items[0..n], path[0..n]); for (path[0..kept.shared], 0..) |id, i| w.pool.share(&l, i, id); } try w.streams.append(w.gpa, l); diff --git a/zig/src/core/kv_prompts.zig b/zig/src/core/kv_prompts.zig index 19518d71b40..9f5f696336d 100644 --- a/zig/src/core/kv_prompts.zig +++ b/zig/src/core/kv_prompts.zig @@ -5,6 +5,19 @@ const kv = @import("kv_pool.zig"); /// Where a pass resumes and the marks it keeps states at (`marks` the store's allocation, freed by `free`). pub const Plan = struct { from: u32 = 0, marks: []const u32 = &.{} }; +/// What a pass asks the store about (prompt_radix.Query as a family on the core pool gives it). +pub const Query = struct { + /// The words the store compares: the request's cache key, one word a model row; for text, the prompt. + key: []const u32, + /// The prompt's length; a key of another length reuses nothing. + prompt_len: u32, + + /// A query of `key` alone, as long as its prompt. + pub fn of(key: []const u32) Query { + return .{ .key = key, .prompt_len = @intCast(key.len) }; + } +}; + /// The store as a family sees it: a prompt pass begins, keeps at its marks, and reports once its prompt is in. pub const Prompts = struct { ctx: *anyopaque, @@ -12,20 +25,20 @@ pub const Prompts = struct { pub const VTable = struct { /// The deepest state `prompt` resumes: its pages into `adopt` (held once more), the rest restored into `owner`. - begin: *const fn (ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, + begin: *const fn (ctx: *anyopaque, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, /// At `at` the store takes the pages of `mine` it lacks and the state; `path`'s first n replace the stream's. - keep: *const fn (ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, + keep: *const fn (ctx: *anyopaque, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, /// One log line after a prompt pass: where it resumed and how many states it kept. report: *const fn (ctx: *anyopaque, prompt: usize, from: u32, kept: u64) void, free: *const fn (ctx: *anyopaque, marks: []const u32) void, }; - pub fn begin(p: Prompts, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - return p.vtable.begin(p.ctx, prompt, history_len, shared, owner, adopt); + pub fn begin(p: Prompts, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + return p.vtable.begin(p.ctx, q, history_len, shared, owner, adopt); } - pub fn keep(p: Prompts, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return p.vtable.keep(p.ctx, prompt, at, owner, mine, path); + pub fn keep(p: Prompts, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return p.vtable.keep(p.ctx, q, at, owner, mine, path); } pub fn report(p: Prompts, prompt: usize, from: u32, kept: u64) void { @@ -37,8 +50,8 @@ pub const Prompts = struct { } /// `begin`, its marks in `gpa` for a lane to own; a failure gives `adopt`'s pages back. - pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - const got = try p.begin(prompt, history_len, shared, owner, adopt); + pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + const got = try p.begin(q, history_len, shared, owner, adopt); defer p.free(got.marks); const marks = gpa.dupe(u32, got.marks) catch |e| { for (adopt.items) |id| pool.release(id); @@ -48,13 +61,13 @@ pub const Prompts = struct { } /// At mark `at` the store takes the stream's pages and state; pages it held replace `l`'s, the whole table written. - pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, prompt: []const u32, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { + pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, q: Query, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { const n = at / page_tokens; const path = try gpa.alloc(u32, n); defer gpa.free(path); const mine = try gpa.dupe(u32, l.ids.items); defer gpa.free(mine); - const shared = p.keep(prompt, @intCast(at), owner, mine[0..n], path); + const shared = p.keep(q, @intCast(at), owner, mine[0..n], path); if (std.mem.eql(u32, mine[0..shared], path[0..shared])) return shared == n; @memcpy(mine[0..shared], path[0..shared]); // the same tokens' rows: a resume from any row is exact try table.write(table.ctx, mine, 0); @@ -163,10 +176,10 @@ const Stub = struct { fn of(ctx: *anyopaque) *Stub { return @ptrCast(@alignCast(ctx)); } - fn begin(_: *anyopaque, _: []const u32, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { + fn begin(_: *anyopaque, _: Query, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { return .{ .from = 256, .marks = try testing.allocator.dupe(u32, &.{512}) }; } - fn keep(ctx: *anyopaque, _: []const u32, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { + fn keep(ctx: *anyopaque, _: Query, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { @memcpy(path, Stub.of(ctx).held[0..mine.len]); return mine.len; } @@ -203,7 +216,7 @@ test "keeping swaps the store's pages in and writes the whole table: a table of const p: Prompts = .{ .ctx = &stub, .vtable = &Stub.vtable }; var w: Written = .{}; var owner: u8 = 0; - const whole = try p.keepPages(testing.allocator, &pool, &l, &.{ 1, 2, 3 }, 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); + const whole = try p.keepPages(testing.allocator, &pool, &l, .of(&.{ 1, 2, 3 }), 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); try testing.expect(whole); try testing.expectEqual(@as(usize, 0), w.from); try testing.expectEqual(@as(usize, 4), w.cap); @@ -221,7 +234,7 @@ test "a plan's marks are the lane's own; a blob's pages come from the pool and g var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); var owner: u8 = 0; - const got = try p.plan(testing.allocator, &pool, &.{1}, 1, &.{}, &owner, &adopt); + const got = try p.plan(testing.allocator, &pool, .of(&.{1}), 1, &.{}, &owner, &adopt); defer testing.allocator.free(got.marks); try testing.expectEqual(@as(u32, 256), got.from); try testing.expectEqualSlices(u32, &.{512}, got.marks); diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index e1fd5586e16..54720dc4da8 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -41,6 +41,29 @@ pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { } pub const Plan = pc.Plan; + +/// What a pass asks the store about: the words the tree compares, and what reuse must also match (enforced from C6). +pub const Query = struct { + /// The tree's words: the request's cache key, one word a model row; for text, the prompt. + key: []const u32, + /// The prompt's length; a key of another length reuses nothing. + prompt_len: u32, + /// The exact identities of the request's images. + images: []const @import("cache_images.zig").Identity = &.{}, + /// The request's image frame. + frame: u32 = 0, + /// What the rows depend on beyond key, images and frame. + namespace: u64 = 0, + /// The request's own chunk starts; empty: the store's grid. + starts: []const u32 = &.{}, + /// Rows this pass prefills; null: the prompt's length. + prefilled: ?u32 = null, + + /// A query of `key` alone, as long as its prompt. + pub fn of(key: []const u32) Query { + return .{ .key = key, .prompt_len = @intCast(key.len) }; + } +}; /// A paged store's `keep`: whether the mark holds a state, and how many leading pages the stream swaps for the store's. pub const Kept = struct { held: bool, shared: usize }; @@ -92,7 +115,8 @@ pub const Store = struct { } /// Restores the longest state `prompt` resumes into `owner`, its pages into `adopt` (held once more); plans marks. - pub fn begin(s: *Store, a: Allocator, prompt: []const u32, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { + pub fn begin(s: *Store, a: Allocator, q: Query, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { + const prompt = q.key; if (!s.on()) return .{ .marks = try a.alloc(u32, 0) }; const found = s.find(prompt); if (found == null) s.tree.counts.misses += 1; @@ -213,7 +237,8 @@ pub const Store = struct { } /// At `at` the tree takes the pages of `mine` it lacks and the snapshot, evicting to fit; `path` gets its pages. - pub fn keep(s: *Store, prompt: []const u32, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + const prompt = q.key; const t = &s.tree; if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != at / s.rules.page or path.len != mine.len) return .{ .held = false, .shared = 0 }; t.clock += 1; diff --git a/zig/src/core/prompt_radix_test.zig b/zig/src/core/prompt_radix_test.zig index 55bb93bf3ab..bcefb3b4af6 100644 --- a/zig/src/core/prompt_radix_test.zig +++ b/zig/src/core/prompt_radix_test.zig @@ -84,7 +84,7 @@ fn seq(comptime n: usize, from: u32) [n]u32 { /// A pass over `prompt` as a backend runs it: begin, then a keep at each mark, the stream taking the tree's pages. fn pass(f: *Fake, s: *Store, a: Allocator, prompt: []const u32, history: u32, shared: []const u32, mine: *std.ArrayList(u32)) !Plan { mine.clearRetainingCapacity(); - const plan = try s.begin(a, prompt, history, shared, null, mine); + const plan = try s.begin(a, .of(prompt), history, shared, null, mine); var fresh: [32]u32 = undefined; const more = (prompt.len + 3) / 4 - mine.items.len; f.lane(fresh[0..more]); @@ -92,7 +92,7 @@ fn pass(f: *Fake, s: *Store, a: Allocator, prompt: []const u32, history: u32, sh for (plan.marks) |m| { var path: [32]u32 = undefined; const n = m / 4; - const kept = s.keep(prompt, m, null, mine.items[0..n], path[0..n]); + const kept = s.keep(.of(prompt), m, null, mine.items[0..n], path[0..n]); for (path[0..kept.shared], 0..) |id, i| if (mine.items[i] != id) { f.refs[id] += 1; f.refs[mine.items[i]] -= 1; @@ -237,7 +237,7 @@ test "a state past the limits is refused, and a failed copy keeps no state" { // the same state again takes no second copy var path: [5]u32 = undefined; const held = [_]u32{ 0, 1, 2, 3, 4 }; - _ = s.keep(&t, 20, null, &held, &path); + _ = s.keep(.of(&t), 20, null, &held, &path); try std.testing.expectEqual(@as(usize, 1), f.live); // a restore that fails drops the entry and the pass starts from the start f.fail_restore = true; @@ -412,7 +412,7 @@ fn sharedAfter(branch: usize, lookahead: u32) !usize { var own: [5]u32 = undefined; f.lane(&own); var path: [5]u32 = undefined; - const kept = s.keep(&second, 20, null, &own, &path); + const kept = s.keep(.of(&second), 20, null, &own, &path); f.leave(&own); try std.testing.expect(kept.held); try std.testing.expectEqual(s.tree.snaps, f.live); @@ -461,9 +461,9 @@ test "planned: a keep off the grid is refused, one on it kept" { defer f.leave(&own); var path: [4]u32 = undefined; const prompt = seq(24, 0); - try std.testing.expect(!s.keep(&prompt, 12, null, own[0..3], path[0..3]).held); + try std.testing.expect(!s.keep(.of(&prompt), 12, null, own[0..3], path[0..3]).held); try std.testing.expectEqual(@as(usize, 0), s.tree.snaps); - try std.testing.expect(s.keep(&prompt, 16, null, &own, &path).held); + try std.testing.expect(s.keep(.of(&prompt), 16, null, &own, &path).held); try std.testing.expectEqual(@as(usize, 1), s.tree.snaps); } diff --git a/zig/src/families/dory/cuda_backend.zig b/zig/src/families/dory/cuda_backend.zig index 023e886fb50..6e84aaba8f6 100644 --- a/zig/src/families/dory/cuda_backend.zig +++ b/zig/src/families/dory/cuda_backend.zig @@ -242,7 +242,7 @@ pub const Dory = struct { x.restore_slot = slot; defer x.restore_slot = null; x.restored = 0; - const got = try pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + const got = try pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); if (got.from == x.restored) return got; for (adopt.items) |id| x.pool.?.release(id); x.gpa.free(got.marks); @@ -309,7 +309,7 @@ pub const Dory = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Dory, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), pages.tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), pages.tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row is in. `_: u32`. diff --git a/zig/src/families/gemma4/cuda_shared.zig b/zig/src/families/gemma4/cuda_shared.zig index b43232c5152..3dd1dde0c38 100644 --- a/zig/src/families/gemma4/cuda_shared.zig +++ b/zig/src/families/gemma4/cuda_shared.zig @@ -108,7 +108,7 @@ pub fn opening(comptime Lane: type, x: anytype, s: *Stream) !*Lane { var adopted = false; errdefer if (!adopted) for (adopt.items) |id| x.pool.?.release(id); const pr = if (s.images.len == 0) x.prompts else null; - const plan: core.kv_prompts.Plan = if (pr) |p| try p.begin(prompt, s.history_len, s.shared_prefixes, s, &adopt) else .{}; + const plan: core.kv_prompts.Plan = if (pr) |p| try p.begin(.of(prompt), s.history_len, s.shared_prefixes, s, &adopt) else .{}; var l: Lane = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .marks = plan.marks, .shared = pr != null }; errdefer { x.m.caps[slot] = 0; @@ -147,7 +147,7 @@ pub fn marked(x: anytype, s: *Stream, l: anytype, at: usize) !void { if (std.mem.indexOfScalar(u32, l.marks, @intCast(at)) != null) { const M = Map(@TypeOf(x)); var map: M = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, s, fwd.pages.tokens, .{ .ctx = &map, .write = M.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, s, fwd.pages.tokens, .{ .ctx = &map, .write = M.write })) l.kept += 1; } if (at == s.prompt().len) pr.report(at, s.cached, l.kept); } diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index 76de1632cfa..a84a958630a 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -424,13 +424,13 @@ const Store = struct { fn reclaim(ctx: *anyopaque, want: usize) bool { return of(ctx).store.reclaim(want); } - fn begin(ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!kp.Plan { + fn begin(ctx: *anyopaque, q: kp.Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!kp.Plan { const x = of(ctx); - const plan = try x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + const plan = try x.store.begin(x.gpa, .{ .key = q.key, .prompt_len = q.prompt_len }, history_len, shared, owner, adopt); return .{ .from = plan.from, .marks = plan.marks }; } - fn keep(ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return of(ctx).store.keep(prompt, at, owner, mine, path).shared; + fn keep(ctx: *anyopaque, q: kp.Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return of(ctx).store.keep(.{ .key = q.key, .prompt_len = q.prompt_len }, at, owner, mine, path).shared; } fn report(_: *anyopaque, _: usize, _: u32, _: u64) void {} fn free(ctx: *anyopaque, marks: []const u32) void { @@ -460,7 +460,7 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l var a: [3000]u32 = undefined; for (&a, 1..) |*t, i| t.* = @intCast(i); // a cold pass: a mark on the grid below its end, where the tree takes the pages below it and the state - const cold = try prompts.begin(&a, 0, &.{}, &owner, &adopt); + const cold = try prompts.begin(.of(&a), 0, &.{}, &owner, &adopt); defer prompts.free(cold.marks); try std.testing.expectEqualSlices(u32, &.{2048}, cold.marks); try expect(cold.from == 0 and owner.pending == null and adopt.items.len == 0); @@ -474,7 +474,7 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l // the same 2,049 tokens, another tail: the pass resumes at 2,048 on the tree's pages, its state pending var b = a; for (b[2049..]) |*t| t.* += 10_000; - const warm = try prompts.begin(&b, 0, &.{}, &owner, &adopt); + const warm = try prompts.begin(.of(&b), 0, &.{}, &owner, &adopt); defer prompts.free(warm.marks); try expectEqual(@as(u32, 2048), warm.from); try expectEqual(@as(u32, 2048), owner.pending.?.at); @@ -487,14 +487,14 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l var c = a; for (c[2048..]) |*t| t.* += 10_000; owner = .{}; - const miss = try prompts.begin(&c, 0, &.{}, &owner, &adopt); + const miss = try prompts.begin(.of(&c), 0, &.{}, &owner, &adopt); defer prompts.free(miss.marks); try expect(miss.from == 0 and owner.pending == null and adopt.items.len == 0); // a stream that needs the tree's pages: its promise evicts the state, and the pages are the stream's _ = try p.openWith(0, c.len, 100, &.{}); // 7 pages: the 4 free and 3 of the tree's try expectEqual(@as(usize, 0), states.live); p.close(0); - const gone = try prompts.begin(&a, 0, &.{}, &owner, &adopt); + const gone = try prompts.begin(.of(&a), 0, &.{}, &owner, &adopt); defer prompts.free(gone.marks); try expect(gone.from == 0 and owner.pending == null and adopt.items.len == 0); try expect(p.pool.idle() and p.pool.census().free == 8); diff --git a/zig/src/families/glm/cuda_share.zig b/zig/src/families/glm/cuda_share.zig index 365a1f8b27e..75ee6b282a8 100644 --- a/zig/src/families/glm/cuda_share.zig +++ b/zig/src/families/glm/cuda_share.zig @@ -79,7 +79,7 @@ pub const Owner = struct { pending: ?*engine.Kept = null }; /// Before a prompt pass: the deepest kept state's pages in `adopt` (held for the stream), the state in `owner`. pub fn begin(e: *engine.Engine, s: *const Stream, owner: *Owner, adopt: *std.ArrayList(u32)) !kp.Plan { const prompts = of(e) orelse return error.NoPromptStore; - const plan = try prompts.begin(cacheKey(s) orelse return error.ImageCacheKeyRequired, s.history_len, s.shared_prefixes, owner, adopt); + const plan = try prompts.begin(.of(cacheKey(s) orelse return error.ImageCacheKeyRequired), s.history_len, s.shared_prefixes, owner, adopt); defer prompts.free(plan.marks); const mine = e.gpa.dupe(u32, plan.marks) catch |err| { // the lane's own marks, freed by its allocator e.pages.?.give(adopt.items); @@ -111,7 +111,7 @@ pub const Keeps = struct { const path = k.e.gpa.alloc(u32, n) catch return; defer k.e.gpa.free(path); var owner: Owner = .{}; - _ = prompts.keep(k.prompt, mark, &owner, ids[0..n], path); + _ = prompts.keep(.of(k.prompt), mark, &owner, ids[0..n], path); } }; diff --git a/zig/src/families/kolibri1/backend.zig b/zig/src/families/kolibri1/backend.zig index f6fafb94e29..078e9fa13f5 100644 --- a/zig/src/families/kolibri1/backend.zig +++ b/zig/src/families/kolibri1/backend.zig @@ -250,7 +250,7 @@ pub const Kolibri = struct { fn plan(x: *Kolibri, s: *Stream, adopt: *std.ArrayList(u32)) !kp.Plan { const pr = x.prompts orelse return .{}; x.pending = null; - return pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + return pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); } /// A new lane's table starts with the kept pages, held once more by its lease (rows below `from`, never written). @@ -328,7 +328,7 @@ pub const Kolibri = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Kolibri, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row is in. `_: u32`. diff --git a/zig/src/families/qwen3_5/cuda_prefix.zig b/zig/src/families/qwen3_5/cuda_prefix.zig index b33ae5dee87..0d3e40ec4df 100644 --- a/zig/src/families/qwen3_5/cuda_prefix.zig +++ b/zig/src/families/qwen3_5/cuda_prefix.zig @@ -97,7 +97,7 @@ pub const Prefix = struct { /// Before a pass: the deepest kept boundary (pages in `adopt`, state in `owner.pending`) and the marks to keep. pub fn begin(x: *Prefix, prompt: []const u32, history_len: u32, shared: []const u32, owner: *Owner, adopt: *std.ArrayList(u32)) !radix.Plan { - return x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + return x.store.begin(x.gpa, .of(prompt), history_len, shared, owner, adopt); } /// The pending state copied into the stream's cache (its lane open and its pages adopted). @@ -121,7 +121,7 @@ pub const Prefix = struct { const mine = try x.gpa.dupe(u32, cache.paged.?.ids[0..n]); defer x.gpa.free(mine); var owner: Owner = .{ .cache = cache }; - const kept = x.store.keep(prompt, @intCast(at), &owner, mine, path); + const kept = x.store.keep(.of(prompt), @intCast(at), &owner, mine, path); try x.pages.swap(cache, path[0..kept.shared]); } @@ -130,7 +130,7 @@ pub const Prefix = struct { var owner: Owner = .{}; var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(x.gpa); - var plan = try x.store.begin(x.gpa, ids, history_len, shared, &owner, &adopt); + var plan = try x.store.begin(x.gpa, .of(ids), history_len, shared, &owner, &adopt); // the adopted pages are held for this call until openWith takes them plan.marks = x.onGrid(plan) catch |err| { for (adopt.items) |id| x.pages.pool.release(id); diff --git a/zig/src/families/spark25/backend.zig b/zig/src/families/spark25/backend.zig index 1da71fd557d..1b0c667461c 100644 --- a/zig/src/families/spark25/backend.zig +++ b/zig/src/families/spark25/backend.zig @@ -246,7 +246,7 @@ pub const Spark = struct { fn plan(x: *Spark, s: *Stream, adopt: *std.ArrayList(u32)) !kp.Plan { const pr = x.prompts orelse return .{}; x.pending = null; - return pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + return pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); } /// A new lane's table starts with the kept pages, held once more by its lease (rows below `from`, never written). @@ -325,7 +325,7 @@ pub const Spark = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Spark, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark. `_: u32` keeps its grid. diff --git a/zig/src/hip_serve/prefix.zig b/zig/src/hip_serve/prefix.zig index 092081b7f2f..26ce50605ef 100644 --- a/zig/src/hip_serve/prefix.zig +++ b/zig/src/hip_serve/prefix.zig @@ -68,7 +68,7 @@ pub fn Prefix(comptime F: type) type { o: *F.Owner, adopt: *std.ArrayList(u32), ) !radix.Plan { - return p.store.begin(p.gpa, prompt, history_len, shared, o, adopt); + return p.store.begin(p.gpa, .of(prompt), history_len, shared, o, adopt); } /// At `at` the tree takes the pages it lacks and the state; the stream swaps in the pages the tree had. @@ -78,7 +78,7 @@ pub fn Prefix(comptime F: type) type { const path = try p.gpa.alloc(u32, n); defer p.gpa.free(path); const table = F.table(o); - const kept = p.store.keep(prompt, @intCast(at), o, table[0..n], path); + const kept = p.store.keep(.of(prompt), @intCast(at), o, table[0..n], path); const own = path[0..kept.shared]; if (!std.mem.eql(u32, own, table[0..kept.shared])) { try F.adopt(p.f, o, own); diff --git a/zig/src/native/cuda_prompt_pages.zig b/zig/src/native/cuda_prompt_pages.zig index 3869d2a92fa..47b67828b7a 100644 --- a/zig/src/native/cuda_prompt_pages.zig +++ b/zig/src/native/cuda_prompt_pages.zig @@ -78,14 +78,19 @@ pub const Adapter = struct { return of(ctx).store.reclaim(want); } - fn begin(ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!hooks.Plan { + fn begin(ctx: *anyopaque, q: hooks.Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!hooks.Plan { const x = of(ctx); - const p = try x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + const p = try x.store.begin(x.gpa, query(q), history_len, shared, owner, adopt); return .{ .from = p.from, .marks = p.marks }; } - fn keep(ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return of(ctx).store.keep(prompt, at, owner, mine, path).shared; + fn keep(ctx: *anyopaque, q: hooks.Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return of(ctx).store.keep(query(q), at, owner, mine, path).shared; + } + + /// The family's query as the store takes it. + fn query(q: hooks.Query) radix.Query { + return .{ .key = q.key, .prompt_len = q.prompt_len }; } fn report(ctx: *anyopaque, prompt: usize, from: u32, kept: u64) void { @@ -141,12 +146,12 @@ test "TF_KV_SHARE's adapter: a kept prompt's pages resume a later one, adopted; try f.pool.cover(&lease, 3, null); var path: [2]u32 = undefined; var owner: u8 = 0; - try testing.expectEqual(@as(usize, 2), p.keep(&prompt, 8, &owner, lease.ids.items[0..2], &path)); + try testing.expectEqual(@as(usize, 2), p.keep(.of(&prompt), 8, &owner, lease.ids.items[0..2], &path)); try testing.expectEqualSlices(u32, lease.ids.items[0..2], &path); f.pool.finish(&lease); var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); - const plan = try p.begin(&prompt, 10, &.{}, &owner, &adopt); + const plan = try p.begin(.of(&prompt), 10, &.{}, &owner, &adopt); defer p.free(plan.marks); try testing.expectEqual(@as(u32, 8), plan.from); try testing.expectEqualSlices(u32, &path, adopt.items); From 3154901d1908a0263986ac3da5463e4eb2b81b23 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:59:32 +0000 Subject: [PATCH 116/186] refactor(radix): one store hold a page through a page registry The checkpoint manifest of #2444 (C3) lets several nodes and entries name one page, so the tree now holds its pages through Holds, a registry that takes one pool reference on a page's first name and gives it back on the last. The new prompt_radix_parts.zig also declares the manifest's types for the next steps. Node pages are distinct today, so the pool sees the same retains and releases in the same order and every store user is unchanged. --- zig/src/core/prompt_radix.zig | 9 ++- zig/src/core/prompt_radix_parts.zig | 96 ++++++++++++++++++++++++ zig/src/core/prompt_radix_parts_test.zig | 74 ++++++++++++++++++ zig/src/core/prompt_radix_test.zig | 6 +- zig/src/core/prompt_radix_tree.zig | 21 ++++-- zig/tests/qwen35/check_radix.zig | 4 +- 6 files changed, 194 insertions(+), 16 deletions(-) create mode 100644 zig/src/core/prompt_radix_parts.zig create mode 100644 zig/src/core/prompt_radix_parts_test.zig diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 54720dc4da8..663b2cdce45 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -100,7 +100,7 @@ pub const Store = struct { pub fn limit(s: *Store, limits: Limits) void { s.limits = limits; while (s.tree.snaps > limits.snaps) s.tree.evict(s.victim(&.{}, null) orelse break); - while (s.tree.held_pages > limits.pages) if (!s.tree.evictLeaf()) break; + while (s.tree.heldPages() > limits.pages) if (!s.tree.evictLeaf()) break; } /// Whether anything can be kept. @@ -168,7 +168,7 @@ pub const Store = struct { const last = marks_[marks_.len - 1]; while (true) { const new = last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead]); - if (s.tree.held_pages + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; + if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; s.tree.evict(s.victim(prompt, from) orelse return); } } @@ -254,7 +254,7 @@ pub const Store = struct { t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.held_pages + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + if (t.heldPages() + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); @@ -332,7 +332,7 @@ pub const Store = struct { pub fn report(s: *const Store, prompt: usize, from: u32, kept: u64) void { if (@import("builtin").is_test) return; const t = &s.tree; - std.log.info("prompt cache: {d} tokens, resumed at {d}, kept {d}; {d} states and {d} pages ({d} MiB) of {d} states and {d} pages (hits {d}, misses {d}, evicted {d}, refused {d}, failed {d})", .{ prompt, from, kept, t.snaps, t.held_pages, t.held() >> 20, s.limits.snaps, s.limits.pages, t.counts.hits, t.counts.misses, t.counts.evicted, t.counts.refused, t.counts.failed }); + std.log.info("prompt cache: {d} tokens, resumed at {d}, kept {d}; {d} states and {d} pages ({d} MiB) of {d} states and {d} pages (hits {d}, misses {d}, evicted {d}, refused {d}, failed {d})", .{ prompt, from, kept, t.snaps, t.heldPages(), t.held() >> 20, s.limits.snaps, s.limits.pages, t.counts.hits, t.counts.misses, t.counts.evicted, t.counts.refused, t.counts.failed }); } }; @@ -344,4 +344,5 @@ fn note(comptime fmt: []const u8, args: anytype) void { test { _ = @import("prompt_radix_test.zig"); + _ = @import("prompt_radix_parts.zig"); } diff --git a/zig/src/core/prompt_radix_parts.zig b/zig/src/core/prompt_radix_parts.zig new file mode 100644 index 00000000000..8985549a2f2 --- /dev/null +++ b/zig/src/core/prompt_radix_parts.zig @@ -0,0 +1,96 @@ +//! The radix store's manifest (C3, #2444): a family's cache groups, a checkpoint's parts, and one hold a page. +const std = @import("std"); +const pc = @import("prompt_cache.zig"); +const Allocator = std.mem.Allocator; + +/// What a cache group holds; the core pool's kinds (kv_groups.zig), named here for the store. +pub const Kind = enum { full, window, state, draft }; + +/// How a checkpoint group's part is made: the stream's own read-only pages, or store-held pages its bytes go to. +pub const Make = enum { retained, copied }; + +/// A checkpoint group: its kind, how its part is made, and (retained) its rows a page. +pub const PartSpec = struct { kind: Kind, make: Make, rows: u32 = 0 }; + +/// A family's cache groups as its store sees them, declared once when the store opens. +pub const Groups = struct { + /// Rows a page of each prefix group; empty: one prefix group of the store's page. + prefix_rows: []const u32 = &.{}, + /// The checkpoint groups, in order. + parts: []const PartSpec = &.{}, + /// A legacy snapshot (the family's Snapshots) beside the parts. + legacy: bool = true, + + /// Prefix groups: page ids a logical page on the tree's path. + pub fn prefixCount(g: Groups) usize { + return @max(1, g.prefix_rows.len); + } + + /// No checkpoint at all (pure full attention): any matched page edge resumes. + pub fn none(g: Groups) bool { + return g.parts.len == 0 and !g.legacy; + } +}; + +/// One checkpoint group's pages in one entry; page `i` sits at logical index `base + i` in the group's table. +pub const Part = struct { base: u32, ids: []u32, bytes: u64 }; + +/// A keep's input for one retained group: its pages from logical index `base`, the last partial if rows fall short. +pub const Retained = struct { base: u32, ids: []const u32 }; + +/// A copy a stream makes from a held page into one of its own: rows [0, rows) of `src` to the group's page `index`. +pub const Copy = struct { group: u32, index: u32, src: u32, rows: u32 }; + +/// The store's references: one pool reference a page, however many nodes and entries name it. +pub const Holds = struct { + counts: std.AutoHashMapUnmanaged(u32, u32) = .empty, + + pub fn deinit(h: *Holds, gpa: Allocator) void { + h.counts.deinit(gpa); + } + + /// Room for `n` more pages, so the registrations after it cannot fail. + pub fn reserve(h: *Holds, gpa: Allocator, n: usize) error{OutOfMemory}!void { + try h.counts.ensureUnusedCapacity(gpa, @intCast(n)); + } + + /// The store names a page a stream holds: its first name takes one pool reference. + pub fn shareAssumeCapacity(h: *Holds, pages: pc.Pages, id: u32) void { + const slot = h.counts.getOrPutAssumeCapacity(id); + if (slot.found_existing) { + slot.value_ptr.* += 1; + return; + } + slot.value_ptr.* = 1; + pages.vtable.retain(pages.ptr, id); + } + + /// The store names a page it took itself (Pages.hold): that hold is its pool reference. + pub fn adoptHeldAssumeCapacity(h: *Holds, id: u32) void { + const slot = h.counts.getOrPutAssumeCapacity(id); + std.debug.assert(!slot.found_existing); // a page just held is new to the store + slot.value_ptr.* = 1; + } + + /// One name fewer; the last gives the pool reference back. + pub fn drop(h: *Holds, pages: pc.Pages, id: u32) void { + const slot = h.counts.getPtr(id) orelse std.debug.panic("prompt cache: page {d} dropped while not held", .{id}); + slot.* -= 1; + if (slot.* > 0) return; + _ = h.counts.remove(id); + pages.vtable.release(pages.ptr, id); + } + + pub fn contains(h: *const Holds, id: u32) bool { + return h.counts.contains(id); + } + + /// Distinct pages the store holds: what its page limit counts. + pub fn count(h: *const Holds) usize { + return h.counts.count(); + } +}; + +test { + _ = @import("prompt_radix_parts_test.zig"); +} diff --git a/zig/src/core/prompt_radix_parts_test.zig b/zig/src/core/prompt_radix_parts_test.zig new file mode 100644 index 00000000000..1341a6ad546 --- /dev/null +++ b/zig/src/core/prompt_radix_parts_test.zig @@ -0,0 +1,74 @@ +//! The manifest's registry alone: one pool reference a page, whatever names it, given back on the last drop. +const std = @import("std"); +const parts = @import("prompt_radix_parts.zig"); +const pc = @import("prompt_cache.zig"); +const testing = std.testing; + +/// Pages whose holders are counted, as a pool's are. +const Refs = struct { + refs: [8]u32 = @splat(0), + calls: u32 = 0, // retains and releases seen + + fn of(p: *anyopaque) *Refs { + return @ptrCast(@alignCast(p)); + } + fn retain(p: *anyopaque, id: u32) void { + of(p).refs[id] += 1; + of(p).calls += 1; + } + fn release(p: *anyopaque, id: u32) void { + of(p).refs[id] -= 1; + of(p).calls += 1; + } + fn holders(p: *anyopaque, id: u32) u32 { + return of(p).refs[id]; + } + fn available(_: *anyopaque) usize { + return 0; + } + fn pages(r: *Refs) pc.Pages { + return .{ .ptr = r, .bytes = 1, .vtable = &.{ .retain = retain, .release = release, .holders = holders, .available = available } }; + } +}; + +test "the store's registry: a page named twice holds one reference, given back on its last drop" { + var r: Refs = .{}; + r.refs[3] = 1; // a stream's page + var h: parts.Holds = .{}; + defer h.deinit(testing.allocator); + try h.reserve(testing.allocator, 2); + h.shareAssumeCapacity(r.pages(), 3); + h.shareAssumeCapacity(r.pages(), 3); // a second entry names it + try testing.expectEqual(@as(u32, 2), r.refs[3]); + try testing.expectEqual(@as(usize, 1), h.count()); + h.drop(r.pages(), 3); + try testing.expectEqual(@as(u32, 2), r.refs[3]); // still named once + h.drop(r.pages(), 3); + try testing.expectEqual(@as(u32, 1), r.refs[3]); + try testing.expectEqual(@as(usize, 0), h.count()); + try testing.expectEqual(@as(u32, 2), r.calls); // one retain, one release +} + +test "the store's registry: a page it held itself is its reference, released once" { + var r: Refs = .{}; + r.refs[5] = 1; // Pages.hold's reference + var h: parts.Holds = .{}; + defer h.deinit(testing.allocator); + try h.reserve(testing.allocator, 1); + h.adoptHeldAssumeCapacity(5); + try testing.expect(h.contains(5)); + try testing.expectEqual(@as(u32, 1), r.refs[5]); // no retain + h.drop(r.pages(), 5); + try testing.expectEqual(@as(u32, 0), r.refs[5]); +} + +test "a family's groups: one prefix group by default, a page-only store has no checkpoint" { + const d: parts.Groups = .{}; + try testing.expectEqual(@as(usize, 1), d.prefixCount()); + try testing.expect(!d.none()); + const two: parts.Groups = .{ .prefix_rows = &.{ 4, 4 }, .legacy = false }; + try testing.expectEqual(@as(usize, 2), two.prefixCount()); + try testing.expect(two.none()); + const state: parts.Groups = .{ .parts = &.{.{ .kind = .state, .make = .copied }}, .legacy = false }; + try testing.expect(!state.none()); +} diff --git a/zig/src/core/prompt_radix_test.zig b/zig/src/core/prompt_radix_test.zig index bcefb3b4af6..761bb229424 100644 --- a/zig/src/core/prompt_radix_test.zig +++ b/zig/src/core/prompt_radix_test.zig @@ -276,7 +276,7 @@ test "reclaim frees the pages of the least valuable leaf nobody holds, and spare try std.testing.expectEqual(@as(usize, 1), s.tree.snaps); for (y_pages) |id| f.refs[id] -= 1; s.limit(.{ .pages = 40, .snaps = 0 }); - try std.testing.expectEqual(@as(usize, 0), s.tree.held_pages); + try std.testing.expectEqual(@as(usize, 0), s.tree.heldPages()); try std.testing.expectEqual(@as(usize, 0), f.live); } @@ -333,7 +333,7 @@ fn goldenRun(a: Allocator, f: *Fake, s: *Store, out: *std.ArrayList(Row)) !void for (passes) |p| { const plan = try pass(f, s, a, p.prompt, p.history, p.shared, &mine); const c = s.tree.counts; - try out.append(a, .{ .from = plan.from, .marks = plan.marks, .snaps = s.tree.snaps, .pages = s.tree.held_pages, .bytes = s.tree.held(), .hits = c.hits, .misses = c.misses, .kept = c.kept, .evicted = c.evicted, .refused = c.refused }); + try out.append(a, .{ .from = plan.from, .marks = plan.marks, .snaps = s.tree.snaps, .pages = s.tree.heldPages(), .bytes = s.tree.held(), .hits = c.hits, .misses = c.misses, .kept = c.kept, .evicted = c.evicted, .refused = c.refused }); f.leave(mine.items); } } @@ -502,7 +502,7 @@ test "GLM's shape, a grid of four pages and lookahead 1: a chat resumes at its h try std.testing.expectEqualSlices(u32, &.{48}, plan.marks); f.leave(mine.items); try std.testing.expectEqual(@as(usize, 2), s.tree.snaps); - try std.testing.expectEqual(@as(usize, 12), s.tree.held_pages); + try std.testing.expectEqual(@as(usize, 12), s.tree.heldPages()); try std.testing.expectEqual(s.tree.snaps, f.live); } diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index 2c6770ee58f..3eff0969810 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -4,6 +4,7 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const pc = @import("prompt_cache.zig"); +const parts = @import("prompt_radix_parts.zig"); pub const Saved = pc.Saved; pub const Snapshots = pc.Snapshots; @@ -50,8 +51,9 @@ pub const Tree = struct { ahead: usize, root: *Node, clock: u64 = 0, - /// Pages and snapshots the tree holds. - held_pages: usize = 0, + /// The pages the tree holds, one pool reference each however many nodes name them. + holds: parts.Holds = .{}, + /// Snapshots the tree holds. snaps: usize = 0, counts: Counts = .{}, @@ -63,6 +65,12 @@ pub const Tree = struct { pub fn deinit(t: *Tree) void { t.drop(t.root); + t.holds.deinit(t.gpa); + } + + /// Distinct pages the tree holds: what its page limit counts. + pub fn heldPages(t: *const Tree) usize { + return t.holds.count(); } fn drop(t: *Tree, n: *Node) void { @@ -72,7 +80,7 @@ pub const Tree = struct { t.family.vtable.drop(t.family.ptr, e.saved); t.gpa.free(e.last); } - for (n.pages) |id| t.pages.vtable.release(t.pages.ptr, id); + for (n.pages) |id| t.holds.drop(t.pages, id); t.gpa.free(n.tokens); t.gpa.free(n.pages); t.gpa.destroy(n); @@ -163,10 +171,10 @@ pub const Tree = struct { leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi..]) }; errdefer t.gpa.free(leaf.tokens); errdefer t.gpa.free(leaf.pages); + try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); - for (leaf.pages) |id| t.pages.vtable.retain(t.pages.ptr, id); + for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); @memcpy(path[pi..], leaf.pages); - t.held_pages += leaf.pages.len; t.touch(leaf); return leaf; }; @@ -224,7 +232,6 @@ pub const Tree = struct { if (n.kids.items.len > 0 or n.entry != null) return; const slot = std.mem.indexOfScalar(*Node, parent.kids.items, n).?; _ = parent.kids.swapRemove(slot); - t.held_pages -= n.pages.len; t.drop(n); n = parent; } @@ -260,7 +267,7 @@ pub const Tree = struct { /// Bytes the tree holds: its pages and its snapshots. pub fn held(t: *const Tree) u64 { - var n: u64 = t.held_pages * t.pages.bytes; + var n: u64 = t.heldPages() * t.pages.bytes; sum(t.root, &n); return n; } diff --git a/zig/tests/qwen35/check_radix.zig b/zig/tests/qwen35/check_radix.zig index f715078b268..a970125a9f2 100644 --- a/zig/tests/qwen35/check_radix.zig +++ b/zig/tests/qwen35/check_radix.zig @@ -161,11 +161,11 @@ fn memoryLine(c: *Ctx) void { const done = c.session.run(c.arena, .{ .prompts = set, .max_new = 1, .solo = true }) catch |err| return c.report.broke(name, err); const tree = &h.prefix.store.tree; const copied = copiedBytes(c, set, done.replies); - const unique = pagesBytes(c, tree.held_pages); + const unique = pagesBytes(c, tree.heldPages()); const states = tree.held() - unique; var cached: usize = 0; for (done.replies) |r| cached += r.cached; - c.report.pass(name, "{d} requests sharing {d} system tokens, {d} tokens each: the tree holds {d:.1} MiB ({d} pages {d:.1} MiB, {d} snapshots {d:.1} MiB); copies of the caches at each cut held {d:.1} MiB; {d} prompt tokens came from the tree", .{ set.len, system_len, set[0].ids.len, mib(tree.held()), tree.held_pages, mib(unique), tree.snaps, mib(states), mib(copied), cached }); + c.report.pass(name, "{d} requests sharing {d} system tokens, {d} tokens each: the tree holds {d:.1} MiB ({d} pages {d:.1} MiB, {d} snapshots {d:.1} MiB); copies of the caches at each cut held {d:.1} MiB; {d} prompt tokens came from the tree", .{ set.len, system_len, set[0].ids.len, mib(tree.held()), tree.heldPages(), mib(unique), tree.snaps, mib(states), mib(copied), cached }); } fn pagesBytes(c: *Ctx, n: usize) u64 { From cfc9004df16864321b5a3618273b728f2ddff7d7 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:08:46 +0000 Subject: [PATCH 117/186] feat(radix): page-major prefix groups on the store's path A family whose cache has several prefix groups (a target and a drafter's full attention on one pool, #2444 C3) needs one page id of each group for every logical page on the tree's path. The tree now keeps per ids a logical page, page-major, counts logical pages in its walks and ids in its limits, and initGroups declares the groups, refusing groups with other rows a page than the store's. With one group, today's case, every index and count is as before. --- zig/src/core/prompt_radix.zig | 37 +++++-- zig/src/core/prompt_radix_groups_test.zig | 125 ++++++++++++++++++++++ zig/src/core/prompt_radix_tree.zig | 39 ++++--- 3 files changed, 174 insertions(+), 27 deletions(-) create mode 100644 zig/src/core/prompt_radix_groups_test.zig diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 663b2cdce45..3376ad0adaa 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -3,6 +3,7 @@ const std = @import("std"); const pc = @import("prompt_cache.zig"); const tree_mod = @import("prompt_radix_tree.zig"); +const parts = @import("prompt_radix_parts.zig"); const Allocator = std.mem.Allocator; pub const Saved = tree_mod.Saved; @@ -11,6 +12,7 @@ pub const Pages = tree_mod.Pages; pub const Counts = tree_mod.Counts; pub const Node = tree_mod.Node; pub const Tree = tree_mod.Tree; +pub const Groups = parts.Groups; /// The rules of prompt_cache.zig, paged: `page` set; a planned store cuts and resumes only on `grid`, a multiple of it. pub const Rules = pc.Rules; @@ -64,7 +66,7 @@ pub const Query = struct { return .{ .key = key, .prompt_len = @intCast(key.len) }; } }; -/// A paged store's `keep`: whether the mark holds a state, and how many leading pages the stream swaps for the store's. +/// A paged store's `keep`: whether the mark holds a state, and how many leading ids the stream swaps for the store's. pub const Kept = struct { held: bool, shared: usize }; pub const Store = struct { @@ -77,8 +79,19 @@ pub const Store = struct { const SHARED_KEYS = 64; pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules) !Store { + return initGroups(gpa, family, pages, rules, .{}); + } + + /// `init` for a family that declares its cache groups (C3, #2444): `P` prefix groups page-major on the path. + pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; - return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead), .rules = rules }; + for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all + return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules }; + } + + /// Page ids a logical page: one per prefix group. + fn per(s: *const Store) usize { + return s.tree.per; } /// Where a pass may cut and resume: a page edge, or a planned store's grid. @@ -167,7 +180,7 @@ pub const Store = struct { if (marks_.len == 0) return; const last = marks_[marks_.len - 1]; while (true) { - const new = last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead]); + const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; s.tree.evict(s.victim(prompt, from) orelse return); } @@ -198,9 +211,9 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { + if (m / s.rules.page * s.per() <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { s.tree.counts.refused += 1; - note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page, s.limits.pages }); + note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); } } return out.toOwnedSlice(a); @@ -236,29 +249,30 @@ pub const Store = struct { return out.toOwnedSlice(a); } - /// At `at` the tree takes the pages of `mine` it lacks and the snapshot, evicting to fit; `path` gets its pages. + /// At `at` the tree takes the ids of `mine` (page-major, `per` a page) it lacks, and the snapshot; `path` its ids. pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { const prompt = q.key; const t = &s.tree; - if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != at / s.rules.page or path.len != mine.len) return .{ .held = false, .shared = 0 }; + const np = at / s.rules.page; // logical pages below `at` + if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != np * s.per() or path.len != mine.len) return .{ .held = false, .shared = 0 }; t.clock += 1; const tokens = prompt[0 .. at + s.rules.lookahead]; // a page past which the lookahead differs is another page const flagged = std.mem.indexOfScalar(u64, s.shared_keys.items, sharedKey(tokens)) != null; while (true) { var last: ?*Node = null; - const got = t.heldPrefix(tokens, mine.len, path, &last); - if (got == mine.len and last != null and last.?.end == at and last.?.entry != null) { + const got = t.heldPrefix(tokens, np, path, &last); + if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy e.shared = e.shared or flagged; t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.heldPages() + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + if (t.heldPages() + (np - got) * s.per() <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); - return .{ .held = false, .shared = got }; + return .{ .held = false, .shared = got * s.per() }; }; t.evict(v); } @@ -345,4 +359,5 @@ fn note(comptime fmt: []const u8, args: anytype) void { test { _ = @import("prompt_radix_test.zig"); _ = @import("prompt_radix_parts.zig"); + _ = @import("prompt_radix_groups_test.zig"); } diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig new file mode 100644 index 00000000000..827672bbb12 --- /dev/null +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -0,0 +1,125 @@ +//! The store over two prefix groups (C3, #2444): ids page-major on the path, limits counted in ids. +const std = @import("std"); +const radix = @import("prompt_radix.zig"); +const testing = std.testing; + +/// Pages whose holders are counted, and states that are the position they were kept at; a page is 4 tokens. +const Fake = struct { + refs: [32]u32 = @splat(0), + states: [8]u32 = @splat(0), + nstates: usize = 0, + restored: ?u32 = null, + + fn of(p: *anyopaque) *Fake { + return @ptrCast(@alignCast(p)); + } + fn retain(p: *anyopaque, id: u32) void { + of(p).refs[id] += 1; + } + fn release(p: *anyopaque, id: u32) void { + of(p).refs[id] -= 1; + } + fn holders(p: *anyopaque, id: u32) u32 { + return of(p).refs[id]; + } + fn available(p: *anyopaque) usize { + var n: usize = 0; + for (of(p).refs) |r| n += @intFromBool(r == 0); + return n; + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(p: *anyopaque, _: ?*anyopaque, at: u32) anyerror!radix.Saved { + const f = of(p); + f.states[f.nstates] = at; + f.nstates += 1; + return @ptrCast(&f.states[f.nstates - 1]); + } + fn restore(p: *anyopaque, _: ?*anyopaque, saved: radix.Saved) anyerror!void { + const at: *u32 = @ptrCast(@alignCast(saved)); + of(p).restored = at.*; + } + fn dropState(_: *anyopaque, _: radix.Saved) void {} + + fn pages(f: *Fake) radix.Pages { + return .{ .ptr = f, .bytes = 1, .vtable = &.{ .retain = retain, .release = release, .holders = holders, .available = available } }; + } + fn family(f: *Fake) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = dropState } }; + } + /// A stream's own ids: taken pages, one holder each. + fn take(f: *Fake, ids: []const u32) void { + for (ids) |id| f.refs[id] = 1; + } +}; + +const two: radix.Groups = .{ .prefix_rows = &.{ 4, 4 } }; + +fn open(f: *Fake, limit_ids: usize) !radix.Store { + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two); + s.limit(.{ .pages = limit_ids, .snaps = 4 }); + return s; +} + +test "two prefix groups: a keep holds each logical page's two ids, and a resume adopts them page-major" { + var f: Fake = .{}; + var s = try open(&f, 16); + defer s.deinit(); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + const mine = [_]u32{ 0, 1, 2, 3 }; // page 0: ids 0 (group 0) and 1 (group 1); page 1: 2 and 3 + f.take(&mine); + var path: [4]u32 = undefined; + const kept = s.keep(.of(&prompt), 8, null, &mine, &path); + try testing.expect(kept.held); + try testing.expectEqual(@as(usize, 4), kept.shared); // ids, not logical pages + try testing.expectEqualSlices(u32, &mine, &path); + try testing.expectEqual(@as(usize, 4), s.tree.heldPages()); + for (mine) |id| try testing.expectEqual(@as(u32, 2), f.refs[id]); // the stream and the store + for (mine) |id| f.refs[id] -= 1; // the stream ends + + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&prompt), 10, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 8), plan.from); + try testing.expectEqual(@as(?u32, 8), f.restored); + try testing.expectEqualSlices(u32, &mine, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; +} + +test "two prefix groups: a longer prompt shares the first page's ids and adds its own, split page-major" { + var f: Fake = .{}; + var s = try open(&f, 16); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 0 }; + const b = [_]u32{ 1, 2, 3, 4, 9, 9, 9, 9, 0 }; + const mine_a = [_]u32{ 0, 1, 2, 3 }; + const mine_b = [_]u32{ 4, 5, 6, 7 }; + f.take(&mine_a); + f.take(&mine_b); + var path: [4]u32 = undefined; + try testing.expect(s.keep(.of(&a), 8, null, &mine_a, &path).held); + const kept = s.keep(.of(&b), 8, null, &mine_b, &path); + try testing.expect(kept.held); + try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); // the shared first page is a's ids; b's second is its own + try testing.expectEqual(@as(usize, 6), s.tree.heldPages()); +} + +test "two prefix groups: the page limit counts ids, so a keep needing more ids than it holds evicts or refuses" { + var f: Fake = .{}; + var s = try open(&f, 4); // two logical pages of two groups + defer s.deinit(); + const long = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2, 3, 4, 5 }; + f.take(&mine); + var path: [6]u32 = undefined; + try testing.expect(!s.keep(.of(&long), 12, null, &mine, &path).held); // 6 ids pass the limit of 4 + try testing.expectEqual(@as(usize, 0), s.tree.heldPages()); + try testing.expect(s.keep(.of(&long), 8, null, mine[0..4], path[0..4]).held); // 4 ids fit +} + +test "prefix groups with other rows a page than the store's are refused when it opens" { + var f: Fake = .{}; + try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } })); +} diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index 3eff0969810..eee106ebf31 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -47,6 +47,8 @@ pub const Tree = struct { pages: Pages, /// Tokens a page holds. page: usize, + /// Prefix groups: page ids a logical page, page-major on each node (`pages[p * per + g]`). + per: usize = 1, /// Tokens past a page that its rows read (the family's lookahead): a page matches only when they match too. ahead: usize, root: *Node, @@ -57,10 +59,15 @@ pub const Tree = struct { snaps: usize = 0, counts: Counts = .{}, - pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, page: usize, ahead: usize) !Tree { + pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, page: usize, ahead: usize, per: usize) !Tree { const root = try gpa.create(Node); root.* = .{ .parent = null, .end = 0, .tokens = &.{}, .pages = &.{} }; - return .{ .gpa = gpa, .family = family, .pages = pages, .page = page, .ahead = ahead, .root = root }; + return .{ .gpa = gpa, .family = family, .pages = pages, .page = page, .per = per, .ahead = ahead, .root = root }; + } + + /// Logical pages of `n`: its ids over the prefix groups. + pub fn npages(t: *const Tree, n: *const Node) usize { + return n.pages.len / t.per; } pub fn deinit(t: *Tree) void { @@ -101,7 +108,7 @@ pub const Tree = struct { fn common(t: *const Tree, n: *const Node, tokens: []const u32, most: usize) usize { const w = t.page + t.ahead; var p: usize = 0; - while (p < n.pages.len and p < most and std.mem.eql(u32, n.tokens[p * t.page ..][0..w], tokens[p * t.page ..][0..w])) p += 1; + while (p < t.npages(n) and p < most and std.mem.eql(u32, n.tokens[p * t.page ..][0..w], tokens[p * t.page ..][0..w])) p += 1; return p; } @@ -113,7 +120,7 @@ pub const Tree = struct { while (at + t.page <= limit_at) { const c = t.child(node, prompt[at..]) orelse break; const p = t.common(c, prompt[at..], (limit_at - at) / t.page); - if (p < c.pages.len) break; + if (p < t.npages(c)) break; at += p * t.page; node = c; if (c.entry != null and c.end % step == 0) best = c; @@ -121,17 +128,17 @@ pub const Tree = struct { return best; } - /// Pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, into `path` if any; `last` the end node. + /// Logical pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, their ids into `path` if any. pub fn heldPrefix(t: *const Tree, tokens: []const u32, np: usize, path: ?[]u32, last: *?*Node) usize { var node = t.root; var pi: usize = 0; while (pi < np) { const c = t.child(node, tokens[pi * t.page ..]) orelse break; const p = t.common(c, tokens[pi * t.page ..], np - pi); - if (path) |out| @memcpy(out[pi..][0..p], c.pages[0..p]); + if (path) |out| @memcpy(out[pi * t.per ..][0 .. p * t.per], c.pages[0 .. p * t.per]); pi += p; node = c; - if (p < c.pages.len) break; + if (p < t.npages(c)) break; } last.* = if (node == t.root) null else node; return pi; @@ -149,7 +156,7 @@ pub const Tree = struct { defer path.deinit(t.gpa); var w: ?*Node = n; while (w) |x| : (w = x.parent) if (x.parent != null) try path.append(t.gpa, x); - try adopt.ensureUnusedCapacity(t.gpa, n.end / t.page); + try adopt.ensureUnusedCapacity(t.gpa, n.end / t.page * t.per); while (path.pop()) |x| { for (x.pages) |id| { t.pages.vtable.retain(t.pages.ptr, id); @@ -159,28 +166,28 @@ pub const Tree = struct { } } - /// Puts `tokens` (whole pages, `mine` the caller's, and `ahead` more) in the tree; the end node, `path` its pages. + /// Puts `tokens` (whole pages, `mine` the caller's ids page-major, `ahead` more) in the tree; `path` gets its ids. pub fn insert(t: *Tree, tokens: []const u32, mine: []const u32, path: []u32) !*Node { - const np = mine.len; + const np = mine.len / t.per; var node = t.root; var pi: usize = 0; while (pi < np) { const c = t.child(node, tokens[pi * t.page ..]) orelse { const leaf = try t.gpa.create(Node); errdefer t.gpa.destroy(leaf); - leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi..]) }; + leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi * t.per ..]) }; errdefer t.gpa.free(leaf.tokens); errdefer t.gpa.free(leaf.pages); try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); - @memcpy(path[pi..], leaf.pages); + @memcpy(path[pi * t.per ..], leaf.pages); t.touch(leaf); return leaf; }; const p = t.common(c, tokens[pi * t.page ..], np - pi); - const up = if (p < c.pages.len) try t.split(c, p) else c; - @memcpy(path[pi..][0..p], up.pages); + const up = if (p < t.npages(c)) try t.split(c, p) else c; + @memcpy(path[pi * t.per ..][0 .. p * t.per], up.pages); t.touch(up); node = up; pi += p; @@ -194,11 +201,11 @@ pub const Tree = struct { errdefer t.gpa.destroy(up); const head_tokens = try t.gpa.dupe(u32, c.tokens[0 .. p * t.page + t.ahead]); // its pages, the tail's first `ahead` errdefer t.gpa.free(head_tokens); - const head_pages = try t.gpa.dupe(u32, c.pages[0..p]); + const head_pages = try t.gpa.dupe(u32, c.pages[0 .. p * t.per]); errdefer t.gpa.free(head_pages); const tail_tokens = try t.gpa.dupe(u32, c.tokens[p * t.page ..]); errdefer t.gpa.free(tail_tokens); - const tail_pages = try t.gpa.dupe(u32, c.pages[p..]); + const tail_pages = try t.gpa.dupe(u32, c.pages[p * t.per ..]); errdefer t.gpa.free(tail_pages); const parent = c.parent.?; up.* = .{ .parent = parent, .end = parent.end + @as(u32, @intCast(p * t.page)), .tokens = head_tokens, .pages = head_pages, .used = c.used }; From af86722a190c519a07b866622feb617ca878f3e6 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:35:35 +0000 Subject: [PATCH 118/186] feat(radix): checkpoint parts made, restored and reclaimed C3 of #2444: a family may declare checkpoint groups whose pages an entry holds beside its legacy snapshot. A copied part takes store-held pool pages its live bytes are copied into; a retained part names the stream's own pages and copies a partial last page. Entries are published only whole, with every copy fenced before a page goes back; a node is pinned while a keep builds it or a resume restores it, so reclaim never frees its own source; reclaim first frees idle parts. Also fixes a leak in insert when its second allocation failed. Stores without parts are unchanged. --- zig/src/core/kv_pool_checks.zig | 1 + zig/src/core/kv_pool_parts_test.zig | 393 ++++++++++++++++++++++ zig/src/core/prompt_cache.zig | 2 + zig/src/core/prompt_radix.zig | 137 +++++++- zig/src/core/prompt_radix_groups_test.zig | 4 +- zig/src/core/prompt_radix_make.zig | 156 +++++++++ zig/src/core/prompt_radix_parts.zig | 19 ++ zig/src/core/prompt_radix_tree.zig | 49 ++- 8 files changed, 732 insertions(+), 29 deletions(-) create mode 100644 zig/src/core/kv_pool_parts_test.zig create mode 100644 zig/src/core/prompt_radix_make.zig diff --git a/zig/src/core/kv_pool_checks.zig b/zig/src/core/kv_pool_checks.zig index fb0678a4a73..a938dcc392d 100644 --- a/zig/src/core/kv_pool_checks.zig +++ b/zig/src/core/kv_pool_checks.zig @@ -1,6 +1,7 @@ //! The page pool's checks beside engine_api: under the paged prompt store and a lane host (kv_pool is its own module). test { _ = @import("kv_pool_store_test.zig"); + _ = @import("kv_pool_parts_test.zig"); _ = @import("kv_pool_host_test.zig"); _ = @import("kv_pool_ledger_test.zig"); } diff --git a/zig/src/core/kv_pool_parts_test.zig b/zig/src/core/kv_pool_parts_test.zig new file mode 100644 index 00000000000..eebc961d2a7 --- /dev/null +++ b/zig/src/core/kv_pool_parts_test.zig @@ -0,0 +1,393 @@ +//! The paged store's checkpoint parts (C3, #2444) over a real pool: made all or nothing, held once, reclaimed, pinned. +const std = @import("std"); +const kv = @import("kv_pool"); +const radix = @import("engine_api").prompt_radix; +const Pool = kv.Pool; +const Lease = kv.Lease; +const testing = std.testing; +const expect = testing.expect; +const expectEqual = testing.expectEqual; +const expectEqualSlices = testing.expectEqualSlices; + +const page = 4; // tokens a page of the prefix group +const page_bytes = 8; + +/// A family whose device memory is one word a page: a copy writes the live state's word, a restore reads it back. +const Family = struct { + pool: *Pool, + mem: [64]u32 = @splat(0), + live: [2]u32 = .{ 0, 0 }, // each group's live state + restored: [2]?u32 = .{ null, null }, + fences: u32 = 0, + copies: u32 = 0, + rows_copied: u32 = 0, + fail_copy: bool = false, + fail_rows: bool = false, + queued: []const u32 = &.{}, // the pages the last copy wrote (queued work, until a fence) + held_at_fence: bool = true, // every fence saw the queued pages still held: nothing went back before it + hold_on_restore: usize = 0, // pages the restore takes from the pool (its reclaim runs) + + fn of(p: *anyopaque) *Family { + return @ptrCast(@alignCast(p)); + } + fn bytes(_: *anyopaque, _: u32, _: u32) u64 { + return 20; // three pages of 8 bytes + } + fn copy(p: *anyopaque, _: ?*anyopaque, group: u32, _: u32, dst: radix.Parts.Part) anyerror!void { + const f = of(p); + for (dst.ids) |id| f.mem[id] = f.live[group]; + f.queued = dst.ids; + if (f.fail_copy) return error.CopyFailed; // after its work is queued + f.copies += 1; + } + fn restore(p: *anyopaque, _: ?*anyopaque, group: u32, src: radix.Parts.Part) anyerror!void { + const f = of(p); + if (f.hold_on_restore > 0) { + var ids: [8]u32 = undefined; + if (!f.pool.hold(ids[0..f.hold_on_restore])) return error.NoPages; + for (ids[0..f.hold_on_restore]) |id| f.pool.release(id); + } + f.restored[group] = f.mem[src.ids[0]]; + } + fn copyRows(p: *anyopaque, _: ?*anyopaque, _: u32, src: u32, dst: u32, rows: u32) anyerror!void { + const f = of(p); + if (f.fail_rows) return error.CopyFailed; + f.mem[dst] = f.mem[src]; + f.rows_copied += rows; + } + fn fence(p: *anyopaque, _: ?*anyopaque) void { + const f = of(p); + f.fences += 1; + for (f.queued) |id| f.held_at_fence = f.held_at_fence and f.pool.holders(id) > 0; + f.queued = &.{}; + } + fn hooks(f: *Family) radix.PartHooks { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .copy = copy, .restore = restore, .copy_rows = copyRows, .fence = fence } }; + } + + fn snapBytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + return error.NoLegacy; + } + fn snapRestore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: radix.Saved) void {} + fn snapshots(f: *Family) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = snapBytes, .save = save, .restore = snapRestore, .drop = drop } }; + } +}; + +const copied_state = [_]radix.Parts.PartSpec{.{ .kind = .state, .make = .copied }}; +const window3 = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained, .rows = 3 }}; +const both = [_]radix.Parts.PartSpec{ .{ .kind = .window, .make = .retained, .rows = 3 }, .{ .kind = .state, .make = .copied } }; + +const rules: radix.Rules = .{ .page = page, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }; + +const World = struct { + pool: Pool, + f: Family = undefined, + store: radix.Store = undefined, + reclaims: u32 = 0, // the pool asked the store + + fn init(w: *World, specs: []const radix.Parts.PartSpec, pages: usize) !void { + w.pool = try Pool.init(testing.allocator, pages, page_bytes); + w.f = .{ .pool = &w.pool }; + w.store = try radix.Store.initGroups(testing.allocator, w.f.snapshots(), radix.pagesOf(Pool, &w.pool, page_bytes), rules, .{ .parts = specs, .legacy = false }, w.f.hooks()); + w.store.limit(.{ .pages = pages, .snaps = 8 }); + try w.pool.joinReclaim(.{ .ctx = &w.store, .call = reclaim }); + } + fn deinit(w: *World) void { + w.pool.leaveReclaim(&w.store); + w.store.deinit(); + expect(w.pool.idle() and w.pool.available() == w.pool.total()) catch @panic("pages left after close"); + w.pool.deinit(); + } + fn reclaim(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + @as(*World, @fieldParentPtr("store", s)).reclaims += 1; + return s.reclaim(want); + } + /// A stream's lease of `n` pages. + fn lease(w: *World, n: usize) !Lease { + var l: Lease = .{}; + try w.pool.promise(&l, n); + try w.pool.cover(&l, n, null); + return l; + } +}; + +const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14 }; + +test "a copied part: held from the pool at keep, restored and fenced at a resume, back at eviction" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 77; + var path: [2]u32 = undefined; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(kept.held); + try expectEqual(@as(usize, 5), w.store.tree.heldPages()); // two path pages, three part pages + w.pool.finish(&l); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 8), plan.from); + try expectEqual(@as(?u32, 77), w.f.restored[0]); + try expect(w.f.fences >= 2); // the keep's copies, then the restore's + for (adopt.items) |id| w.pool.release(id); + const before = w.pool.available(); + w.store.limit(.{ .pages = 16, .snaps = 0 }); // the entry goes, its parts with it + try expectEqual(before + 5, w.pool.available()); // the part pages and the path the entry alone kept +} + +test "a retained window with other rows a page: whole pages shared, the partial one copied, mapped at their logical index" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); // the prefix group + var win = try w.lease(3); // the window: rows 0..8 on 3-row pages, the last one partial (2 rows) + for (win.ids.items, 0..) |id, i| w.f.mem[id] = @intCast(100 + i); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{.{ .base = 1, .ids = win.ids.items[1..3] }}; // its window omits page 0 + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + try expectEqual(@as(u32, 2), w.f.rows_copied); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[1])); // the stream and the store + try expectEqual(@as(u32, 1), w.pool.holders(win.ids.items[2])); // the stream's partial page: never the store's + const last_page = win.ids.items[2]; + w.f.mem[last_page] = 999; // the stream writes on: the checkpoint does not change + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 1), plan.parts.len); + try expectEqual(@as(u32, 1), plan.parts[0].base); + try expectEqual(@as(usize, 1), plan.parts[0].ids.len); // page 1: whole + try expectEqual(@as(usize, 1), plan.copies.len); + try expectEqual(@as(u32, 2), plan.copies[0].index); // logical page 2, rows 0..2 + try expectEqual(@as(u32, 2), plan.copies[0].rows); + try expectEqual(@as(u32, 102), w.f.mem[plan.copies[0].src]); // the rows as they were at the keep + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + const fences = w.f.fences; + w.store.finish(null, plan.copies); + try expectEqual(fences + 1, w.f.fences); +} + +test "a part copy that fails publishes nothing: the path stays, the part pages go back after a fence" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + for ([_]bool{ true, false }) |rows_fail| { + w.f.fail_rows = rows_fail; + w.f.fail_copy = !rows_fail; + const fences = w.f.fences; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + try expect(!kept.held); + try expectEqual(@as(usize, 2), kept.shared); // the path is the store's + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // and nothing else + try expectEqual(fences + 1, w.f.fences); + try expect(w.f.held_at_fence); // the fence came before the pages went back + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); + } + try expectEqual(@as(u64, 2), w.store.tree.counts.failed); +} + +test "two entries whose windows overlap hold each shared window page once" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(3); + defer w.pool.finish(&l); + var win = try w.lease(4); // rows 0..12 on 3-row pages + defer w.pool.finish(&win); + var path: [3]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items[0..1], &.{.{ .base = 0, .ids = win.ids.items[0..2] }}, path[0..1]).held); + try expect(w.store.keepParts(.of(&prompt), 12, null, l.ids.items[0..3], &.{.{ .base = 0, .ids = win.ids.items[0..4] }}, path[0..3]).held); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[0])); // named by both entries, held once by the store +} + +test "reclaim frees a copied part while its path and retained part are a stream's; the stream keeps its pages" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items[0..3] }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const before = w.pool.available(); + try expect(w.store.reclaim(before + 3)); // the copied part's 3 pages (and the partial copy) are store-only + try expect(w.pool.available() >= before + 3); + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (l.ids.items) |id| try expect(w.pool.holders(id) >= 1); // the stream's pages stay + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); +} + +test "a restore that takes pool pages runs reclaim, which never evicts the entry being restored" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 8); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 5; + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); + w.f.hold_on_restore = w.pool.available() + 1; // one more than free: the pool asks the store + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 0), plan.from); // the restore could not get its pages: a prefill from the start + try expectEqual(@as(?u32, null), w.f.restored[0]); + try expectEqual(@as(usize, 0), adopt.items.len); +} + +test "parts need hooks, a pool that holds pages, and rows for a retained group" { + var pool = try Pool.init(testing.allocator, 4, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + const pages = radix.pagesOf(Pool, &pool, page_bytes); + try testing.expectError(error.PartsNeedHooks, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &copied_state, .legacy = false }, null)); + var no_hold = pages; + var vt = pages.vtable.*; + vt.hold = null; + no_hold.vtable = &vt; + try testing.expectError(error.PartsNeedHold, radix.Store.initGroups(testing.allocator, f.snapshots(), no_hold, rules, .{ .parts = &copied_state, .legacy = false }, f.hooks())); + const zero = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained }}; + try testing.expectError(error.RetainedRows, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &zero, .legacy = false }, f.hooks())); +} + +test "an allocation failing at each step of a keep with parts: before the path, nothing changes; after it, the path stays alone" { + var k: usize = 0; + while (true) : (k += 1) { + var pool = try Pool.init(testing.allocator, 32, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + var fa = std.testing.FailingAllocator.init(testing.allocator, .{}); + var store = try radix.Store.initGroups(fa.allocator(), f.snapshots(), radix.pagesOf(Pool, &pool, page_bytes), rules, .{ .parts = &both, .legacy = false }, f.hooks()); + store.limit(.{ .pages = 32, .snaps = 8 }); + var l: Lease = .{}; + try pool.promise(&l, 2); + try pool.cover(&l, 2, null); + var win: Lease = .{}; + try pool.promise(&win, 3); + try pool.cover(&win, 3, null); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + fa.fail_index = fa.alloc_index + k; + const kept = store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + const failed = fa.has_induced_failure; + if (failed) { + try expect(!kept.held); + try expectEqual(@as(usize, 0), store.tree.snaps); + const path_pages: usize = if (kept.shared == 2) 2 else 0; // after the insert the path is the store's + try expectEqual(path_pages, store.tree.heldPages()); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), pool.holders(id)); // no part outlived the failure + } else try expect(kept.held); + pool.finish(&l); + pool.finish(&win); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); + if (!failed) break; + } + try expect(k > 3); // the insert, the parts, the copies and the prompt each failed once +} + +test "a retained group without its pages through the mark is refused, not counted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 16); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var path: [2]u32 = undefined; + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 3, .ids = &.{} }}, &path).held); + try expectEqual(@as(usize, 0), w.store.tree.heldPages()); +} + +test "a keep whose parts need more than the free pages reclaims, but never the path it is building" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 6); + defer w.deinit(); + var l = try w.lease(2); + var path: [2]u32 = undefined; + w.f.fail_copy = true; // a failed checkpoint leaves its path alone in the tree + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // the path is now the store's alone: idle, reclaimable + var other = try w.lease(2); // 6 pages: 2 the path's, 2 this stream's, 2 free; the part needs 3 + defer w.pool.finish(&other); + w.f.fail_copy = false; + const asked = w.reclaims; + const kept = w.store.keepParts(.of(&prompt), 8, null, other.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(w.reclaims > asked); // the part's hold ran the reclaim chain + try expect(!kept.held); // the only reclaimable pages were the path being built: refused, not pruned + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // the path is still the tree's + try expectEqual(@as(usize, 2), kept.shared); +} + +test "a plan's retained ids are its own: an eviction between begin and adoption leaves them valid" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const whole = [_]u32{ win.ids.items[0], win.ids.items[1] }; + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + w.store.limit(.{ .pages = 32, .snaps = 0 }); // the entry goes before the stream maps anything + try expectEqualSlices(u32, &whole, plan.parts[0].ids); + for (plan.parts[0].ids) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); // the stream's holds alone + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + for (adopt.items) |id| w.pool.release(id); +} + +test "a begin whose hand-out allocation fails gives back the pages it adopted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = win.ids.items }}, &path).held); + w.pool.finish(&l); + w.pool.finish(&win); + var k: usize = 0; + while (true) : (k += 1) { + var fa = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = k }); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const census = w.pool.census(); + const plan = w.store.begin(fa.allocator(), .of(&prompt), 14, &.{}, null, &adopt) catch { + try expectEqual(@as(usize, 0), adopt.items.len); + try expectEqual(census, w.pool.census()); // every hold this begin took went back + continue; + }; + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + plan.deinit(fa.allocator()); + break; + } + try expect(k >= 3); // the marks, the parts, their ids and the copies each failed once +} diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index ec73923dc58..54d945e3dd0 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -72,6 +72,8 @@ pub const Pages = struct { holders: *const fn (ptr: *anyopaque, id: u32) u32, /// Pages free and not promised to a stream. available: *const fn (ptr: *anyopaque) usize, + /// `out.len` free pages the store holds itself, reclaiming first (a checkpoint's copies); null: none. + hold: ?*const fn (ptr: *anyopaque, out: []u32) bool = null, }; }; /// What a family's states depend on beyond their position. diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 3376ad0adaa..0e0249607ff 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -4,6 +4,7 @@ const std = @import("std"); const pc = @import("prompt_cache.zig"); const tree_mod = @import("prompt_radix_tree.zig"); const parts = @import("prompt_radix_parts.zig"); +const make_mod = @import("prompt_radix_make.zig"); const Allocator = std.mem.Allocator; pub const Saved = tree_mod.Saved; @@ -11,6 +12,7 @@ pub const Snapshots = tree_mod.Snapshots; pub const Pages = tree_mod.Pages; pub const Counts = tree_mod.Counts; pub const Node = tree_mod.Node; +pub const Entry = tree_mod.Entry; pub const Tree = tree_mod.Tree; pub const Groups = parts.Groups; @@ -20,7 +22,7 @@ pub const Rules = pc.Rules; /// What the tree may hold: pages, and snapshots. pub const Limits = struct { pages: usize = 0, snaps: usize = 0 }; -/// An owner's pages as the store holds them: `T` has retain, release, holders and available (kv_pool.Pool). +/// An owner's pages as the store holds them: `T` has retain, release, holders, available and maybe hold (kv_pool.Pool). pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { const V = struct { fn of(p: *anyopaque) *T { @@ -38,11 +40,35 @@ pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { fn available(p: *anyopaque) usize { return of(p).available(); } + fn hold(p: *anyopaque, out: []u32) bool { + return of(p).hold(out); + } }; - return .{ .ptr = owner, .bytes = bytes, .vtable = &.{ .retain = V.retain, .release = V.release, .holders = V.holders, .available = V.available } }; + const has_hold = @hasDecl(T, "hold"); + return .{ .ptr = owner, .bytes = bytes, .vtable = &.{ .retain = V.retain, .release = V.release, .holders = V.holders, .available = V.available, .hold = if (has_hold) V.hold else null } }; } -pub const Plan = pc.Plan; +/// Where a pass resumes and the marks it keeps at (in the caller's allocator), and its retained parts to map (C3). +pub const Plan = struct { + from: u32 = 0, + marks: []const u32 = &.{}, + /// Each retained group's whole pages, held once for the stream: it adopts each at `base + i` (or releases it). + parts: []parts.Retained = &.{}, + /// Partial pages to copy into pages of the stream's promise; `finish` ends them. + copies: []parts.Copy = &.{}, + + /// The plan's slices back to `a` (the ids' holds are the stream's to adopt or release). + pub fn deinit(p: Plan, a: Allocator) void { + a.free(p.marks); + for (p.parts) |r| a.free(r.ids); + a.free(p.parts); + a.free(p.copies); + } +}; +pub const Parts = parts; +pub const PartHooks = parts.PartHooks; +pub const Retained = parts.Retained; +pub const Copy = parts.Copy; /// What a pass asks the store about: the words the tree compares, and what reuse must also match (enforced from C6). pub const Query = struct { @@ -75,18 +101,31 @@ pub const Store = struct { limits: Limits = .{}, /// The tokens of recently planned shared cuts, hashed (at most SHARED_KEYS). shared_keys: std.ArrayList(u64) = .empty, + /// The checkpoint groups (C3) and the family's hooks for them; none for a store of legacy snapshots. + decl: make_mod.Decl = .{}, + /// Whether an entry holds a legacy snapshot. + legacy: bool = true, const SHARED_KEYS = 64; pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules) !Store { - return initGroups(gpa, family, pages, rules, .{}); + return initGroups(gpa, family, pages, rules, .{}, null); } - /// `init` for a family that declares its cache groups (C3, #2444): `P` prefix groups page-major on the path. - pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups) !Store { + /// `init` with the family's cache groups (C3, #2444) and its part hooks; `groups`' slices outlive the store. + pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups, hooks: ?PartHooks) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all - return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules }; + if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; + if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds + for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; + if (groups.none()) return error.PageOnlyStore; // a store with no checkpoint at all: not yet served + return .{ + .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), + .rules = rules, + .decl = .{ .specs = groups.parts, .hooks = hooks }, + .legacy = groups.legacy, + }; } /// Page ids a logical page: one per prefix group. @@ -143,19 +182,48 @@ pub const Store = struct { } s.reserve(prompt, found, marks_); var from: u32 = 0; + var out: Plan = .{ .marks = marks_ }; if (found) |n| { - const family = s.tree.family; - const ok = if (family.vtable.restore(family.ptr, owner, n.entry.?.saved)) |_| true else |err| blk: { + const ok = if (s.restore(owner, n)) |_| true else |err| blk: { note("restoring {d} tokens failed ({s}); prefilling from the start", .{ n.end, @errorName(err) }); break :blk false; }; if (ok) { + const had = adopt.items.len; try s.tree.take(n, adopt); from = n.end; + if (n.entry.?.parts.len > 0) { + const handed = s.decl.handOut(&s.tree, a, n.end, n.entry.?.parts) catch |err| { + for (adopt.items[had..]) |id| s.tree.pages.vtable.release(s.tree.pages.ptr, id); // this begin's + adopt.shrinkRetainingCapacity(had); + return err; + }; + out.parts = handed.parts; + out.copies = handed.copies; + } } s.resumed(n, prompt, ok); } - return .{ .from = from, .marks = marks_ }; + out.from = from; + return out; + } + + /// The entry at `n` into `owner`, snapshot then copied parts; pinned, so reclaim cannot evict its own source. + fn restore(s: *Store, owner: ?*anyopaque, n: *Node) !void { + const e = &n.entry.?; + n.pins += 1; + defer n.pins -= 1; + const family = s.tree.family; + if (e.saved) |saved| try family.vtable.restore(family.ptr, owner, saved); + try s.decl.restore(owner, e.parts); + } + + /// A resume's copies are done (the stream mapped its own pages from them): each source's hold goes back. + pub fn finish(s: *Store, owner: ?*anyopaque, copies: []const Copy) void { + if (copies.len == 0) return; + const h = s.decl.hooks.?; + h.vtable.fence(h.ptr, owner); // the stream's copies out of the sources complete first + for (copies) |c| s.tree.pages.vtable.release(s.tree.pages.ptr, c.src); } /// A restored state: it is now the prompt's (a failed restore drops it). @@ -251,10 +319,16 @@ pub const Store = struct { /// At `at` the tree takes the ids of `mine` (page-major, `per` a page) it lacks, and the snapshot; `path` its ids. pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + return s.keepParts(q, at, owner, mine, &.{}, path); + } + + /// `keep` with the retained groups' pages through `at` (one entry a declared part; copied ones empty). + pub fn keepParts(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, retained: []const Retained, path: []u32) Kept { const prompt = q.key; const t = &s.tree; const np = at / s.rules.page; // logical pages below `at` if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != np * s.per() or path.len != mine.len) return .{ .held = false, .shared = 0 }; + if (!s.decl.fits(at, retained)) return .{ .held = false, .shared = 0 }; t.clock += 1; const tokens = prompt[0 .. at + s.rules.lookahead]; // a page past which the lookahead differs is another page const flagged = std.mem.indexOfScalar(u64, s.shared_keys.items, sharedKey(tokens)) != null; @@ -268,7 +342,8 @@ pub const Store = struct { t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.heldPages() + (np - got) * s.per() <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + const ids = (np - got) * s.per() + s.decl.ids(t, at, retained); // every id the keep adds to the holds + if (t.heldPages() + ids <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); @@ -277,13 +352,22 @@ pub const Store = struct { t.evict(v); } const node = t.insert(tokens, mine, path) catch |err| return s.fail(at, err, 0); + node.pins += 1; // the parts' holds may reclaim: not this node, nor its path + defer node.pins -= 1; + // the path stays whatever follows; the entry is published only whole + const made = s.decl.make(t, owner, at, retained) catch |err| return s.fail(at, err, mine.len); const family = t.family; - const saved = family.vtable.save(family.ptr, owner, at) catch |err| return s.fail(at, err, mine.len); + const saved: ?Saved = if (s.legacy) family.vtable.save(family.ptr, owner, at) catch |err| { + t.dropParts(made); + return s.fail(at, err, mine.len); + } else null; const last = t.gpa.dupe(u32, prompt) catch { - family.vtable.drop(family.ptr, saved); + if (saved) |sv| family.vtable.drop(family.ptr, sv); + t.dropParts(made); return s.fail(at, error.OutOfMemory, mine.len); }; - node.entry = .{ .saved = saved, .bytes = family.vtable.bytes(family.ptr, at), .born = @intCast(prompt.len), .used = t.clock, .last = last, .shared = flagged }; + const bytes = if (s.legacy) family.vtable.bytes(family.ptr, at) else 0; + node.entry = .{ .saved = saved, .parts = made, .bytes = bytes, .born = @intCast(prompt.len), .used = t.clock, .last = last, .shared = flagged }; t.snaps += 1; t.counts.kept += 1; return .{ .held = true, .shared = mine.len }; @@ -318,7 +402,7 @@ pub const Store = struct { } fn pick(n: *Node, prompt: []const u32, skip: ?*Node, moved: *?*Node, oldest: *?*Node) void { - if (n.entry) |e| if (n != skip) { + if (n.entry) |e| if (n != skip and n.pins == 0) { // a state its conversation moved past goes first; a shared cut or a branch point only as the least valuable const gone = !e.shared and !e.branched and ((prompt.len > e.born and prompt.len > n.end and extends(n, prompt)) or superseded(n, e.born)); if (gone and (moved.* == null or e.used < moved.*.?.entry.?.used)) moved.* = n; @@ -338,10 +422,33 @@ pub const Store = struct { /// Evicts until `want` pool pages are free and not promised; false when the tree cannot give that many. pub fn reclaim(s: *Store, want: usize) bool { const pages = s.tree.pages; + if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { // parts first: their path may be in use + s.tree.evict(s.partVictim() orelse break); + }; while (pages.vtable.available(pages.ptr) < want) if (!s.tree.evictLeaf()) return false; return true; } + /// The least valuable entry with a part page only the store holds, never a pinned one. + fn partVictim(s: *const Store) ?*Node { + var best: ?*Node = null; + s.pickParts(s.tree.root, &best); + return best; + } + + fn pickParts(s: *const Store, n: *Node, best: *?*Node) void { + if (n.entry) |e| if (n.pins == 0 and s.freesPart(e) and (best.* == null or worse(n, best.*.?))) { + best.* = n; + }; + for (n.kids.items) |k| s.pickParts(k, best); + } + + fn freesPart(s: *const Store, e: Entry) bool { + const pages = s.tree.pages; + for (e.parts) |part| for (part.ids) |id| if (pages.vtable.holders(pages.ptr, id) == 1) return true; + return false; + } + /// One log line after a prompt pass: where it resumed, how many states it kept, and what the tree holds. pub fn report(s: *const Store, prompt: usize, from: u32, kept: u64) void { if (@import("builtin").is_test) return; diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index 827672bbb12..d28f4d5084b 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -57,7 +57,7 @@ const Fake = struct { const two: radix.Groups = .{ .prefix_rows = &.{ 4, 4 } }; fn open(f: *Fake, limit_ids: usize) !radix.Store { - var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two); + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two, null); s.limit(.{ .pages = limit_ids, .snaps = 4 }); return s; } @@ -121,5 +121,5 @@ test "two prefix groups: the page limit counts ids, so a keep needing more ids t test "prefix groups with other rows a page than the store's are refused when it opens" { var f: Fake = .{}; - try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } })); + try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } }, null)); } diff --git a/zig/src/core/prompt_radix_make.zig b/zig/src/core/prompt_radix_make.zig new file mode 100644 index 00000000000..4a6803dba4f --- /dev/null +++ b/zig/src/core/prompt_radix_make.zig @@ -0,0 +1,156 @@ +//! A checkpoint's parts (C3, #2444): made at a keep, restored at a resume, all or nothing, the copies fenced. +const std = @import("std"); +const parts = @import("prompt_radix_parts.zig"); +const tree_mod = @import("prompt_radix_tree.zig"); +const Allocator = std.mem.Allocator; + +const Part = parts.Part; +const Retained = parts.Retained; +const Copy = parts.Copy; +const Tree = tree_mod.Tree; + +/// A store's part declaration: its groups' specs, and the family's hooks for them. +pub const Decl = struct { + specs: []const parts.PartSpec = &.{}, + hooks: ?parts.PartHooks = null, + + /// Logical pages of a retained group below or at `at`: whole ones, and one more when its rows fall short. + fn spans(rows: u32, at: u32) u32 { + return (at + rows - 1) / rows; + } + + /// Whether `retained` gives each retained group its pages through `at` (one entry a part, in order). + pub fn fits(d: Decl, at: u32, retained: []const Retained) bool { + if (retained.len != d.specs.len) return false; + for (d.specs, retained) |spec, r| switch (spec.make) { + .copied => if (r.ids.len != 0) return false, + .retained => if (r.ids.len == 0 or r.base + r.ids.len != spans(spec.rows, at)) return false, + }; + return true; + } + + /// Ids a keep at `at` adds to the tree's holds: copied parts' pages, retained ids not held yet, partial copies. + pub fn ids(d: Decl, t: *const Tree, at: u32, retained: []const Retained) usize { + var n: usize = 0; + for (d.specs, retained, 0..) |spec, r, g| switch (spec.make) { + .copied => n += copiedPages(d, t, @intCast(g), at), + .retained => { + const partial = at % spec.rows != 0; + const whole = r.ids[0 .. r.ids.len - @intFromBool(partial)]; + for (whole) |id| n += @intFromBool(!t.holds.contains(id)); + n += @intFromBool(partial); + }, + }; + return n; + } + + fn copiedPages(d: Decl, t: *const Tree, g: u32, at: u32) usize { + const h = d.hooks.?; + const bytes = h.vtable.bytes(h.ptr, g, at); + return @intCast((bytes + t.pages.bytes - 1) / t.pages.bytes); + } + + /// The parts of a checkpoint at `at`, each id named in the tree's holds; on failure, what it took goes back. + pub fn make(d: Decl, t: *Tree, owner: ?*anyopaque, at: u32, retained: []const Retained) ![]Part { + if (d.specs.len == 0) return &.{}; + const h = d.hooks.?; + const out = try t.gpa.alloc(Part, d.specs.len); + var made: usize = 0; + errdefer { + h.vtable.fence(h.ptr, owner); // queued copies into the pages complete before they go back + t.dropEach(out[0..made]); + t.gpa.free(out); + } + try t.holds.reserve(t.gpa, d.ids(t, at, retained)); + for (d.specs, retained, 0..) |spec, r, gi| { + const g: u32 = @intCast(gi); + switch (spec.make) { + .copied => { + const n = copiedPages(d, t, g, at); + const held = try t.gpa.alloc(u32, n); + if (!hold(t, held)) { + t.gpa.free(held); + return error.NoCachePages; + } + for (held) |id| t.holds.adoptHeldAssumeCapacity(id); + out[made] = .{ .base = 0, .ids = held, .bytes = h.vtable.bytes(h.ptr, g, at) }; + made += 1; + try h.vtable.copy(h.ptr, owner, g, at, out[made - 1]); + }, + .retained => { + const rows_left = at % spec.rows; + const whole = r.ids.len - @intFromBool(rows_left != 0); + var copy: [1]u32 = undefined; // the partial page: a copy the store holds + if (rows_left != 0 and !hold(t, ©)) return error.NoCachePages; + const own = t.gpa.alloc(u32, r.ids.len) catch |err| { + if (rows_left != 0) t.pages.vtable.release(t.pages.ptr, copy[0]); + return err; + }; + for (r.ids[0..whole], own[0..whole]) |id, *o| { + t.holds.shareAssumeCapacity(t.pages, id); + o.* = id; + } + if (rows_left != 0) { + own[whole] = copy[0]; + t.holds.adoptHeldAssumeCapacity(copy[0]); + } + out[made] = .{ .base = r.base, .ids = own, .bytes = 0 }; + made += 1; + if (rows_left != 0) try h.vtable.copy_rows(h.ptr, owner, g, r.ids[whole], copy[0], rows_left); + }, + } + } + h.vtable.fence(h.ptr, owner); // every copy into the parts is done: they are the checkpoint's from now on + return out; + } + + fn hold(t: *Tree, out: []u32) bool { + const f = t.pages.vtable.hold orelse return false; + return f(t.pages.ptr, out); + } + + /// The copied parts of `e` into the owner's live state, fenced; the store keeps their pages. + pub fn restore(d: Decl, owner: ?*anyopaque, ps: []const Part) !void { + if (d.specs.len == 0) return; + const h = d.hooks.?; + defer h.vtable.fence(h.ptr, owner); // the copies out of the store's pages are done, failed or not + for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .copied) try h.vtable.restore(h.ptr, owner, @intCast(g), part); + } + + /// The retained parts a resume at `at` maps: whole pages held once for the stream, partial ones as copies to make. + pub fn handOut(d: Decl, t: *Tree, a: Allocator, at: u32, ps: []const Part) !struct { parts: []Retained, copies: []Copy } { + var nr: usize = 0; + var nc: usize = 0; + for (d.specs) |spec| if (spec.make == .retained) { + nr += 1; + nc += @intFromBool(at % spec.rows != 0); + }; + const out = try a.alloc(Retained, nr); + var owned: usize = 0; // retained entries whose ids `a` owns + errdefer { + for (out[0..owned]) |r| a.free(r.ids); + a.free(out); + } + const copies = try a.alloc(Copy, nc); + errdefer a.free(copies); + nr = 0; + for (d.specs, ps) |spec, part| if (spec.make == .retained) { + const whole = part.ids.len - @intFromBool(at % spec.rows != 0); + out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; // the plan's own: eviction frees part.ids + nr += 1; + owned = nr; + }; + nr = 0; + nc = 0; + for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .retained) { + const rows_left = at % spec.rows; + const whole = part.ids.len - @intFromBool(rows_left != 0); + for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); // the stream's hold on each, until it adopts or finishes + nr += 1; + if (rows_left == 0) continue; + copies[nc] = .{ .group = @intCast(g), .index = at / spec.rows, .src = part.ids[whole], .rows = rows_left }; + nc += 1; + }; + return .{ .parts = out, .copies = copies }; + } +}; diff --git a/zig/src/core/prompt_radix_parts.zig b/zig/src/core/prompt_radix_parts.zig index 8985549a2f2..39f6991a7a9 100644 --- a/zig/src/core/prompt_radix_parts.zig +++ b/zig/src/core/prompt_radix_parts.zig @@ -41,6 +41,25 @@ pub const Retained = struct { base: u32, ids: []const u32 }; /// A copy a stream makes from a held page into one of its own: rows [0, rows) of `src` to the group's page `index`. pub const Copy = struct { group: u32, index: u32, src: u32, rows: u32 }; +/// What the store asks a family that declares parts: the bytes, and the device copies (each on the owner's stream). +pub const PartHooks = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// A copied part's bytes at `at` tokens. + bytes: *const fn (ptr: *anyopaque, group: u32, at: u32) u64, + /// The owner's live bytes of `group` after `at` tokens into the held pages of `dst`. + copy: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, at: u32, dst: Part) anyerror!void, + /// The held pages of `src` into the owner's live `group` (the next prompt chunk starts at its position). + restore: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, src: Part) anyerror!void, + /// Rows [0, rows) of every layer of `group` from page `src` into page `dst`. + copy_rows: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, src: u32, dst: u32, rows: u32) anyerror!void, + /// Every copy queued for `owner` has completed, on every rank that runs it. + fence: *const fn (ptr: *anyopaque, owner: ?*anyopaque) void, + }; +}; + /// The store's references: one pool reference a page, however many nodes and entries name it. pub const Holds = struct { counts: std.AutoHashMapUnmanaged(u32, u32) = .empty, diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index eee106ebf31..cd9b50d1e3f 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -13,7 +13,10 @@ pub const Counts = pc.Counts; /// The state kept at a node's end. pub const Entry = struct { - saved: Saved, + /// The family's legacy snapshot; null when it declares none (its state is in `parts`). + saved: ?Saved, + /// The checkpoint groups' pages (C3), each id named in the tree's `holds`. + parts: []parts.Part = &.{}, bytes: u64, /// The length of the prompt that kept it: a later turn's is longer. born: u32, @@ -39,6 +42,8 @@ pub const Node = struct { kids: std.ArrayList(*Node) = .empty, entry: ?Entry = null, used: u64 = 0, + /// A keep or a restore in progress here: no eviction or prune takes this node, its entry or its path. + pins: u32 = 0, }; pub const Tree = struct { @@ -83,10 +88,7 @@ pub const Tree = struct { fn drop(t: *Tree, n: *Node) void { for (n.kids.items) |k| t.drop(k); n.kids.deinit(t.gpa); - if (n.entry) |e| { - t.family.vtable.drop(t.family.ptr, e.saved); - t.gpa.free(e.last); - } + if (n.entry) |e| t.dropEntry(e); for (n.pages) |id| t.holds.drop(t.pages, id); t.gpa.free(n.tokens); t.gpa.free(n.pages); @@ -175,9 +177,11 @@ pub const Tree = struct { const c = t.child(node, tokens[pi * t.page ..]) orelse { const leaf = try t.gpa.create(Node); errdefer t.gpa.destroy(leaf); - leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi * t.per ..]) }; - errdefer t.gpa.free(leaf.tokens); - errdefer t.gpa.free(leaf.pages); + const own_tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]); + errdefer t.gpa.free(own_tokens); + const own_pages = try t.gpa.dupe(u32, mine[pi * t.per ..]); + errdefer t.gpa.free(own_pages); + leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = own_tokens, .pages = own_pages }; try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); @@ -221,11 +225,32 @@ pub const Tree = struct { return up; } + /// An entry's snapshot, parts and prompt go. + fn dropEntry(t: *Tree, e: Entry) void { + if (e.saved) |saved| t.family.vtable.drop(t.family.ptr, saved); + t.dropParts(e.parts); + t.gpa.free(e.last); + } + + /// Parts made by a keep go back: each id's name in `holds`, then the slices. + pub fn dropParts(t: *Tree, ps: []parts.Part) void { + t.dropEach(ps); + t.gpa.free(ps); + } + + /// Each part's ids go back and its id slice is freed; the caller frees `ps` itself. + pub fn dropEach(t: *Tree, ps: []const parts.Part) void { + for (ps) |part| { + for (part.ids) |id| t.holds.drop(t.pages, id); + t.gpa.free(part.ids); + } + } + /// Frees the state at `n`, and the nodes below nothing keeps. pub fn evict(t: *Tree, n: *Node) void { const e = n.entry orelse return; - t.family.vtable.drop(t.family.ptr, e.saved); - t.gpa.free(e.last); + std.debug.assert(n.pins == 0); // a keep or a restore in progress + t.dropEntry(e); n.entry = null; t.snaps -= 1; t.counts.evicted += 1; @@ -236,7 +261,7 @@ pub const Tree = struct { pub fn prune(t: *Tree, start: *Node) void { var n = start; while (n.parent) |parent| { - if (n.kids.items.len > 0 or n.entry != null) return; + if (n.kids.items.len > 0 or n.entry != null or n.pins > 0) return; const slot = std.mem.indexOfScalar(*Node, parent.kids.items, n).?; _ = parent.kids.swapRemove(slot); t.drop(n); @@ -259,7 +284,7 @@ pub const Tree = struct { } fn weakestLeaf(t: *const Tree, n: *Node, best: *?*Node) void { - if (n.parent != null and n.kids.items.len == 0 and t.idle(n) and (best.* == null or lessValuable(n, best.*.?))) best.* = n; + if (n.parent != null and n.kids.items.len == 0 and n.pins == 0 and t.idle(n) and (best.* == null or lessValuable(n, best.*.?))) best.* = n; for (n.kids.items) |k| t.weakestLeaf(k, best); } From 51d01245c207a1186ad4eb8b444749f15059c8ad Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:55:28 +0000 Subject: [PATCH 119/186] feat(radix): a page-only store resumes at any matched page edge C3 of #2444 defines the store for a family with no checkpoint at all (pure full attention, as vLLM's prefix cache): it keeps the path's pages with no snapshot slot, resumes at the deepest page edge a prompt repeats, inside a node too, and makes room by freeing idle leaves, never the path a pass resumes. No family on the store is such a family today, so every store user is unchanged; the first one (MiniMax M3, P7) needs no store change. --- zig/src/core/prompt_radix.zig | 51 +++++++++++--- zig/src/core/prompt_radix_groups_test.zig | 86 +++++++++++++++++++++++ zig/src/core/prompt_radix_tree.zig | 47 +++++++++++++ 3 files changed, 175 insertions(+), 9 deletions(-) diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 0e0249607ff..c399486b078 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -105,6 +105,8 @@ pub const Store = struct { decl: make_mod.Decl = .{}, /// Whether an entry holds a legacy snapshot. legacy: bool = true, + /// No checkpoint at all (pure full attention): any matched page edge resumes, and no entry is kept. + page_only: bool = false, const SHARED_KEYS = 64; @@ -119,12 +121,12 @@ pub const Store = struct { if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; - if (groups.none()) return error.PageOnlyStore; // a store with no checkpoint at all: not yet served return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules, .decl = .{ .specs = groups.parts, .hooks = hooks }, .legacy = groups.legacy, + .page_only = groups.none(), }; } @@ -157,7 +159,14 @@ pub const Store = struct { /// Whether anything can be kept. pub fn on(s: *const Store) bool { - return s.limits.pages > 0 and s.limits.snaps > 0; + return s.limits.pages > 0 and (s.page_only or s.limits.snaps > 0); + } + + /// The least valuable thing to free: an entry, or for a store without entries an idle leaf; false when none. + fn evictOne(s: *Store, prompt: []const u32, skip: ?*Node) bool { + if (s.page_only) return s.tree.evictLeaf(); + s.tree.evict(s.victim(prompt, skip) orelse return false); + return true; } /// The deepest node whose state `prompt` resumes exactly: its tokens and lookahead a prefix, a token left. @@ -170,6 +179,7 @@ pub const Store = struct { pub fn begin(s: *Store, a: Allocator, q: Query, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { const prompt = q.key; if (!s.on()) return .{ .marks = try a.alloc(u32, 0) }; + if (s.page_only) return s.beginPages(a, prompt, history_len, shared, adopt); const found = s.find(prompt); if (found == null) s.tree.counts.misses += 1; const resume_at: u32 = if (found) |n| n.end else 0; @@ -208,6 +218,25 @@ pub const Store = struct { return out; } + /// `begin` for a store with no checkpoint: the deepest matched page edge, its path adopted, nothing restored. + fn beginPages(s: *Store, a: Allocator, prompt: []const u32, history_len: u32, shared: []const u32, adopt: *std.ArrayList(u32)) !Plan { + const hit: ?tree_mod.Hit = if (prompt.len < s.past() + 1) null else s.tree.deepestPage(prompt, (prompt.len - s.past()) / s.rules.page * s.rules.page); + const from: u32 = if (hit) |h| h.at(&s.tree) else 0; + const planned = try s.marks(a, prompt, from, history_len, shared, &.{}); + const marks_ = try s.fitting(a, planned); + errdefer a.free(marks_); + if (hit) |h| h.node.pins += 1; // the room made for the marks never prunes the path this pass resumes + s.reserve(prompt, null, marks_); + if (hit) |h| h.node.pins -= 1; + const h = hit orelse { + s.tree.counts.misses += 1; + return .{ .marks = marks_ }; + }; + try s.tree.takeTo(h, adopt); + s.tree.counts.hits += 1; + return .{ .from = from, .marks = marks_ }; + } + /// The entry at `n` into `owner`, snapshot then copied parts; pinned, so reclaim cannot evict its own source. fn restore(s: *Store, owner: ?*anyopaque, n: *Node) !void { const e = &n.entry.?; @@ -249,8 +278,8 @@ pub const Store = struct { const last = marks_[marks_.len - 1]; while (true) { const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); - if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; - s.tree.evict(s.victim(prompt, from) orelse return); + if (s.tree.heldPages() + new <= s.limits.pages and (s.page_only or s.tree.snaps + marks_.len <= s.limits.snaps)) return; + if (!s.evictOne(prompt, from)) return; } } @@ -279,7 +308,7 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page * s.per() <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { + if (m / s.rules.page * s.per() <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { s.tree.counts.refused += 1; note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); } @@ -335,6 +364,7 @@ pub const Store = struct { while (true) { var last: ?*Node = null; const got = t.heldPrefix(tokens, np, path, &last); + if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; // the path is the tree's already if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy @@ -343,15 +373,18 @@ pub const Store = struct { return .{ .held = true, .shared = mine.len }; } const ids = (np - got) * s.per() + s.decl.ids(t, at, retained); // every id the keep adds to the holds - if (t.heldPages() + ids <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; - const v = s.victim(prompt, null) orelse { + if (t.heldPages() + ids <= s.limits.pages and (s.page_only or t.snaps + 1 <= s.limits.snaps)) break; + if (!s.evictOne(prompt, null)) { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); return .{ .held = false, .shared = got * s.per() }; - }; - t.evict(v); + } } const node = t.insert(tokens, mine, path) catch |err| return s.fail(at, err, 0); + if (s.page_only) { + t.counts.kept += 1; + return .{ .held = true, .shared = mine.len }; + } node.pins += 1; // the parts' holds may reclaim: not this node, nor its path defer node.pins -= 1; // the path stays whatever follows; the entry is published only whole diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index d28f4d5084b..62a7d29848e 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -123,3 +123,89 @@ test "prefix groups with other rows a page than the store's are refused when it var f: Fake = .{}; try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } }, null)); } + +const page_only: radix.Groups = .{ .legacy = false }; + +fn openPages(f: *Fake, limit_ids: usize) !radix.Store { + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, page_only, null); + s.limit(.{ .pages = limit_ids, .snaps = 0 }); // no checkpoint: no snapshot slots needed + return s; +} + +test "a store with no checkpoint keeps pages with no snapshot slot and resumes at any matched page edge" { + var f: Fake = .{}; + var s = try openPages(&f, 16); + defer s.deinit(); + try testing.expect(s.on()); + const kept_prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&kept_prompt), 12, null, &mine, &path).held); + try testing.expectEqual(@as(usize, 0), s.tree.snaps); + for (mine) |id| f.refs[id] -= 1; + // a prompt that diverges inside the kept node's third page resumes after its second, adopting two pages + const other = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 99, 99, 99, 99, 99 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&other), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 8), plan.from); + try testing.expectEqualSlices(u32, &.{ 0, 1 }, adopt.items); + try testing.expectEqual(@as(?u32, null), f.restored); // nothing to restore + for (adopt.items) |id| f.refs[id] -= 1; + adopt.clearRetainingCapacity(); + // a prompt that ends inside it resumes at the last whole page that leaves a token to prefill + const short = [_]u32{ 1, 2, 3, 4, 5, 6, 7 }; + const plan2 = try s.begin(testing.allocator, .of(&short), 7, &.{}, null, &adopt); + defer testing.allocator.free(plan2.marks); + try testing.expectEqual(@as(u32, 4), plan2.from); + try testing.expectEqualSlices(u32, &.{0}, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; +} + +test "a store with no checkpoint frees an idle leaf to keep another prompt's pages" { + var f: Fake = .{}; + var s = try openPages(&f, 3); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&a), 12, null, &mine, &path).held); + for (mine) |id| f.refs[id] -= 1; + const b = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7, 0 }; + const theirs = [_]u32{ 5, 6 }; + f.take(&theirs); + var path2: [2]u32 = undefined; + try testing.expect(s.keep(.of(&b), 8, null, &theirs, &path2).held); // room: the idle leaf of `a` went + try testing.expect(s.tree.heldPages() <= 3); + for (theirs) |id| f.refs[id] -= 1; +} + +test "a store with no checkpoint making room for a pass's marks never frees the path that pass resumes" { + var f: Fake = .{}; + var s = try openPages(&f, 5); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&a), 12, null, &mine, &path).held); + for (mine) |id| f.refs[id] -= 1; // the tree's alone: idle, and the oldest leaf + const b = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7, 0 }; + const theirs = [_]u32{ 5, 6 }; + f.take(&theirs); + var path2: [2]u32 = undefined; + try testing.expect(s.keep(.of(&b), 8, null, &theirs, &path2).held); + for (theirs) |id| f.refs[id] -= 1; // a newer idle leaf: 5 pages held of 5 + const longer = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); // its mark at 20 wants 2 more pages + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 12), plan.from); + try testing.expectEqualSlices(u32, &mine, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; + try testing.expectEqual(@as(u32, 0), f.refs[5]); // the newer leaf went instead +} diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index cd9b50d1e3f..81ad18aa5d9 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -32,6 +32,17 @@ pub const Entry = struct { branched: bool = false, }; +/// A matched position: `pages` logical pages into `node` (1 to its count). +pub const Hit = struct { + node: *Node, + pages: usize, + + /// The tokens below the position. + pub fn at(h: Hit, t: *const Tree) u32 { + return h.node.end - @as(u32, @intCast((t.npages(h.node) - h.pages) * t.page)); + } +}; + pub const Node = struct { parent: ?*Node, /// Tokens up to the end of this node. @@ -130,6 +141,42 @@ pub const Tree = struct { return best; } + /// The deepest page edge `prompt` repeats up to `limit_at`, inside a node too: where a page-only store resumes. + pub fn deepestPage(t: *const Tree, prompt: []const u32, limit_at: usize) ?Hit { + var node = t.root; + var at: usize = 0; + var best: ?Hit = null; + while (at + t.page <= limit_at) { + const c = t.child(node, prompt[at..]) orelse break; + const p = t.common(c, prompt[at..], (limit_at - at) / t.page); + if (p == 0) break; + best = .{ .node = c, .pages = p }; + if (p < t.npages(c)) break; + at += p * t.page; + node = c; + } + return best; + } + + /// The ids of the path to `h`, each held once more for the caller; the nodes marked used. + pub fn takeTo(t: *Tree, h: Hit, adopt: *std.ArrayList(u32)) !void { + var path: std.ArrayList(*Node) = .empty; + defer path.deinit(t.gpa); + var w: ?*Node = h.node.parent; + while (w) |x| : (w = x.parent) if (x.parent != null) try path.append(t.gpa, x); + try adopt.ensureUnusedCapacity(t.gpa, h.at(t) / t.page * t.per); + while (path.pop()) |x| t.hand(x, x.pages, adopt); + t.hand(h.node, h.node.pages[0 .. h.pages * t.per], adopt); + } + + fn hand(t: *Tree, n: *Node, ids: []const u32, adopt: *std.ArrayList(u32)) void { + for (ids) |id| { + t.pages.vtable.retain(t.pages.ptr, id); + adopt.appendAssumeCapacity(id); + } + t.touch(n); + } + /// Logical pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, their ids into `path` if any. pub fn heldPrefix(t: *const Tree, tokens: []const u32, np: usize, path: ?[]u32, last: *?*Node) usize { var node = t.root; From b1fff1f7681475235a134bd0d3e34c3e21ad6080 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 21:54:14 +0000 Subject: [PATCH 120/186] fix(radix): count part pages when planning a pass's marks A mark whose path and checkpoint parts passed the page limit was planned anyway; making room for it evicted a kept checkpoint before the keep itself was refused. Planning now counts the part pages it knows before the keep (copied parts and partial retained pages), so such a mark is refused up front. Long trailing comments in the lines this branch added move above their lines. --- zig/src/core/kv_pool_parts_test.zig | 18 ++++++++++++++++++ zig/src/core/prompt_radix.zig | 19 +++++++++++++------ zig/src/core/prompt_radix_groups_test.zig | 6 ++++-- zig/src/core/prompt_radix_make.zig | 16 ++++++++++++++-- 4 files changed, 49 insertions(+), 10 deletions(-) diff --git a/zig/src/core/kv_pool_parts_test.zig b/zig/src/core/kv_pool_parts_test.zig index eebc961d2a7..70a567f8997 100644 --- a/zig/src/core/kv_pool_parts_test.zig +++ b/zig/src/core/kv_pool_parts_test.zig @@ -391,3 +391,21 @@ test "a begin whose hand-out allocation fails gives back the pages it adopted" { } try expect(k >= 3); // the marks, the parts, their ids and the copies each failed once } + +test "a mark whose path and parts pass the page limit is refused when planned, so it never evicts a checkpoint first" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + w.store.limit(.{ .pages = 5, .snaps = 8 }); + var l = try w.lease(1); + var path: [1]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // one prefix page and three copied: 4 of 5 + const other = [_]u32{ 9, 9, 9, 9, 8, 8, 8, 8, 7, 7, 7, 7, 6, 6 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&other), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 0), plan.marks.len); // its mark at 12 needs 3 path pages and 3 part pages: 6 of 5 + try expectEqual(@as(usize, 1), w.store.tree.snaps); // the kept checkpoint was not evicted for it +} diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index c399486b078..aec5202eafb 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -117,7 +117,8 @@ pub const Store = struct { /// `init` with the family's cache groups (C3, #2444) and its part hooks; `groups`' slices outlive the store. pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups, hooks: ?PartHooks) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; - for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all + // one key a page for all prefix groups + for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; @@ -276,8 +277,11 @@ pub const Store = struct { fn reserve(s: *Store, prompt: []const u32, from: ?*Node, marks_: []const u32) void { if (marks_.len == 0) return; const last = marks_[marks_.len - 1]; + var parts_ids: usize = 0; // each mark's parts, as far as they are known before the keep + for (marks_) |m| parts_ids += s.decl.planIds(&s.tree, m); while (true) { - const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); + const path_ids = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); + const new = path_ids + parts_ids; if (s.tree.heldPages() + new <= s.limits.pages and (s.page_only or s.tree.snaps + marks_.len <= s.limits.snaps)) return; if (!s.evictOne(prompt, from)) return; } @@ -308,9 +312,10 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page * s.per() <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { + const ids = m / s.rules.page * s.per() + s.decl.planIds(&s.tree, m); // the path and the parts' known pages + if (ids <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { s.tree.counts.refused += 1; - note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); + note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, ids, s.limits.pages }); } } return out.toOwnedSlice(a); @@ -364,7 +369,8 @@ pub const Store = struct { while (true) { var last: ?*Node = null; const got = t.heldPrefix(tokens, np, path, &last); - if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; // the path is the tree's already + // a page-only store: the path is the tree's already + if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy @@ -455,7 +461,8 @@ pub const Store = struct { /// Evicts until `want` pool pages are free and not promised; false when the tree cannot give that many. pub fn reclaim(s: *Store, want: usize) bool { const pages = s.tree.pages; - if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { // parts first: their path may be in use + // parts first: their path may be in use + if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { s.tree.evict(s.partVictim() orelse break); }; while (pages.vtable.available(pages.ptr) < want) if (!s.tree.evictLeaf()) return false; diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index 62a7d29848e..e6c8a8dd77c 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -102,7 +102,8 @@ test "two prefix groups: a longer prompt shares the first page's ids and adds it try testing.expect(s.keep(.of(&a), 8, null, &mine_a, &path).held); const kept = s.keep(.of(&b), 8, null, &mine_b, &path); try testing.expect(kept.held); - try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); // the shared first page is a's ids; b's second is its own + // the shared first page is a's ids; b's second is its own + try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); try testing.expectEqual(@as(usize, 6), s.tree.heldPages()); } @@ -202,7 +203,8 @@ test "a store with no checkpoint making room for a pass's marks never frees the const longer = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21 }; var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); - const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); // its mark at 20 wants 2 more pages + // its mark at 20 wants 2 more pages + const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); defer testing.allocator.free(plan.marks); try testing.expectEqual(@as(u32, 12), plan.from); try testing.expectEqualSlices(u32, &mine, adopt.items); diff --git a/zig/src/core/prompt_radix_make.zig b/zig/src/core/prompt_radix_make.zig index 4a6803dba4f..173384f83cf 100644 --- a/zig/src/core/prompt_radix_make.zig +++ b/zig/src/core/prompt_radix_make.zig @@ -44,6 +44,16 @@ pub const Decl = struct { return n; } + /// Ids a keep at `at` adds that a plan knows before the keep: copied parts' pages, one per retained partial page. + pub fn planIds(d: Decl, t: *const Tree, at: u32) usize { + var n: usize = 0; + for (d.specs, 0..) |spec, g| switch (spec.make) { + .copied => n += copiedPages(d, t, @intCast(g), at), + .retained => n += @intFromBool(at % spec.rows != 0), + }; + return n; + } + fn copiedPages(d: Decl, t: *const Tree, g: u32, at: u32) usize { const h = d.hooks.?; const bytes = h.vtable.bytes(h.ptr, g, at); @@ -136,7 +146,8 @@ pub const Decl = struct { nr = 0; for (d.specs, ps) |spec, part| if (spec.make == .retained) { const whole = part.ids.len - @intFromBool(at % spec.rows != 0); - out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; // the plan's own: eviction frees part.ids + // the plan's own copy: eviction frees part.ids + out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; nr += 1; owned = nr; }; @@ -145,7 +156,8 @@ pub const Decl = struct { for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .retained) { const rows_left = at % spec.rows; const whole = part.ids.len - @intFromBool(rows_left != 0); - for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); // the stream's hold on each, until it adopts or finishes + // the stream's hold on each, until it adopts or finishes + for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); nr += 1; if (rows_left == 0) continue; copies[nc] = .{ .group = @intCast(g), .index = at / spec.rows, .src = part.ids[whole], .rows = rows_left }; From 5ac1801fd4c1937ad12947550bc9227ff260452a Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Wed, 11 Nov 2026 10:30:54 +0000 Subject: [PATCH 121/186] test(radix): a legacy save that reclaims never prunes its leaf The CUDA ring families' snapshot saves take pool pages, so a save can make the pool reclaim from the store in the middle of a keep. On a keep onto a leaf an earlier failed save left with no entry, that reclaim could prune the very leaf the keep then writes its entry to. The keep's node pin prevents it; this test holds a save that needs one page more than free over exactly that leaf, and it crashes with the pin removed. --- zig/src/core/kv_pool_store_test.zig | 74 +++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/zig/src/core/kv_pool_store_test.zig b/zig/src/core/kv_pool_store_test.zig index fa3f5b18123..7f3699267dc 100644 --- a/zig/src/core/kv_pool_store_test.zig +++ b/zig/src/core/kv_pool_store_test.zig @@ -212,3 +212,77 @@ test "page accounting: random admits, keeps, resumes, finishes, cancels and fail } try expect(refused > 0 and failed > 0 and adopted > 0); // the runs met full pools, failed uploads and resumes } + +/// A legacy family whose save holds pool pages for its copy (as the CUDA ring families' blobs do), one more than free. +const HoldingFamily = struct { + pool: *Pool, + fail: bool = false, + live: usize = 0, + + const State = struct { ids: [8]u32, n: usize }; + + fn of(ptr: *anyopaque) *HoldingFamily { + return @ptrCast(@alignCast(ptr)); + } + fn snapshots(f: *HoldingFamily) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = drop } }; + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(ptr: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + const f = of(ptr); + if (f.fail) return error.CopyFailed; + const st = try std.testing.allocator.create(State); + st.n = f.pool.available() + 1; // the pool must reclaim from the store to give them + if (!f.pool.hold(st.ids[0..st.n])) { + std.testing.allocator.destroy(st); + return error.NoCachePages; + } + f.live += 1; + return st; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(ptr: *anyopaque, saved: radix.Saved) void { + const f = of(ptr); + const st: *State = @ptrCast(@alignCast(saved)); + for (st.ids[0..st.n]) |id| f.pool.release(id); + std.testing.allocator.destroy(st); + f.live -= 1; + } +}; + +fn reclaimStore(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + return s.reclaim(want); +} + +test "a legacy save that reclaims pool pages while keeping onto a leaf with no entry never prunes that leaf" { + var pool = try Pool.init(std.testing.allocator, 6, 1); + defer pool.deinit(); + var f: HoldingFamily = .{ .pool = &pool }; + var store = try radix.Store.init(std.testing.allocator, f.snapshots(), radix.pagesOf(Pool, &pool, 1), .{ .page = page, .min_gap = 1, .min_prompt = 0 }); + store.limit(.{ .pages = 6, .snaps = 4 }); + try pool.joinReclaim(.{ .ctx = &store, .call = reclaimStore }); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + var path: [2]u32 = undefined; + var first: Lease = .{}; + try pool.promise(&first, 2); + try pool.cover(&first, 2, null); + f.fail = true; // a failed save leaves the path in the tree, a leaf with no entry + try expect(!store.keep(.of(&prompt), 8, null, first.ids.items, &path).held); + pool.finish(&first); // the leaf's pages are the store's alone: the only ones reclaim could free + try expectEqual(@as(usize, 2), store.tree.heldPages()); + var second: Lease = .{}; + try pool.promise(&second, 2); + try pool.cover(&second, 2, null); + f.fail = false; + const kept = store.keep(.of(&prompt), 8, null, second.ids.items, &path); // onto that leaf; its save reclaims + try expect(!kept.held); // the pinned leaf could not be taken, so the save found no page + try expectEqual(@as(usize, 2), store.tree.heldPages()); // the leaf survived the keep that was writing to it + try expectEqual(@as(usize, 0), f.live); + pool.finish(&second); + pool.leaveReclaim(&store); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); +} From c91cc243ef6ea71be99b7f80d1095cc3d8905c20 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:09:56 +0000 Subject: [PATCH 122/186] changelog: #2875 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 24ad3b85410..14c6fc5b247 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. + - GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. - HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. From 63e955749ec65bdeb0ccce87e2861ac91ec2be05 Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Mon, 9 Nov 2026 23:32:54 +0200 Subject: [PATCH 123/186] qwen27: slot 0 keeps its rewind point through the warm pass, and a background pass's chunks narrow with depth A 20 GiB prompt cache holds two 128k states. The state kept at a reply's end pushed out the one at the history end, so a client whose replayed reply differs from the generated one resumed a turn further back. Slot 0's in-slot rewind buffer now exists whenever there is room, also beside a cache that keeps two full-context states, where needsRewind alone makes none, and its point is carried through passes that only extend slot 0, so the next turn rewinds in place. Beside such a wide cache, which keeps the reply's end, the warm pass over slot 0 runs when it carries that point; beside a smaller cache it still yields, as on main. Background chunks scale down with depth so a waiting request is held for less. --- zig/src/native/qwen27_prompt_rows.zig | 21 +++++--- zig/src/native/qwen27_resume.zig | 72 ++++++++++++++++++++++++++- zig/src/native/qwen27_served.zig | 2 +- zig/src/native/qwen27_slots.zig | 32 +++++++----- zig/src/qwen27_rows_test.zig | 47 +++++++++++++---- 5 files changed, 141 insertions(+), 33 deletions(-) diff --git a/zig/src/native/qwen27_prompt_rows.zig b/zig/src/native/qwen27_prompt_rows.zig index 9f885bb3a33..893c77e99de 100644 --- a/zig/src/native/qwen27_prompt_rows.zig +++ b/zig/src/native/qwen27_prompt_rows.zig @@ -1,9 +1,14 @@ -//! Each prompt chunk chooses its width from the streams that can decode beside it. -pub fn choose(b: anytype, current: anytype) usize { - for (b.owners[0..b.runner.slots]) |owner| if (owner) |stream| { - if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) { - return @min(b.chunk, 512); - } - }; - return b.chunk; +//! Each prompt chunk chooses its width from the streams that can decode beside it, and a background pass's depth. +pub fn choose(b: anytype, current: anytype, at: usize) usize { + const width = for (b.owners[0..b.runner.slots]) |owner| { + const stream = owner orelse continue; + if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) break @min(b.chunk, 512); + } else b.chunk; + return if (current.yield_check != null) @min(width, deep(b.chunk, at)) else width; +} + +/// A background chunk's width at offset `at`: narrower with depth, each about as long (256 at least). +pub fn deep(chunk: usize, at: usize) usize { + const scaled = @min(chunk, chunk * 16384 / @max(at, 16384)); + return @max(256, scaled / 128 * 128); } diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index 29e1e543b53..7a45c8e5f40 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -27,11 +27,13 @@ pub const Standing = struct { tokens: []u32 = &.{}, spans: []modes.Span = &.{}, rewind_at: u32 = 0, // where the rewind buffer holds slot 0's DeltaNet states and taps (0: none) + wide: bool = false, // the prompt cache keeps two full-context states: it holds a reply's end beside slot 0 + /// Frees the record; `wide`, the cache's size, stays. pub fn deinit(st: *Standing, gpa: std.mem.Allocator) void { gpa.free(st.tokens); gpa.free(st.spans); - st.* = .{}; + st.* = .{ .wide = st.wide }; } /// What `l` left in slot 0 of `context` (its rows' spans `spans`), when a later prompt may resume it; else nothing. @@ -41,7 +43,20 @@ pub const Standing = struct { if (!l.kept() or l.at > context.len) return; const canonical = (if (l.at <= l.prompt_len) modes.prefix(gpa, spans, l.at) else modes.reply(gpa, spans, l.prompt_len, l.at)) catch return; const tokens = gpa.dupe(u32, context[0..l.at]) catch return gpa.free(canonical); - st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0 }; + st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .wide = st.wide }; + } + + /// The rewind point a pass resuming slot 0 at `n` keeps while it only extends the slot (0: none). + pub fn carry(st: *const Standing, n: u32) u32 { + return if (st.rewind_at > 0 and st.rewind_at <= n) st.rewind_at else 0; + } + + /// The rewind point a pass that `r` resumes in slot 0 carries through it (0: none, or it does not resume there). + pub fn carried(st: *const Standing, r: Resume) u32 { + return switch (r) { + .none => 0, + .here, .back => |n| st.carry(n), + }; } /// How `ids` (its rows' spans `spans`) resumes, past the cache's own restore point `kept_at`, by the cache's rule. @@ -184,6 +199,59 @@ test "qwen27 resume: slot 0 keeps what a fresh pass would stand at, and the poin try t.expectEqual(@as(u32, 0), st.rewind_at); } +test "qwen27 resume: a warm pass that only extends slot 0 keeps the reply's rewind point for a diverging next turn" { + var st: Standing = .{ .rewind_at = 4 }; // the turn's pass saved slot 0's states at its history's end + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 24, 25, 26 }; // a warm pass past the slot's 7: its history is all + const n = switch (st.plan(&warm, &.{.{ 5, 7 }}, 0)) { + .here => |x| x, + else => return error.TestUnexpectedResult, + }; + try t.expectEqual(@as(u32, 7), n); + const kept = st.carry(n); // as takeStanding: the record is used up, the point stays + st.deinit(t.allocator); + st.rewind_at = kept; + st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 7 }}); + try t.expectEqual(@as(u32, 4), st.rewind_at); // a yield at a landed chunk keeps it too + st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 7 }}); + try t.expectEqualSlices(u32, &warm, st.tokens); + try t.expectEqual(@as(u32, 4), st.rewind_at); + const next = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply rendered without its reasoning + try t.expectEqual(Resume{ .back = 4 }, st.plan(&next, &.{.{ 5, 6 }}, 0)); + try t.expectEqual(Resume{ .here = 10 }, st.plan(&(warm ++ [_]u32{32}), &.{.{ 5, 7 }}, 0)); +} + +test "qwen27 resume: a rewind point past where the pass resumes, or none, is not carried" { + const st: Standing = .{ .rewind_at = 4 }; + try t.expectEqual(@as(u32, 4), st.carry(4)); // a `.back` pass resumes at the point itself + try t.expectEqual(@as(u32, 4), st.carry(9)); + try t.expectEqual(@as(u32, 0), st.carry(3)); + try t.expectEqual(@as(u32, 0), st.carry(0)); + try t.expectEqual(@as(u32, 0), (Standing{}).carry(9)); +} + +test "qwen27 resume: a warm pass over slot 0's conversation runs where it carries the rewind point, else steps aside" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply as the template renders it: it goes back to the point + try t.expectEqual(Resume{ .back = 4 }, st.plan(&warm, &.{.{ 5, 6 }}, 0)); + try t.expectEqual(@as(u32, 4), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0))); // it runs + try t.expectEqual(@as(u32, 0), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside + const kept = st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end + st.deinit(t.allocator); + st.rewind_at = kept; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); + try t.expectEqual(Resume{ .here = 7 }, st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0)); // a client replaying it + try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 20, 21, 22, 23, 40 }, &.{}, 0)); // one that does not + st.rewind_at = 0; // a slot with no rewind point would keep none through the pass: it steps aside + try t.expectEqual(@as(u32, 0), st.carried(st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0))); + try t.expectEqual(@as(u32, 0), st.carried(.none)); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 2e24131464e..b307e195c09 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -421,7 +421,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q if (model.ane) |share| share.footprint_growth = platform.memory_policy.growth(before, cache_fit.footprint()); } ane.end(); - // Slot 0's history rewind when two full-context copies do not fit, wired before cache sizing. + // Slot 0's history rewind, wired before cache sizing: even a cache with two full-context copies may evict it. const rewind_room = laterRoom(h); if (slot_state.rewindInit(h.back, rewind_room, if (o.cache_gib) |g| gibBytes(g) else cache_fit.defaultBudget(rewind_room))) |buf| if (h.resident) |*r| r.addNow(&.{buf}); try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 2df612a9944..c893a9c0788 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -27,13 +27,14 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { const ids = s.prompt(); if (ids.len == 0 or ids.len + s.max_new + b.window > b.runner.capacity) return error.PromptTooLong; if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot - // as does one over the conversation slot 0 holds when the cache cannot keep its end: it would leave nothing there - if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none) return error.Yielded; + const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; + // as does one over slot 0's conversation when the cache cannot keep its end, unless a wide one and the point allow + if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none and + (!b.standing.wide or b.standing.carried(resumable(b, s, kept_at)) == 0)) return error.Yielded; if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); - const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; const resumed = if (slot == 0) try takeStanding(b, s, kept_at) else 0; // slot 0's own state: nothing copied if (resumed == 0) { if (d) |x| try x.reset() else try b.runner.reset(slot); @@ -87,7 +88,7 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill, _: u32) !bool { return error.Yielded; } const at = f.at; - var end = resume_rules.chunkEnd(at, @import("qwen27_prompt_rows.zig").choose(b, s), ids.len, s.chunks, s.decode_spans); + var end = resume_rules.chunkEnd(at, @import("qwen27_prompt_rows.zig").choose(b, s, at), ids.len, s.chunks, s.decode_spans); if (f.checkpoint > at and f.checkpoint < end) end = f.checkpoint; end = resume_rules.rewindEnd(at, end, f.rewind); var decoded = false; @@ -233,21 +234,27 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules return b.standing.plan(s.prompt(), s.decode_spans, kept_at); } -/// Slot 0's position for `s` when it resumes there (0: not); the record is used up, as the pass changes the slot. +/// Slot 0's position for `s` when it resumes there (0: not); the record is used up, its rewind point carried on. fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { - defer b.standing.deinit(b.gpa); - switch (resumable(b, s, kept_at)) { + const plan = resumable(b, s, kept_at); + const kept = b.standing.carried(plan); + b.standing.deinit(b.gpa); + // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it + switch (plan) { .none => return 0, .here => |n| { - const d = b.draft orelse return n; - try d.model.reset(); // a new session over the kept taps, as after a restore - d.model.committed_end = n; + if (b.draft) |d| { + try d.model.reset(); // a new session over the kept taps, as after a restore + d.model.committed_end = n; + } + b.standing.rewind_at = kept; return n; }, .back => |h| { const buf = b.rewind orelse return 0; if (b.draft) |d| try d.model.reset(); q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return 0; // the pass then starts at 0 + b.standing.rewind_at = kept; return h; }, } @@ -274,11 +281,12 @@ fn saveRewind(b: *Metal, at: u32) void { b.standing.rewind_at = at; } -/// The history rewind where two full-context states do not fit in `cache`, provided `room` holds its bytes. +/// The history rewind, where `room` holds it; beside a `cache` that keeps two full-context states it is also `wide`. pub fn rewindInit(b: *Metal, room: u64, cache: u64) ?mtl.Buffer { if (!slotResume(b) or b.runner.ready == 0) return null; const need = q.snapshot.rewindBytes(b.runner, b.ringOf(0)); - if (!resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity)))) return null; + // such a cache keeps the reply's end but may still evict the history's end: slot 0 keeps that one in place + b.standing.wide = !resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); if (need > room) { std.log.info("qwen27: no rewind point for slot 0 ({d} MiB, {d} MiB left under the load limit): a next turn that drops the last reply's reasoning reads its prompt again", .{ need >> 20, room >> 20 }); return null; diff --git a/zig/src/qwen27_rows_test.zig b/zig/src/qwen27_rows_test.zig index 43c94d1f9b1..4d672e3568d 100644 --- a/zig/src/qwen27_rows_test.zig +++ b/zig/src/qwen27_rows_test.zig @@ -43,32 +43,59 @@ test "qwen27 rows: a decoder arriving between chunks caps the next one and leavi var b: Backend = .{ .owners = .{ &prompt, null, null } }; defer b.fills.deinit(a); try b.fills.put(a, &prompt, true); - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); b.owners[1] = &other; try b.fills.put(a, &other, true); - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); try std.testing.expect(b.fills.remove(&other)); - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); other.paused = true; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); other.paused = false; - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); other.finished = true; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); other.finished = false; for ([_]usize{ 16, 128, 256, 512, 1024, 2048 }) |capacity| { b.chunk = capacity; - try std.testing.expectEqual(@min(capacity, @as(usize, 512)), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@min(capacity, @as(usize, 512)), policy.choose(&b, &prompt, 0)); } b.chunk = 2048; b.owners[1] = null; b.owners[2] = &other; b.runner.slots = 2; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); b.runner.slots = 3; - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); b.owners[2] = null; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); +} + +fn never(_: *anyopaque) bool { + return false; +} + +test "qwen27 rows: a background pass narrows its chunks with depth and a foreground one does not" { + var prompt = try Stream.init(a, .{ .id = "prompt", .prompt = &.{ 1, 2 }, .max_new = 4, .drafts = false }); + defer prompt.deinit(a); + var b: Backend = .{ .owners = .{ &prompt, null, null } }; + defer b.fills.deinit(a); + try b.fills.put(a, &prompt, true); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 120000)); // a foreground pass + prompt.yield_check = .{ .ptr = &b, .check = never }; + for ([_]usize{ 0, 16384 }) |at| try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, at)); + try std.testing.expectEqual(@as(usize, 1664), policy.choose(&b, &prompt, 20000)); // 1677, a whole 128 rows + try std.testing.expectEqual(@as(usize, 1024), policy.choose(&b, &prompt, 32768)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 131072)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 262144)); // never below 256 + b.chunk = 128; // nor wider than the chunk + try std.testing.expectEqual(@as(usize, 128), policy.choose(&b, &prompt, 131072)); + b.chunk = 2048; + var other = try Stream.init(a, .{ .id = "other", .prompt = &.{ 1, 2 }, .max_new = 4, .drafts = false }); + defer other.deinit(a); + b.owners[1] = &other; // a decoder beside it caps the chunk first + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 20000)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 131072)); } const floor = @import("native/qwen27_prompt_floor.zig"); From a94a16638fe960e00588bca6bdf17c2109e082b0 Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Tue, 10 Nov 2026 13:04:54 +0200 Subject: [PATCH 124/186] qwen27: the warm pass's yield, the rewind carry and the rewind buffer's choice are tested rules Standing.warmYields holds the warm pass's yield in opening, Standing.take holds takeStanding's use of the record and its carried rewind point, and makesRewind holds rewindInit's choice: the buffer wherever room holds it, tight where needsRewind says the cache cannot keep two full-context states, else wide. Each has a test that fails when its rule is reverted. The changelog takes the line for #2723 under a new Unreleased heading. --- CHANGELOG.md | 2 + zig/src/native/qwen27_resume.zig | 71 +++++++++++++++++++++++++++++--- zig/src/native/qwen27_slots.zig | 26 +++++++----- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14c6fc5b247..c8165118cdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -151,6 +151,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - HIP on gfx12 under `--policy matrix=on`: Qwen3 dense checkpoints load each layer's gate and up weights interleaved as one weight, and one WMMA product writes silu(gate) x up in its epilogue instead of two products and `tf_silu_mul`, with the same output bytes: Qwen3-Reranker-8B ranks 16 documents in 0.657 to 0.661 s instead of 0.687 to 0.697 s and Qwen3-Embedding-8B embeds 100 short inputs in 0.901 to 0.911 s instead of 0.940 to 0.948 s on a Radeon AI PRO R9700. The kernel contract checks both fused blocks byte for byte against the unfused launches (part of #2245). Thanks to @laudney. - Nemotron on Macs: the masked prompt shader calls the shared `tf_paged_row` page helper again, so its prompt sources compile; #2322 had used the name #2283 retired the same day (#2548). +- Qwen3.8-27B on Metal: a follow-up turn whose client sends the last reply without its reasoning resumes slot 0 in place from its rewind point, now also beside a prompt cache large enough for two whole contexts (which can still evict the history's end), instead of a turn further back; there the warm pass keeps that point, and a background pass at depth runs in narrower chunks so a waiting request is held for less (#2723). Thanks to @ThinkOffApp. + ## 1.2.7 - GLM-5.3-Flash on CUDA: with `TF_GLM_LAYERS=N`, serve's start-up memory check counts the weights by tensor from the shard headers, only those the load reads (the first N layers, the MTP layer when drafts are on, the embedding, norm and head), where it counted every checkpoint file. A 20-layer thin load of the 163.6 GiB checkpoint now counts 66.6 GiB without drafts. Unset, the count is the files' as before (#2788, fixes #2745). Thanks to @wojo. diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index 7a45c8e5f40..f8ac6cefce3 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -59,6 +59,18 @@ pub const Standing = struct { }; } + /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point the pass carries (0: none). + pub fn take(st: *Standing, gpa: std.mem.Allocator, r: Resume) void { + const kept = st.carried(r); + st.deinit(gpa); + st.rewind_at = kept; + } + + /// Whether a warm pass over slot 0 (`own`) yields: not if the cache keeps its end, nor `wide` with `from` carrying. + pub fn warmYields(st: *const Standing, keeps_end: bool, own: Resume, from: Resume) bool { + return !keeps_end and own != .none and (!st.wide or st.carried(from) == 0); + } + /// How `ids` (its rows' spans `spans`) resumes, past the cache's own restore point `kept_at`, by the cache's rule. pub fn plan(st: *const Standing, ids: []const u32, spans: []const modes.Span, kept_at: u32) Resume { const n = st.tokens.len; @@ -101,6 +113,15 @@ pub fn rewindEnd(at: usize, end: usize, rewind: u32) usize { return if (rewind > at and rewind < end) rewind else end; } +/// Slot 0's `need`-byte history rewind: none, a `tight` one, or a `wide` one beside a cache that keeps two of them. +pub const Rewind = enum { none, tight, wide }; + +/// Slot 0 makes its rewind wherever `room` holds it: a `cache` keeping two `full`-byte states may still evict it. +pub fn makesRewind(need: u64, room: u64, cache: u64, full: u64) Rewind { + if (need > room) return .none; + return if (needsRewind(cache, full)) .tight else .wide; +} + /// A prompt's next physical chunk end, before cache marks or the slot's rewind point cut it further. pub fn chunkEnd(at: usize, chunk: usize, len: usize, starts: []const u32, spans: []const modes.Span) usize { var end = @min(len, at + chunk); @@ -210,9 +231,7 @@ test "qwen27 resume: a warm pass that only extends slot 0 keeps the reply's rewi else => return error.TestUnexpectedResult, }; try t.expectEqual(@as(u32, 7), n); - const kept = st.carry(n); // as takeStanding: the record is used up, the point stays - st.deinit(t.allocator); - st.rewind_at = kept; + st.take(t.allocator, .{ .here = n }); // as takeStanding: the record is used up, the point stays st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 7 }}); try t.expectEqual(@as(u32, 4), st.rewind_at); // a yield at a landed chunk keeps it too st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 7 }}); @@ -241,9 +260,7 @@ test "qwen27 resume: a warm pass over slot 0's conversation runs where it carrie try t.expectEqual(Resume{ .back = 4 }, st.plan(&warm, &.{.{ 5, 6 }}, 0)); try t.expectEqual(@as(u32, 4), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0))); // it runs try t.expectEqual(@as(u32, 0), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside - const kept = st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end - st.deinit(t.allocator); - st.rewind_at = kept; + st.take(t.allocator, st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); try t.expectEqual(Resume{ .here = 7 }, st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0)); // a client replaying it try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 20, 21, 22, 23, 40 }, &.{}, 0)); // one that does not @@ -252,6 +269,48 @@ test "qwen27 resume: a warm pass over slot 0's conversation runs where it carrie try t.expectEqual(@as(u32, 0), st.carried(.none)); } +test "qwen27 resume: rule a, a warm pass carrying slot 0's rewind point runs beside a wide cache, else it yields" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // back to the point at 4 + const own = st.plan(&warm, &.{.{ 5, 6 }}, 0); + try t.expect(st.warmYields(false, own, own)); // a tight cache would lose the reply's end it leaves: aside + st.wide = true; // the cache keeps two full-context states + try t.expect(!st.warmYields(false, own, own)); // it carries the point: it runs + try t.expect(st.warmYields(false, own, st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside + try t.expect(!st.warmYields(true, own, .none)); // the cache keeps the conversation's end + try t.expect(!st.warmYields(false, .none, .none)); // it would not resume in slot 0 + st.rewind_at = 0; + try t.expect(st.warmYields(false, own, own)); // no point to carry: it would leave nothing in slot 0 + st.deinit(t.allocator); + try t.expect(st.wide); // the record goes, the cache's size stays +} + +test "qwen27 resume: rule b, taking slot 0's record for a pass that extends it keeps the rewind point" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + st.take(t.allocator, st.plan(&[_]u32{ 1, 2, 3, 4, 20, 21, 22, 24 }, &.{.{ 5, 7 }}, 0)); // `.here` at 7 + try t.expectEqual(@as(usize, 0), st.tokens.len); // the record is used up + try t.expectEqual(@as(u32, 4), st.rewind_at); // the point stays + st.take(t.allocator, .none); // a pass that does not resume in slot 0 drops it + try t.expectEqual(@as(u32, 0), st.rewind_at); +} + +test "qwen27 resume: rule c, slot 0 makes its rewind buffer where room holds it, even beside a cache with two states" { + const full: u64 = 2496 << 20; // one full-context state at 33024 rows + try t.expectEqual(Rewind.wide, makesRewind(1 << 30, 2 << 30, 5 << 30, full)); // needsRewind says no: made anyway + try t.expectEqual(Rewind.wide, makesRewind(1 << 30, 2 << 30, 8 << 30, full)); + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 2 << 30, 4 << 30, full)); // needsRewind's own case + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 1 << 30, full, full)); + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 2 << 30, 0, full)); // no cache at all + try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 0, full)); // no room + try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 8 << 30, full)); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index c893a9c0788..e13b6755e1d 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -29,8 +29,8 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; // as does one over slot 0's conversation when the cache cannot keep its end, unless a wide one and the point allow - if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none and - (!b.standing.wide or b.standing.carried(resumable(b, s, kept_at)) == 0)) return error.Yielded; + if (slot == 0 and s.yield_check != null and s.max_new == 0 and + b.standing.warmYields(keepsEnd(s), resumable(b, s, 0), resumable(b, s, kept_at))) return error.Yielded; if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); @@ -237,8 +237,8 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules /// Slot 0's position for `s` when it resumes there (0: not); the record is used up, its rewind point carried on. fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { const plan = resumable(b, s, kept_at); - const kept = b.standing.carried(plan); - b.standing.deinit(b.gpa); + b.standing.take(b.gpa, plan); + errdefer b.standing.rewind_at = 0; // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it switch (plan) { .none => return 0, @@ -247,19 +247,23 @@ fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { try d.model.reset(); // a new session over the kept taps, as after a restore d.model.committed_end = n; } - b.standing.rewind_at = kept; return n; }, .back => |h| { - const buf = b.rewind orelse return 0; + const buf = b.rewind orelse return dropRewind(b); if (b.draft) |d| try d.model.reset(); - q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return 0; // the pass then starts at 0 - b.standing.rewind_at = kept; + q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return dropRewind(b); // starts at 0 return h; }, } } +/// A pass that could not rewind slot 0 starts it afresh, keeping no rewind point (0: resumed nowhere). +fn dropRewind(b: *Metal) u32 { + b.standing.rewind_at = 0; + return 0; +} + /// A stream leaves slot 0: what it stands at stays for the next prompt that extends it, if a fresh pass stands there. fn leave(b: *Metal, s: *const lanes.Stream) void { const at = b.runner.offsets[0]; @@ -285,9 +289,9 @@ fn saveRewind(b: *Metal, at: u32) void { pub fn rewindInit(b: *Metal, room: u64, cache: u64) ?mtl.Buffer { if (!slotResume(b) or b.runner.ready == 0) return null; const need = q.snapshot.rewindBytes(b.runner, b.ringOf(0)); - // such a cache keeps the reply's end but may still evict the history's end: slot 0 keeps that one in place - b.standing.wide = !resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); - if (need > room) { + const kind = resume_rules.makesRewind(need, room, cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); + b.standing.wide = kind == .wide; // such a cache keeps a reply's end, but may evict the history's end + if (kind == .none) { std.log.info("qwen27: no rewind point for slot 0 ({d} MiB, {d} MiB left under the load limit): a next turn that drops the last reply's reasoning reads its prompt again", .{ need >> 20, room >> 20 }); return null; } From 96ea6876f5f8bcf00f37a246689923748165b00b Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:35:51 +0200 Subject: [PATCH 125/186] qwen27: a turn that resumed inside a reply kept as generated makes no new rewind point Slot 0 remembers the rows of the reply it generated. A follow-up whose prompt replays that reply as generated and resumes, by the slot or a cache state, inside it or at its end skips the extra rewind stop at its history end, so its fill is the same as main's. A client that drops the reply's reasoning, a cold turn and a resume that goes back still make a fresh rewind point. --- zig/src/native/qwen27_resume.zig | 104 ++++++++++++++++++++++++++++++- zig/src/native/qwen27_slots.zig | 11 ++-- 2 files changed, 108 insertions(+), 7 deletions(-) diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index f8ac6cefce3..bba7f7d6bb7 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -27,6 +27,7 @@ pub const Standing = struct { tokens: []u32 = &.{}, spans: []modes.Span = &.{}, rewind_at: u32 = 0, // where the rewind buffer holds slot 0's DeltaNet states and taps (0: none) + reply: modes.Span = .{ 0, 0 }, // rows of a reply these tokens hold exactly as it was generated (empty: none) wide: bool = false, // the prompt cache keeps two full-context states: it holds a reply's end beside slot 0 /// Frees the record; `wide`, the cache's size, stays. @@ -39,11 +40,13 @@ pub const Standing = struct { /// What `l` left in slot 0 of `context` (its rows' spans `spans`), when a later prompt may resume it; else nothing. pub fn leave(st: *Standing, gpa: std.mem.Allocator, l: Leaving, context: []const u32, spans: []const modes.Span) void { const rewind = st.rewind_at; // a point this stream's own pass saved + const reply = st.reply; // a reply the pass resumed past and only extended st.deinit(gpa); if (!l.kept() or l.at > context.len) return; const canonical = (if (l.at <= l.prompt_len) modes.prefix(gpa, spans, l.at) else modes.reply(gpa, spans, l.prompt_len, l.at)) catch return; const tokens = gpa.dupe(u32, context[0..l.at]) catch return gpa.free(canonical); - st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .wide = st.wide }; + const own: modes.Span = if (l.at > l.prompt_len) .{ l.prompt_len, l.at } else if (reply[1] <= l.at) reply else .{ 0, 0 }; + st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .reply = own, .wide = st.wide }; } /// The rewind point a pass resuming slot 0 at `n` keeps while it only extends the slot (0: none). @@ -59,11 +62,25 @@ pub const Standing = struct { }; } - /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point the pass carries (0: none). + /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point and reply it carries. pub fn take(st: *Standing, gpa: std.mem.Allocator, r: Resume) void { const kept = st.carried(r); + const reply: modes.Span = switch (r) { + .here => |n| if (st.reply[1] <= n) st.reply else .{ 0, 0 }, + .back, .none => .{ 0, 0 }, + }; st.deinit(gpa); st.rewind_at = kept; + st.reply = reply; + } + + /// Before `take`: where `ids` (spans `spans`) sends back the reply these tokens hold, as it was generated. + pub fn replayed(st: *const Standing, ids: []const u32, spans: []const modes.Span) Replay { + const from, const end = st.reply; + if (end <= from or end > st.tokens.len) return .{}; + const same: u32 = @intCast(std.mem.findDiff(u32, st.tokens[0..end], ids) orelse end); + if (same <= from or !modes.equal(st.spans, spans, same)) return .{}; + return .{ .from = from, .to = if (same == end) std.math.maxInt(u32) else same }; } /// Whether a warm pass over slot 0 (`own`) yields: not if the cache keeps its end, nor `wide` with `from` carrying. @@ -122,6 +139,17 @@ pub fn makesRewind(need: u64, room: u64, cache: u64, full: u64) Rewind { return if (needsRewind(cache, full)) .tight else .wide; } +/// Resume rows (from, to] past a reply's prompt whose reply rows before them a prompt sends back as generated. +pub const Replay = struct { + from: u32 = 0, + to: u32 = 0, + + /// A pass from `at` (slot or cache) resumed in a reply sent back as generated: it makes no point (unused there). + pub fn holds(p: Replay, at: usize) bool { + return at > p.from and at <= p.to; + } +}; + /// A prompt's next physical chunk end, before cache marks or the slot's rewind point cut it further. pub fn chunkEnd(at: usize, chunk: usize, len: usize, starts: []const u32, spans: []const modes.Span) usize { var end = @min(len, at + chunk); @@ -311,6 +339,78 @@ test "qwen27 resume: rule c, slot 0 makes its rewind buffer where room holds it, try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 8 << 30, full)); } +test "qwen27 resume: a follow-up sending the reply back as generated, restored at the cache's reply end, makes no point" { + var st: Standing = .{ .rewind_at = 4 }; // the turn's pass saved slot 0's states at its history's end + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // the stream left one short of its reply's end + const next = context ++ [_]u32{ 30, 31 }; // the reply as generated, then a new user turn + const spans = [_]modes.Span{.{ 5, 8 }}; + try t.expectEqual(@as(Resume, .none), st.plan(&next, &spans, 7)); // the cache keeps the reply's end: its own path + const replay = st.replayed(&next, &spans); // as opening, before takeStanding uses the record up + st.take(t.allocator, .none); + try t.expectEqual(@as(u32, 0), st.rewind_at); + try t.expect(replay.holds(7)); // restored at 7, inside the generated rows 5-7: no stop at the history's end + try t.expect(!replay.holds(4)); // a cache state at the history's mark is prompt rows: the point is made + try t.expect(!replay.holds(0)); + const turn = next ++ [_]u32{ 40, 41, 42 }; // that pass, with no point, and its reply 40-42 + st.leave(t.allocator, .{ .at = 12, .prompt_len = 10, .context_len = 13, .fill_at = 10, .ended = true }, &turn, &spans); + try t.expectEqual(modes.Span{ 10, 12 }, st.reply); + const again = [_]modes.Span{ .{ 5, 8 }, .{ 10, 13 } }; + try t.expect(st.replayed(&(turn ++ [_]u32{50}), &again).holds(12)); // the next follow-up skips it too + const diverging = next ++ [_]u32{ 60, 42, 61 }; // the last reply rendered without its reasoning + const cut = [_]modes.Span{ .{ 5, 8 }, .{ 10, 12 } }; + try t.expectEqual(@as(Resume, .none), st.plan(&diverging, &cut, 9)); // no point: the cache's state, as main + try t.expect(!st.replayed(&diverging, &cut).holds(9)); // and that pass makes one + try t.expectEqual(Replay{}, st.replayed(&diverging, &.{})); // reply rows sent as prompt rows are not as generated +} + +test "qwen27 resume: a follow-up resuming in slot 0 past a reply kept as generated makes no point, carrying the old one" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = context ++ [_]u32{ 30, 31 }; // a warm pass that ran: it extends the slot past the reply + st.take(t.allocator, st.plan(&warm, &.{.{ 5, 8 }}, 0)); // `.here` at 7 + st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 8 }}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // a yield at a landed chunk keeps it + st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 8 }}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // and so does the pass to its prompt's end + const follow = warm ++ [_]u32{32}; + try t.expectEqual(Resume{ .here = 10 }, st.plan(&follow, &.{.{ 5, 8 }}, 0)); + const replay = st.replayed(&follow, &.{.{ 5, 8 }}); + try t.expect(replay.holds(10)); // in the slot, or a cache state at the warm pass's end + try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 50, 23, 51 }, &.{.{ 5, 6 }}, 0)); // reply dropped + st.take(t.allocator, .{ .here = 10 }); + try t.expectEqual(@as(u32, 4), st.rewind_at); // the older point stays for such a turn + st.leave(t.allocator, .{ .at = 3, .prompt_len = 5, .context_len = 5, .fill_at = 3, .ended = false }, &context, &.{}); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); // a reply past where the slot stands is not kept +} + +test "qwen27 resume: a client dropping the reasoning, a cold turn or a pass that resumes nowhere makes a fresh point" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + try t.expect(!st.replayed(&.{ 1, 2, 3 }, &.{}).holds(0)); // cold: nothing stands + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + const reply: Leaving = .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }; + st.leave(t.allocator, reply, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply rendered without its reasoning + try t.expect(!st.replayed(&(warm ++ [_]u32{32}), &.{.{ 5, 6 }}).holds(4)); // its follow-up goes back to 4 + const back = st.plan(&warm, &.{.{ 5, 6 }}, 0); // the server's warm pass over it + try t.expectEqual(Resume{ .back = 4 }, back); + st.take(t.allocator, back); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); // the rows past the point are rewritten + st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); + const follow = warm ++ [_]u32{ 32, 33 }; + try t.expectEqual(Resume{ .here = 7 }, st.plan(&follow, &.{.{ 5, 6 }}, 0)); // past the carried point at 4 + try t.expect(!st.replayed(&follow, &.{.{ 5, 6 }}).holds(7)); // prompt rows: the new history's end is a point + st.leave(t.allocator, reply, &context, &.{}); + st.take(t.allocator, .none); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); + try t.expectEqual(@as(u32, 0), st.rewind_at); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index e13b6755e1d..d82558d64da 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -35,6 +35,7 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); + const replay: resume_rules.Replay = if (slot == 0) b.standing.replayed(ids, s.decode_spans) else .{}; // before take const resumed = if (slot == 0) try takeStanding(b, s, kept_at) else 0; // slot 0's own state: nothing copied if (resumed == 0) { if (d) |x| try x.reset() else try b.runner.reset(slot); @@ -58,7 +59,7 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { s.fill_at.store(@intCast(at), .monotonic); // what the restore kept: 0 when it failed const encoded = if (b.vision) |v| try images.encode(v, b.runner.model.queue, s, slot) else null; if (encoded != null) b.image_fill = s; // the vision rows are this pass's until it closes - return .{ .slot = slot, .at = at, .from = at, .images = encoded, .checkpoint = resume_rules.cachePoint(slot == 0 and slotResume(b) and plain(s) and s.max_new > 0 and s.reuse.hook != null and b.checkpoint_hook != null, at, b.chunk, ids.len, s.history_len, s.reuse.marks, s.chunks, s.decode_spans), .rewind = rewindPoint(b, s, slot, at) }; + return .{ .slot = slot, .at = at, .from = at, .images = encoded, .checkpoint = resume_rules.cachePoint(slot == 0 and slotResume(b) and plain(s) and s.max_new > 0 and s.reuse.hook != null and b.checkpoint_hook != null, at, b.chunk, ids.len, s.history_len, s.reuse.marks, s.chunks, s.decode_spans), .rewind = rewindPoint(b, s, slot, at, replay) }; } /// The pass is over (done, failed or cancelled): its images' rows are free for the next image prompt. @@ -238,7 +239,7 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { const plan = resumable(b, s, kept_at); b.standing.take(b.gpa, plan); - errdefer b.standing.rewind_at = 0; + errdefer b.standing = .{ .wide = b.standing.wide }; // tokens and spans freed by take: the point and reply go too // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it switch (plan) { .none => return 0, @@ -272,9 +273,9 @@ fn leave(b: *Metal, s: *const lanes.Stream) void { b.standing.leave(b.gpa, .{ .at = at, .prompt_len = @intCast(s.prompt_len), .context_len = @intCast(s.context.items.len), .fill_at = s.fill_at.load(.monotonic), .ended = ended }, s.context.items, s.decode_spans); } -/// Where slot 0's pass keeps its rewind point: its rendered history's end, which a next turn shares (0: none). -fn rewindPoint(b: *const Metal, s: *const lanes.Stream, slot: u32, at: usize) u32 { - if (slot != 0 or b.rewind == null or !slotResume(b) or !plain(s)) return 0; +/// Where slot 0's pass from `at` keeps its rewind point: its history's end (0: none, or it resumed in a `replay`). +fn rewindPoint(b: *const Metal, s: *const lanes.Stream, slot: u32, at: usize, replay: resume_rules.Replay) u32 { + if (slot != 0 or b.rewind == null or !slotResume(b) or !plain(s) or replay.holds(at)) return 0; return if (s.history_len > at and s.history_len < s.prompt().len) s.history_len else 0; } From 698e30174230fe71ec1cef0df32f8380c127b14a Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:10:02 +0000 Subject: [PATCH 126/186] changelog: #2723's line from lines/2723.txt in place of its own --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c8165118cdd..cbd9a4b64fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. + - Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. - GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. @@ -151,7 +153,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - HIP on gfx12 under `--policy matrix=on`: Qwen3 dense checkpoints load each layer's gate and up weights interleaved as one weight, and one WMMA product writes silu(gate) x up in its epilogue instead of two products and `tf_silu_mul`, with the same output bytes: Qwen3-Reranker-8B ranks 16 documents in 0.657 to 0.661 s instead of 0.687 to 0.697 s and Qwen3-Embedding-8B embeds 100 short inputs in 0.901 to 0.911 s instead of 0.940 to 0.948 s on a Radeon AI PRO R9700. The kernel contract checks both fused blocks byte for byte against the unfused launches (part of #2245). Thanks to @laudney. - Nemotron on Macs: the masked prompt shader calls the shared `tf_paged_row` page helper again, so its prompt sources compile; #2322 had used the name #2283 retired the same day (#2548). -- Qwen3.8-27B on Metal: a follow-up turn whose client sends the last reply without its reasoning resumes slot 0 in place from its rewind point, now also beside a prompt cache large enough for two whole contexts (which can still evict the history's end), instead of a turn further back; there the warm pass keeps that point, and a background pass at depth runs in narrower chunks so a waiting request is held for less (#2723). Thanks to @ThinkOffApp. ## 1.2.7 From 2c9c2992dfe21ce69f5d5cdeeeea7b4803ec2605 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:05:57 +0900 Subject: [PATCH 127/186] cuda source check: declare the SDK names main's kernels use check_cuda_sources.py failed 66 of its 918 checks on main: eleven units use SDK names that the fixed declarations in tools/zig/cuda_host lack. Declare each as CUDA 13.0 does: __syncthreads_or and __double2float_rz from sm_20_intrinsics.h, __float2int_rn from crt/device_functions.h, __float2bfloat16_rz from cuda_bf16.h, __hmul on half from cuda_fp16.h, ldexpf from crt/math_functions.h, and the __maxnreg__ macro from crt/host_defines.h, whose maxnreg attribute Clang ignores with a warning. Only the check reads these headers: no kernel, engine or test binary changes. --- tools/zig/cuda_host/cuda_bf16.h | 1 + tools/zig/cuda_host/cuda_fp16.h | 1 + tools/zig/cuda_host/cuda_runtime.h | 3 +++ tools/zig/cuda_host/math.h | 1 + 4 files changed, 6 insertions(+) diff --git a/tools/zig/cuda_host/cuda_bf16.h b/tools/zig/cuda_host/cuda_bf16.h index 9f96566052c..3d215164188 100644 --- a/tools/zig/cuda_host/cuda_bf16.h +++ b/tools/zig/cuda_host/cuda_bf16.h @@ -4,6 +4,7 @@ struct __nv_bfloat16 { unsigned short x; }; struct __nv_bfloat162 { __nv_bfloat16 x,y; }; __device__ float __bfloat162float(__nv_bfloat16); __device__ __nv_bfloat16 __float2bfloat16(float); __device__ __nv_bfloat16 __float2bfloat16_rn(float); __device__ __nv_bfloat16 __ushort_as_bfloat16(unsigned short); +__device__ __nv_bfloat16 __float2bfloat16_rz(float); __device__ unsigned short __bfloat16_as_ushort(__nv_bfloat16); __device__ float2 __bfloat1622float2(__nv_bfloat162); __device__ __nv_bfloat162 __floats2bfloat162_rn(float,float); diff --git a/tools/zig/cuda_host/cuda_fp16.h b/tools/zig/cuda_host/cuda_fp16.h index d423d31f12d..44654bcdacc 100644 --- a/tools/zig/cuda_host/cuda_fp16.h +++ b/tools/zig/cuda_host/cuda_fp16.h @@ -14,3 +14,4 @@ __device__ half2 __halves2half2(half,half); __device__ half __low2half(half2); _ __device__ float __low2float(half2); __device__ float __high2float(half2); __device__ half2 __hadd2(half2,half2); __device__ half2 __hsub2(half2,half2); __device__ half2 __hmul2(half2,half2); __device__ half2 __hfma2(half2,half2,half2); __device__ half __hadd(half,half); +__device__ half __hmul(half,half); diff --git a/tools/zig/cuda_host/cuda_runtime.h b/tools/zig/cuda_host/cuda_runtime.h index a6fd8663516..2a63e2b086e 100644 --- a/tools/zig/cuda_host/cuda_runtime.h +++ b/tools/zig/cuda_host/cuda_runtime.h @@ -8,6 +8,7 @@ #define __noinline__ __attribute__((noinline)) #define __align__(n) __attribute__((aligned(n))) #define __launch_bounds__(...) __attribute__((launch_bounds(__VA_ARGS__))) +#define __maxnreg__(a) __attribute__((maxnreg(a))) #define __grid_constant__ #define __CUDACC__ 1 #include @@ -31,6 +32,7 @@ template __device__ T __shfl_xor_sync(unsigned,T,int,int=32); template __device__ T __shfl_up_sync(unsigned,T,unsigned,int=32); template __device__ T __shfl_down_sync(unsigned,T,unsigned,int=32); __device__ void __syncthreads(); __device__ int __syncthreads_count(int); __device__ void __syncwarp(unsigned=~0u); +__device__ int __syncthreads_or(int); __device__ void __threadfence(); __device__ void __threadfence_system(); __device__ void __trap(); __device__ void __nanosleep(unsigned); __device__ unsigned long long clock64(); __device__ unsigned __ballot_sync(unsigned,int); __device__ int __any_sync(unsigned,int); @@ -49,6 +51,7 @@ __device__ int __float_as_int(float); __device__ unsigned __float_as_uint(float); __device__ float __int_as_float(int); __device__ float __uint_as_float(unsigned); __device__ double __longlong_as_double(long long); __device__ long long __double2ll_rz(double); __device__ float __double2float_rn(double); +__device__ float __double2float_rz(double); __device__ int __float2int_rn(float); __device__ float __fadd_rn(float,float); __device__ float __fsub_rn(float,float); __device__ float __fmul_rn(float,float); __device__ float __fdiv_rn(float,float); __device__ float __fsqrt_rn(float); __device__ float __fmaf_rn(float,float,float); __device__ float __expf(float); diff --git a/tools/zig/cuda_host/math.h b/tools/zig/cuda_host/math.h index 4aade403d68..65cf5c4b1c8 100644 --- a/tools/zig/cuda_host/math.h +++ b/tools/zig/cuda_host/math.h @@ -19,5 +19,6 @@ template __host__ __device__ bool isfinite(T); #define M_2_SQRTPI 1.12837916709551257390 #define M_SQRT1_2 0.70710678118654752440 __host__ __device__ float erff(float); __host__ __device__ float rintf(float); +__host__ __device__ float ldexpf(float,int); __host__ __device__ double pow(double,double); __host__ __device__ double sin(double); __host__ __device__ double cos(double); template __host__ __device__ bool signbit(T); From 8fda4b71eb76208e650340f45cac4357624741b2 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:10:06 +0000 Subject: [PATCH 128/186] changelog: #2902 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbd9a4b64fe..92e52fa03c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. + - Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. - Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. From cbf0ceb65592fcb33539c04d6dd429e3a9b95556 Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:35:35 +0700 Subject: [PATCH 129/186] native: the CUDA kept-state identity is written once through any hasher, its keys unchanged --- CHANGELOG.md | 1 + zig/src/cuda/aot_identity.zig | 9 +- zig/src/native/cuda_disk.zig | 59 +++++++-- zig/src/native/cuda_disk_golden_test.zig | 143 +++++++++++++++++++++ zig/src/native/cuda_disk_identity_test.zig | 4 + zig/src/native/disk_identity.zig | 10 +- zig/src/platform/file_identity.zig | 12 +- 7 files changed, 215 insertions(+), 23 deletions(-) create mode 100644 zig/src/native/cuda_disk_golden_test.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e52fa03c7..b6eddd97525 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -126,6 +126,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. +- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. diff --git a/zig/src/cuda/aot_identity.zig b/zig/src/cuda/aot_identity.zig index 6c6e8e001af..a1b21900072 100644 --- a/zig/src/cuda/aot_identity.zig +++ b/zig/src/cuda/aot_identity.zig @@ -3,13 +3,19 @@ const std = @import("std"); const aot = @import("aot.zig"); pub fn hash(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { + var h = std.hash.Wyhash.init(0x616f74); + try write(&h, a, io, dir); + return h.final(); +} + +/// The manifest's text, then each cubin it names by its SHA-256, into any hasher with `update(bytes)`. +pub fn write(h: anytype, a: std.mem.Allocator, io: std.Io, dir: []const u8) !void { const path = try std.fs.path.join(a, &.{ dir, "aot.json" }); defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 24)); defer a.free(text); const parsed = try aot.parseSpecs(a, text); defer parsed.deinit(); - var h = std.hash.Wyhash.init(0x616f74); h.update(text); for (parsed.value.kernels) |k| { const cubin = try std.fmt.allocPrint(a, "{s}/cubins/{s}.cubin", .{ dir, k.hash }); @@ -17,7 +23,6 @@ pub fn hash(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { const digest = try @import("sys").file_identity.hash(io, cubin, 1 << 26); h.update(&digest); } - return h.final(); } test "persistence identity: captured CUDA cubin replacement moves an unchanged manifest key" { diff --git a/zig/src/native/cuda_disk.zig b/zig/src/native/cuda_disk.zig index 4ff1ee1e14b..0c6c03d8219 100644 --- a/zig/src/native/cuda_disk.zig +++ b/zig/src/native/cuda_disk.zig @@ -133,24 +133,57 @@ pub fn keptKey(comptime K: type, a: Allocator, io: std.Io, g: *const cuda.Contex /// The identity's hash: the checkpoint, kernels, device, driver, a family's facts, then its probe's hash. pub fn identityOf(a: Allocator, io: std.Io, g: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, facts: []const u64, probed: u64, profile: disk_identity.Profile) !u64 { var h = std.hash.Wyhash.init(0x6375); - try @import("sys").file_identity.current(&h, a, io); - try disk_identity.checkpoint(&h, a, io, dir); + const binary = try @import("sys").file_identity.running(a, io); + var name_buf: [256]u8 = undefined; + try components(&h, a, io, .{ .binary = binary, .checkpoint = dir, .kernels = kernels, .device = try deviceOf(g, &name_buf), .facts = facts, .probed = probed, .profile = profile }); + return h.final(); +} + +/// The GPU's name (or "GPU" when the driver gives none), compute capability and driver version. +pub const Device = struct { name: []const u8, capability: u64, driver: u64 }; + +pub fn deviceOf(g: *const cuda.Context, name_buf: []u8) !Device { + const name = g.name(name_buf) catch "GPU"; + return .{ .name = name, .capability = try g.capability(), .driver = @intCast(try g.d.version()) }; +} + +/// What a kept key binds; a null part is bound otherwise by its caller (the witness binds bytes, not stamps). +pub const Components = struct { + binary: ?[32]u8, + checkpoint: ?[]const u8, + kernels: ?[]const u8, + device: Device, + facts: []const u64, + probed: ?u64, + profile: disk_identity.Profile, +}; + +/// The components in one order into any hasher with `update(bytes)`: a kept key's Wyhash, or a SHA-256. +pub fn components(h: anytype, a: Allocator, io: std.Io, c: Components) !void { + if (c.binary) |digest| @import("sys").file_identity.bind(h, digest); + if (c.checkpoint) |dir| try disk_identity.checkpoint(h, a, io, dir); inline for (comptime std.meta.declarations(cuda.kernels)) |d| { const v = @field(cuda.kernels, d); if (comptime @TypeOf(v) == []const u8) h.update(v); } - if (kernels) |k| { - const captured = try cuda.aot.identity.hash(a, io, k); - h.update(std.mem.asBytes(&captured)); - } - var name_buf: [256]u8 = undefined; - h.update(g.name(&name_buf) catch "GPU"); - const device = [_]u64{ try g.capability(), @intCast(try g.d.version()) }; + if (c.kernels) |k| try captured(h, a, io, k); + h.update(c.device.name); + const device = [_]u64{ c.device.capability, c.device.driver }; h.update(std.mem.sliceAsBytes(&device)); - h.update(std.mem.sliceAsBytes(facts)); - h.update(std.mem.asBytes(&probed)); - disk_identity.profile(&h, profile); - return h.final(); + h.update(std.mem.sliceAsBytes(c.facts)); + if (c.probed) |p| h.update(std.mem.asBytes(&p)); + disk_identity.profile(h, c.profile); +} + +/// A captured set as a nested digest: a kept key folds its 64-bit Wyhash, a SHA-256 its SHA-256. +fn captured(h: anytype, a: Allocator, io: std.Io, dir: []const u8) !void { + if (@TypeOf(h.*) == std.hash.Wyhash) { + const v = try cuda.aot.identity.hash(a, io, dir); + return h.update(std.mem.asBytes(&v)); + } + var s = std.crypto.hash.sha2.Sha256.init(.{}); + try cuda.aot.identity.write(&s, a, io, dir); + h.update(&s.finalResult()); } /// A fixed prompt's state at a mark inside its last chunk, hashed (a kernel, driver or arithmetic change moves it). diff --git a/zig/src/native/cuda_disk_golden_test.zig b/zig/src/native/cuda_disk_golden_test.zig new file mode 100644 index 00000000000..9e5ce7b5a12 --- /dev/null +++ b/zig/src/native/cuda_disk_golden_test.zig @@ -0,0 +1,143 @@ +//! Kept keys stay byte for byte: the identity, however it is written, matches the sequence main hashed. +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("engine_api"); +const disk = @import("cuda_disk.zig"); +const identity = @import("disk_identity.zig"); +const t = std.testing; + +const Driver = struct { + fn name(out: [*]u8, _: c_int, _: cuda.abi.Device) callconv(.c) cuda.abi.Result { + @memcpy(out[0..5], "fake\x00"); + return 0; + } + fn attribute(out: *c_int, which: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result { + out.* = if (which == .compute_capability_major) 12 else 1; + return 0; + } + fn version(out: *c_int) callconv(.c) cuda.abi.Result { + out.* = 13000; + return 0; + } + fn environment(_: [:0]const u8) ?[:0]const u8 { + return "1"; + } +}; + +const row: api.profile.Entry = .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = &.{"qwen3_5"}, .class = .arithmetic, .syntax = .one, .effect = "test arithmetic", .evidence = "test fixture" }; +const profile: identity.Profile = .{ .family = "qwen3_5", .backend = "cuda", .entries = &.{row}, .environment = Driver.environment }; +/// The key main (93a2e93ec) gives the fixed inputs below. +const main_key: u64 = 9113307822166948912; +const manifest = "{\"kernels\":[{\"fn\":\"f\",\"hash\":\"h\",\"name\":\"k\",\"num_warps\":4,\"params\":[],\"consts\":{}}]}"; + +/// main's identityOf, update for update, with the binary's digest and the device's facts given. +fn reference(a: std.mem.Allocator, binary: [32]u8, dir: []const u8, kernels: ?[]const u8, facts: []const u64, probed: u64) !u64 { + var h = std.hash.Wyhash.init(0x6375); + @import("sys").file_identity.bind(&h, binary); + try identity.checkpoint(&h, a, t.io, dir); + inline for (comptime std.meta.declarations(cuda.kernels)) |d| { + const v = @field(cuda.kernels, d); + if (comptime @TypeOf(v) == []const u8) h.update(v); + } + if (kernels) |k| { + const captured = try cuda.aot.identity.hash(a, t.io, k); + h.update(std.mem.asBytes(&captured)); + } + h.update("fake"); + const device = [_]u64{ 121, 13000 }; + h.update(std.mem.sliceAsBytes(&device)); + h.update(std.mem.sliceAsBytes(facts)); + h.update(std.mem.asBytes(&probed)); + identity.profile(&h, profile); + return h.final(); +} + +/// A checkpoint folder whose identity reads no stamps (config.json alone), and a captured set beside it. +const Fixture = struct { + tmp: t.TmpDir, + dir: []const u8, + kernels: []const u8, + + fn init(a: std.mem.Allocator) !Fixture { + var tmp = t.tmpDir(.{}); + errdefer tmp.cleanup(); + try tmp.dir.createDir(t.io, "model", .default_dir); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model/config.json", .data = "{\"model_type\":\"qwen3_5\"}" }); + try tmp.dir.createDir(t.io, "aot", .default_dir); + try tmp.dir.createDir(t.io, "aot/cubins", .default_dir); + try tmp.dir.writeFile(t.io, .{ .sub_path = "aot/aot.json", .data = manifest }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "aot/cubins/h.cubin", .data = "cubin bytes" }); + const root = try tmp.dir.realPathFileAlloc(t.io, ".", a); + return .{ .tmp = tmp, .dir = try std.fs.path.join(a, &.{ root, "model" }), .kernels = try std.fs.path.join(a, &.{ root, "aot" }) }; + } +}; + +test "disk identity: identityOf hashes main's sequence over this binary, a checkpoint and a captured set" { + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + try f.tmp.dir.writeFile(t.io, .{ .sub_path = "model/model.safetensors", .data = "shard" }); // its stamps are hashed too + var driver: cuda.Driver = undefined; + driver.api.cuDeviceGetName = Driver.name; + driver.api.cuDeviceGetAttribute = Driver.attribute; + driver.api.cuDriverGetVersion = Driver.version; + const gpu: cuda.Context = .{ .d = &driver, .device = 0, .handle = null }; + const exe = try std.process.executablePathAlloc(t.io, a); + const binary = try @import("sys").file_identity.hash(t.io, exe, 1 << 31); + for ([_]?[]const u8{ null, f.kernels }) |kernels| + try t.expectEqual(try reference(a, binary, f.dir, kernels, &.{ 1, 2, 3 }, 42), try disk.identityOf(a, t.io, &gpu, f.dir, kernels, &.{ 1, 2, 3 }, 42, profile)); +} + +test "disk identity: a fixed binary, checkpoint, captured set, device, facts, probe and profile give main's key" { + if (cuda.kernels.available) return error.SkipZigTest; // the constant is a host build's, whose kernel images are empty + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + try t.expectEqual(main_key, try reference(a, @splat(0x5a), f.dir, f.kernels, &.{ 1, 2, 3 }, 42)); + var h = std.hash.Wyhash.init(0x6375); + try disk.components(&h, a, t.io, fixed(f)); + try t.expectEqual(main_key, h.final()); +} + +fn fixed(f: Fixture) disk.Components { + return .{ .binary = @splat(0x5a), .checkpoint = f.dir, .kernels = f.kernels, .device = .{ .name = "fake", .capability = 121, .driver = 13000 }, .facts = &.{ 1, 2, 3 }, .probed = 42, .profile = profile }; +} + +fn sha(a: std.mem.Allocator, c: disk.Components) ![32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + try disk.components(&h, a, t.io, c); + return h.finalResult(); +} + +test "disk identity: the same components through SHA-256 move with a captured cubin, a profile value and the device" { + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + var c = fixed(f); + c.binary = null; // the witness binds the binary and the checkpoint's bytes elsewhere + c.checkpoint = null; + c.probed = null; + const base = try sha(a, c); + try t.expectEqualSlices(u8, &base, &try sha(a, c)); + try f.tmp.dir.writeFile(t.io, .{ .sub_path = "aot/cubins/h.cubin", .data = "other cubin" }); + const cubin = try sha(a, c); + try t.expect(!std.mem.eql(u8, &base, &cubin)); + var off = c; + off.profile.environment = Unset.environment; + try t.expect(!std.mem.eql(u8, &cubin, &try sha(a, off))); + var other = c; + other.device.capability = 80; + try t.expect(!std.mem.eql(u8, &cubin, &try sha(a, other))); +} + +const Unset = struct { + fn environment(_: [:0]const u8) ?[:0]const u8 { + return null; + } +}; diff --git a/zig/src/native/cuda_disk_identity_test.zig b/zig/src/native/cuda_disk_identity_test.zig index 107fd664053..7bfcdf9a5b6 100644 --- a/zig/src/native/cuda_disk_identity_test.zig +++ b/zig/src/native/cuda_disk_identity_test.zig @@ -6,6 +6,10 @@ const disk = @import("cuda_disk.zig"); const identity = @import("disk_identity.zig"); const t = std.testing; +test { + _ = @import("cuda_disk_golden_test.zig"); +} + const Recording = struct { var calls: [16]enum { name, major, minor, version } = undefined; var count: usize = 0; diff --git a/zig/src/native/disk_identity.zig b/zig/src/native/disk_identity.zig index 3126a14864b..46ff8cb44a2 100644 --- a/zig/src/native/disk_identity.zig +++ b/zig/src/native/disk_identity.zig @@ -4,7 +4,7 @@ const api = @import("engine_api"); const Allocator = std.mem.Allocator; /// Config/index bytes, stamps of every indexed shard (IndexedShardMissing if absent) and top-level safetensors. -pub fn checkpoint(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8) !void { +pub fn checkpoint(h: anytype, a: Allocator, io: std.Io, dir: []const u8) !void { var arena: std.heap.ArenaAllocator = .init(a); defer arena.deinit(); const scratch = arena.allocator(); @@ -46,7 +46,7 @@ pub fn checkpoint(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8 } /// This Mac's OS build and chip: a macOS update or another chip can compile a kernel to other bits. -pub fn machine(h: *std.hash.Wyhash, chip: []const u8) void { +pub fn machine(h: anytype, chip: []const u8) void { var os: [64]u8 = undefined; h.update(api.prompt_imprint.osBuild(&os)); h.update(chip); @@ -65,7 +65,7 @@ pub const Profile = struct { }; /// Only applicable non-identical entries partition kept state; defaults and `on_with` resolve before values are hashed. -pub fn profile(h: *std.hash.Wyhash, p: Profile) void { +pub fn profile(h: anytype, p: Profile) void { var began = false; for (p.entries) |e| { if (e.class == .identical or !e.applies(p.family, p.backend)) continue; @@ -95,7 +95,7 @@ pub fn profile(h: *std.hash.Wyhash, p: Profile) void { } } -fn valueField(h: *std.hash.Wyhash, syntax: api.profile.Syntax, value: ?[]const u8) void { +fn valueField(h: anytype, syntax: api.profile.Syntax, value: ?[]const u8) void { switch (syntax) { .one => field(h, if (value != null and std.mem.eql(u8, value.?, "1")) "1" else "0"), .auto => field(h, value orelse "auto"), @@ -117,7 +117,7 @@ fn valueField(h: *std.hash.Wyhash, syntax: api.profile.Syntax, value: ?[]const u } } -fn field(h: *std.hash.Wyhash, bytes: []const u8) void { +fn field(h: anytype, bytes: []const u8) void { var len: [8]u8 = undefined; std.mem.writeInt(u64, &len, @intCast(bytes.len), .little); h.update(&len); diff --git a/zig/src/platform/file_identity.zig b/zig/src/platform/file_identity.zig index fc40f934da8..4c126ed767a 100644 --- a/zig/src/platform/file_identity.zig +++ b/zig/src/platform/file_identity.zig @@ -22,14 +22,20 @@ pub fn hash(io: std.Io, path: []const u8, maximum: u64) ![32]u8 { return h.finalResult(); } -pub fn bind(h: *std.hash.Wyhash, digest: [32]u8) void { +/// `h`: any hasher with `update(bytes)`, a kept key's Wyhash or a SHA-256. +pub fn bind(h: anytype, digest: [32]u8) void { h.update("serving-binary-v1"); h.update(&digest); } /// The running binary includes compiled host code, embedded GPU images, compiler options and snapshot layouts. -pub fn current(h: *std.hash.Wyhash, a: std.mem.Allocator, io: std.Io) !void { +pub fn current(h: anytype, a: std.mem.Allocator, io: std.Io) !void { + bind(h, try running(a, io)); +} + +/// The running binary's SHA-256, read by its path. +pub fn running(a: std.mem.Allocator, io: std.Io) ![32]u8 { const path = try std.process.executablePathAlloc(io, a); defer a.free(path); - bind(h, try hash(io, path, 1 << 31)); + return hash(io, path, 1 << 31); } From 1885a1723a64813068584cedcc55127fc926250f Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:54:52 +0700 Subject: [PATCH 130/186] changelog: #2903's line names its pull request --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b6eddd97525..c2619030b6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -127,6 +127,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. - Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). +- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (#2903, refs #2741). Thanks to @kotevcode. - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. From 3f8a87a9671945e5acaf0b65204d538cfb7e403b Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:10:11 +0000 Subject: [PATCH 131/186] changelog: #2903's line from lines/2903.txt in place of its own --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c2619030b6f..bebfd2ca485 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. + - Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. - Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. @@ -126,8 +128,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. -- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). -- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (#2903, refs #2741). Thanks to @kotevcode. - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. From 78b814727ab29b0958fd39f5836fa96b5597811c Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 14:27:55 +0000 Subject: [PATCH 132/186] docs(xpu): the B70 prefill row predates 1.1.0's precise prompt DeltaNet (#2859) --- docs/xpu/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/xpu/README.md b/docs/xpu/README.md index bd649f64412..edddef11648 100644 --- a/docs/xpu/README.md +++ b/docs/xpu/README.md @@ -90,6 +90,11 @@ Options through environment variables: | 2 / 4 / 8 simultaneous streams, chat, greedy (aggregate) | 151 / 239 / 300 | 148 / 215 / 208 | 162 / 230 / 224 | 162 / 249 / 227 | | `bench_openai` decode, 1 / 2 / 4 / 8 requests (aggregate) | 173 / 243 / 343 / 350 | 165 / 300 / 337 / 439 | 165 / 336 / 367 / 456 | 179 / 325 / 364 / 463 | +The prefill row predates 1.1.0, which feeds the prompt DeltaNet's fp32 intermediates to the matrix units at close to fp32 +precision (`489baa2673`, `17facba3f8`). On a B70 with GPTQ g128 that costs about 5% of a 32,768-token prompt in the +raw runner (`xpu-qwen27 --bench=32768`: 1,647 to 1,571 tok/s, tokens unchanged; #2859). The table's sizes have not been +re-measured since. + Before the per-stream DFlash2 drafts (copies only in the shared rounds), 2 / 4 / 8 code streams ran at 64 / 120 / 182 tok/s in GPTQ. Validation: `zig build test`, oracle, drafted == serial, multi-stream == serial (chains and trees), and `bench_concurrent --alone --serial` on 1, 2, 4 and 8 streams, greedy and sampled, for the four From ef7d6926c6ff58a84fc2879204a59bcb72739cbe Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:10:16 +0000 Subject: [PATCH 133/186] changelog: #2906 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bebfd2ca485..9381b4e5820 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Docs: the Intel Arc (XPU) page notes that its B70 prefill row predates 1.1.0's more precise prompt DeltaNet, which costs about 5% of a 32,768-token prompt with tokens unchanged (#2906, measured in #2859). Thanks to @dinadur. + - Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. - Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. From cbdc1fe03d869882e950956ac14460df11241251 Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:31:24 -0800 Subject: [PATCH 134/186] tools: tfvideo_frames.py checks libtfvideo's RGB frames against PyAV's The video byte gate, tf-glm-video with glm_video_fixtures.py, builds only on macOS and needs a checkpoint's processor. This tool needs only a libtfvideo build and clips: it decodes every frame, GLM's 2 frames a second, and the first three and last frames, hashes each frame's RGB rows, and exits nonzero when any frame or return code differs from PyAV's or from a second build's. A clip that will not open or decode fails and the run goes on. --time-only times the decoder alone. Nothing in the engine or the build calls it, so no served path changes. --- tools/zig/tfvideo_frames.py | 177 ++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 tools/zig/tfvideo_frames.py diff --git a/tools/zig/tfvideo_frames.py b/tools/zig/tfvideo_frames.py new file mode 100644 index 00000000000..d53b11506a6 --- /dev/null +++ b/tools/zig/tfvideo_frames.py @@ -0,0 +1,177 @@ +"""libtfvideo's RGB frames against PyAV's (read_video_pyav) or another libtfvideo build, byte for byte, per clip.""" + +import argparse +import ctypes +import hashlib +import json +import sys +import time +from pathlib import Path + +TF_DECODE, TF_SHORT = -5, -6 + + +class Info(ctypes.Structure): + _fields_ = [("frames", ctypes.c_int64), ("rate_num", ctypes.c_int), ("rate_den", ctypes.c_int), + ("width", ctypes.c_int), ("height", ctypes.c_int), ("counted", ctypes.c_int)] + + +FRAME_FN = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, + ctypes.POINTER(ctypes.c_uint8), ctypes.c_int, ctypes.c_int, ctypes.c_int) + + +class Library: + """One libtfvideo build, loaded by path so two builds can share the process.""" + + def __init__(self, path: Path): + self.name = str(path) + self.lib = ctypes.CDLL(str(path.resolve()), mode=ctypes.RTLD_LOCAL) + self.lib.tf_video_open.argtypes = [ctypes.c_char_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(Info)] + self.lib.tf_video_frames.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.c_int, FRAME_FN, + ctypes.c_void_p] + self.lib.tf_video_close.argtypes = [ctypes.c_void_p] + self.lib.tf_video_version.restype = ctypes.c_char_p + self.version = self.lib.tf_video_version().decode() + + def open(self, data: bytes) -> tuple[int, Info]: + """tf_video_open's return code and the stream's facts.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc == 0: + self.lib.tf_video_close(handle) + return rc, info + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """The frames delivered (hashed, or index only), the seconds tf_video_frames took and its return code.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc != 0: + return [], 0.0, rc + out: list[tuple[int, str]] = [] + + def take(_ctx, _k, index, rgb, width, height, stride): + if not hashed: # timing: the bytes are left alone, so the decoder's threads never wait on Python + out.append((index, "")) + return 0 + raw = ctypes.string_at(rgb, stride * (height - 1) + width * 3) + out.append((index, hashlib.sha256(b"".join(raw[r * stride:r * stride + width * 3] + for r in range(height))).hexdigest())) + return 0 + + callback = FRAME_FN(take) + start = time.perf_counter() + rc = self.lib.tf_video_frames(handle, (ctypes.c_int * len(indices))(*indices), len(indices), callback, None) + seconds = time.perf_counter() - start + self.lib.tf_video_close(handle) + return out, seconds, rc + + +class PyAV: + """transformers' read_video_pyav: decode from the start, keep the wanted indices, to_ndarray(format="rgb24").""" + + def __init__(self): + import av + self.av, self.name, self.version = av, f"PyAV {av.__version__}", str(av.ffmpeg_version_info) + + def open(self, data: bytes) -> tuple[int, None]: + """0 when PyAV opens the container with a video stream, -2 (TF_NO_STREAM) without one, else -1 (TF_INVALID).""" + import io + try: + with self.av.open(io.BytesIO(data)) as c: + return (0 if c.streams.video else -2), None + except self.av.error.FFmpegError: + return -1, None + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """As Library.frames; an FFmpeg error counts as TF_DECODE and an early end as TF_SHORT.""" + import io + want, out, rc = set(indices), [], 0 + start = time.perf_counter() + try: + with self.av.open(io.BytesIO(data)) as c: + c.seek(0) + for i, frame in enumerate(c.decode(video=0)): + if i > indices[-1]: + break + if i in want: + rgb = frame.to_ndarray(format="rgb24") + out.append((i, hashlib.sha256(rgb.tobytes()).hexdigest() if hashed else "")) + except self.av.error.FFmpegError: + rc = TF_DECODE + seconds = time.perf_counter() - start + return out, seconds, rc or (len(out) if len(out) == len(indices) else TF_SHORT) + + +def index_sets(info: Info) -> dict[str, list[int]]: + """Every frame, GLM's 2 frames a second, the first three (an early stop) and the last (the end-of-stream drain).""" + n, rate = int(info.frames), info.rate_num / info.rate_den if info.rate_den else 0.0 + two = sorted({min(n - 1, round(i * rate / 2)) for i in range(int(n / rate * 2) + 1)}) if rate else [0] + return {"all": list(range(n)), "2fps": two, "first3": list(range(min(3, n))), "last": [n - 1]} + + +def best(source, data: bytes, indices: list[int], repeat: int, hashed: bool) -> tuple[list, float, int]: + frames, seconds, rc = source.frames(data, indices, hashed) + for _ in range(repeat - 1): + seconds = min(seconds, source.frames(data, indices, hashed)[1]) + return frames, seconds, rc + + +def compare(lib, ref, clip: Path, repeat: int, rows: dict) -> int: + """One clip's sets against the reference (or timed alone when ref is None); returns the sets that differ.""" + data = clip.read_bytes() + rc, info = lib.open(data) + ref_rc, ref_info = ref.open(data) if ref else (rc, info) + if rc != 0 or ref_rc != 0 or info.frames <= 0: # a damaged clip may not open, or may count no frames + rows["open"] = {"test_rc": rc, "reference_rc": ref_rc, "frames": info.frames} + theirs = (ref_info or info).frames # PyAV counts nothing here: libtfvideo's count stands for both + ok = ref is not None and rc == ref_rc and (info.frames > 0) == (theirs > 0) + print(f"{'ok ' if ok else 'FAIL'} {clip} open: test rc {rc}, reference rc {ref_rc}, {info.frames} frames") + return int(not ok) + failures = 0 + for name, indices in index_sets(info).items(): + mine, mine_s, mine_rc = best(lib, data, indices, repeat, ref is not None) + theirs, theirs_s, theirs_rc = best(ref, data, indices, repeat, True) if ref else (mine, 0.0, mine_rc) + same = sum(a == b for a, b in zip(mine, theirs)) + ok = mine == theirs and mine_rc == theirs_rc and (ref is not None or mine_rc == len(indices)) + failures += not ok + rows[name] = {"indices": indices, "test": mine, "reference": theirs if ref else None, "test_rc": mine_rc, + "reference_rc": theirs_rc, "test_s": mine_s, "reference_s": theirs_s} + codes = f", rc {mine_rc}/{theirs_rc}" if (mine_rc, theirs_rc) != (len(indices), len(indices)) else "" + timing = f"test {mine_s:.3f} s" + (f", reference {theirs_s:.3f} s" if ref else "") + what = "frames equal" if ref else "frames delivered" + print(f"{'ok ' if ok else 'FAIL'} {clip} {name}: {same}/{len(indices)} {what}{codes}, {timing}") + return failures + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--lib", required=True, type=Path, help="the libtfvideo build under test") + ap.add_argument("--against", type=Path, help="another libtfvideo build; PyAV when omitted") + ap.add_argument("--time-only", action="store_true", + help="no reference and no hashing: time tf_video_frames alone (its callback leaves the bytes)") + ap.add_argument("--repeat", type=int, default=1, help="decodes per set; the fastest time is reported") + ap.add_argument("--json", type=Path, help="write every frame's hash, return code and each set's time here") + ap.add_argument("clips", nargs="+", type=Path) + args = ap.parse_args() + lib = Library(args.lib) + ref = None if args.time_only else Library(args.against) if args.against else PyAV() + print(f"test {lib.name} (FFmpeg {lib.version}); reference " + + (f"{ref.name} (FFmpeg {ref.version})" if ref else "none (--time-only)")) + failures, report = 0, {"lib": lib.name, "reference": ref.name if ref else None, "clips": {}} + for clip in args.clips: + rows = report["clips"][str(clip)] = {} + try: + failures += compare(lib, ref, clip, args.repeat, rows) + except Exception as e: # one unreadable clip is a failed clip, not the end of the run + rows["error"] = repr(e) + failures += 1 + print(f"FAIL {clip}: {e!r}") + if args.json: + args.json.write_text(json.dumps(report, indent=1)) + print(f"{failures} set(s) differ" if failures else "every frame equal") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) From 752aa6b2f69dfa0223ed4dc699482d7cbcf9b21e Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:31:24 -0800 Subject: [PATCH 135/186] video: libtfvideo.so exports only tf_video_*, so FFmpeg's x86_64 asm links FFmpeg's x86_64 assembly addresses tables defined in C, such as ff_pw_* and ff_h264_cabac_tables, with 32-bit PC-relative relocations. Linked from the static libav*.a into a shared library, those tables are exported and so preemptible, and the link stopped on R_X86_64_PC32 errors (5,046 with gcc 13.3) although build_ffmpeg.sh passes --enable-pic. On ELF targets the library now links with a version script that exports the four tf_video_* symbols core/tfvideo.zig looks up and makes every other symbol local, so the tables bind inside the library with no text relocations. check-video-exports-linux checks this for aarch64 and x86_64 Linux with a planted table. The loader, build_ffmpeg.sh and its FFmpeg pin, and Mach-O builds are unchanged. --- CHANGELOG.md | 2 + tools/zig/check_tfvideo_exports.py | 70 ++++++++++++++++++++++++++++++ zig/build/video.zig | 17 ++++++++ zig/tests/tfvideo_planted.c | 14 ++++++ 4 files changed, 103 insertions(+) create mode 100644 tools/zig/check_tfvideo_exports.py create mode 100644 zig/tests/tfvideo_planted.c diff --git a/CHANGELOG.md b/CHANGELOG.md index 9381b4e5820..01ca8503596 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -218,6 +218,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the optimization ledger marks the rows that have open PRs and adds DSpine as row 29 (#2681). Thanks to @jschmied. +- Video on Linux: `libtfvideo.so` links with a version script that exports only its four `tf_video_*` functions, so on x86_64 `zig build video` links `tools/zig/build_ffmpeg.sh`'s FFmpeg with its assembly, where before the link stopped on thousands of `R_X86_64_PC32` relocations against FFmpeg's own tables. Its frames equal PyAV 19.0.1's byte for byte on 23 clips. On aarch64 the library stops exporting FFmpeg's API; real video there is unchecked. macOS is unchanged (#2861). Thanks to @swbiggart. + - GLM-5.3-Flash on Metal with `--vision`: a prompt chunk sends the vision tower only the images that end after its start, so two interleaved image prompts no longer re-encode images their earlier chunks already injected (12 images through the tower instead of 18 in a host test); each image is still encoded on its own and read through its own span, so the rows a chunk reads are unchanged (#2744). - Flash Next on Metal: `TF_FLASHNEXT_SCALE_INDEX=1` (off by default; M1 through M4) has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs: the same values from 10% fewer weight bytes, so replies keep their tokens. On an M3 Ultra with the MLX 4-bit checkpoint, drafted decode on one stream is 1.3% to 3.9% faster (prose 137.3 to 139.1 tok/s, a copy-heavy edit 262.8 to 273.2) and `"draft": false` 0.8% to 2.2% faster, with every token unchanged. The index sits beside the stored scales and biases, which prompt passes still read: 7.3 GiB more in use and about 4 s more load (#2760, part of #2751). diff --git a/tools/zig/check_tfvideo_exports.py b/tools/zig/check_tfvideo_exports.py new file mode 100644 index 00000000000..c1a9d57db62 --- /dev/null +++ b/tools/zig/check_tfvideo_exports.py @@ -0,0 +1,70 @@ +"""An ELF libtfvideo.so exports exactly the four tf_video_* functions, keeps a planted table local, has no TEXTREL.""" +import struct +import sys +from pathlib import Path + +WANT = {'tf_video_open', 'tf_video_frames', 'tf_video_close', 'tf_video_version'} +SHT_SYMTAB, SHT_DYNAMIC, SHT_DYNSYM = 2, 6, 11 +DT_TEXTREL, DT_FLAGS, DF_TEXTREL = 22, 30, 0x4 + + +def sections(data: bytes) -> list[tuple]: + """(type, offset, size, link, entsize) for each section of a little-endian ELF64 file.""" + if data[:4] != b'\x7fELF' or data[4] != 2 or data[5] != 1: + raise SystemExit('not a little-endian ELF64 file') + shoff, = struct.unpack_from(' dict[str, tuple[int, int]]: + """Defined symbols of the given table: name -> (binding, section index).""" + out = {} + for t, off, size, link, ent in secs: + if t != kind: + continue + strtab = secs[link][1] + for at in range(off + ent, off + size, ent): + name, info, _, shndx = struct.unpack_from('> 4, shndx) + return out + + +def textrel(data: bytes, secs: list[tuple]) -> bool: + """True when .dynamic has DT_TEXTREL or DF_TEXTREL in DT_FLAGS.""" + for t, off, size, _, _ in secs: + if t == SHT_DYNAMIC: + for at in range(off, off + size, 16): + tag, val = struct.unpack_from(' int: + failed = 0 + for arg in sys.argv[1:]: + data = Path(arg).read_bytes() + secs = sections(data) + exported = {n for n, (bind, _) in symbols(data, secs, SHT_DYNSYM).items() if bind in (1, 2)} + planted = symbols(data, secs, SHT_SYMTAB).get('tfvideo_planted_table') + problems = [f'exports {sorted(exported ^ WANT)} beyond or short of the four'] if exported != WANT else [] + if planted is None or planted[0] != 0: + problems.append(f'the planted table is not a local symbol: {planted}') + if textrel(data, secs): + problems.append('text relocations') + machine = {62: 'x86_64', 183: 'aarch64'}.get(struct.unpack_from(' + +const uint16_t tfvideo_planted_table[4] = {1, 2, 3, 4}; + +uint32_t tfvideo_planted_read(void) { + uint32_t v; +#if defined(__x86_64__) + __asm__("movzwl tfvideo_planted_table(%%rip), %0" : "=r"(v)); +#else + v = tfvideo_planted_table[0]; +#endif + return v; +} From 1c110c8978ed3847a28ad2acf9b68679077e86c8 Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Wed, 11 Nov 2026 02:03:21 -0800 Subject: [PATCH 136/186] ci: the macOS adapters run check-video-exports-linux beside check-video-layout-linux --- zig/build/host_ci.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index feaa56d56f0..fed6dc007a9 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -251,6 +251,7 @@ const macos_adapters = [_][]const u8{ "check-nemotron-prefill-widths-host", "check-qwen35-cache-host", "check-sessions-linux", + "check-video-exports-linux", "check-video-layout-linux", "compile-disk-identity-metal-host", "compile-dsv41-engram-alloc-aarch64", From bf9b1f1a1f6a5ed02b073d0f1aca311254d856a5 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:10:21 +0000 Subject: [PATCH 137/186] changelog: #2861's line from lines/2861.txt in place of its own --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01ca8503596..31f6983e72c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Video on Linux: `libtfvideo.so` exports only its `tf_video_*` functions, so FFmpeg's x86_64 assembly links into it, and a CI check refuses a build that exports more (#2861). Thanks to @swbiggart. + - Docs: the Intel Arc (XPU) page notes that its B70 prefill row predates 1.1.0's more precise prompt DeltaNet, which costs about 5% of a 32,768-token prompt with tokens unchanged (#2906, measured in #2859). Thanks to @dinadur. - Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. @@ -218,8 +220,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the optimization ledger marks the rows that have open PRs and adds DSpine as row 29 (#2681). Thanks to @jschmied. -- Video on Linux: `libtfvideo.so` links with a version script that exports only its four `tf_video_*` functions, so on x86_64 `zig build video` links `tools/zig/build_ffmpeg.sh`'s FFmpeg with its assembly, where before the link stopped on thousands of `R_X86_64_PC32` relocations against FFmpeg's own tables. Its frames equal PyAV 19.0.1's byte for byte on 23 clips. On aarch64 the library stops exporting FFmpeg's API; real video there is unchecked. macOS is unchanged (#2861). Thanks to @swbiggart. - - GLM-5.3-Flash on Metal with `--vision`: a prompt chunk sends the vision tower only the images that end after its start, so two interleaved image prompts no longer re-encode images their earlier chunks already injected (12 images through the tower instead of 18 in a host test); each image is still encoded on its own and read through its own span, so the rows a chunk reads are unchanged (#2744). - Flash Next on Metal: `TF_FLASHNEXT_SCALE_INDEX=1` (off by default; M1 through M4) has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs: the same values from 10% fewer weight bytes, so replies keep their tokens. On an M3 Ultra with the MLX 4-bit checkpoint, drafted decode on one stream is 1.3% to 3.9% faster (prose 137.3 to 139.1 tok/s, a copy-heavy edit 262.8 to 273.2) and `"draft": false` 0.8% to 2.2% faster, with every token unchanged. The index sits beside the stored scales and biases, which prompt passes still read: 7.3 GiB more in use and about 4 s more load (#2760, part of #2751). From d83ad8e2fea3abf280caf0acdebf0066584707f6 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:10 +0000 Subject: [PATCH 138/186] refactor(core): move the memory counts to engine_status.zig engine_api.zig is at the 600-line limit, and the OCR reader types that follow add to it. Memory and MemorySource describe a backend's status, as HealthSource and Status do, so they move beside them in engine_status.zig and engine_api re-exports them under the same names; no caller changes. --- zig/src/core/engine_api.zig | 10 ++-------- zig/src/core/engine_status.zig | 8 ++++++++ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e2ff1527b18..920af8cd006 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -371,18 +371,12 @@ pub const Opened = struct { follow: ?*const fn (*anyopaque) anyerror!void = null, }; -pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; - /// A backend's words for its own refusals (a request it cannot serve); null: the error's name. pub const Explain = struct { ctx: ?*anyopaque = null, text: *const fn (ctx: ?*anyopaque, err: anyerror) ?[]const u8 }; -/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. -pub const MemorySource = struct { - ctx: ?*anyopaque = null, - read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, -}; - pub const HealthSource = @import("engine_status.zig").HealthSource; +pub const Memory = @import("engine_status.zig").Memory; +pub const MemorySource = @import("engine_status.zig").MemorySource; pub const Status = @import("engine_status.zig").Status; pub const SubmitError = error{ Closed, Busy, InvalidSpans, Stale, InvalidRequest, MemoryPressure }; // Stale: put-back has not reached the device diff --git a/zig/src/core/engine_status.zig b/zig/src/core/engine_status.zig index 3a024ac4243..0700fd90a23 100644 --- a/zig/src/core/engine_status.zig +++ b/zig/src/core/engine_status.zig @@ -12,6 +12,14 @@ pub const FsrSource = struct { read: *const fn (*anyopaque) ?Status.FsrCache, }; +pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; + +/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. +pub const MemorySource = struct { + ctx: ?*anyopaque = null, + read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, +}; + pub const Status = struct { healthy: bool = true, /// Engine-thread publication; all retained states, including long-message interior cuts and spare storage. From 6057b7f89f179349c3c19678fd45efb7956ebdb7 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:11 +0000 Subject: [PATCH 139/186] refactor(server): move the JPEG encoder out of pdf.zig The libjpeg-turbo loader and the BGRA page encoder move to their own module, jpeg_out.zig, with the optional-library loader that PDFium shares, so the OCR route can encode image crops and pdf.zig gains the room its page selection needs later; jpegPixels becomes jpegSize, which returns the width and height a frame header names, for the OCR page planner. Chat PDF input behaves exactly as before, and the new encodeRgb is covered by a round-trip test. --- zig/src/server/jpeg_out.zig | 101 +++++++++++++++++++++++++++++ zig/src/server/pdf.zig | 75 ++++----------------- zig/src/server/pdf_worker.zig | 8 ++- zig/src/server/pdf_worker_test.zig | 4 +- zig/src/server/root.zig | 1 + 5 files changed, 120 insertions(+), 69 deletions(-) create mode 100644 zig/src/server/jpeg_out.zig diff --git a/zig/src/server/jpeg_out.zig b/zig/src/server/jpeg_out.zig new file mode 100644 index 00000000000..0e851bc95b1 --- /dev/null +++ b/zig/src/server/jpeg_out.zig @@ -0,0 +1,101 @@ +//! JPEG output via libjpeg-turbo, loaded on first use: PDF page images (pdf.zig) and OCR image crops (encodeRgb). +const std = @import("std"); +const sys = @import("sys"); +const log = @import("log.zig"); +const Allocator = std.mem.Allocator; + +/// The TurboJPEG 3 entry points the encoders call. +pub const Fns = struct { + tj3Init: *const fn (c_int) callconv(.c) ?*anyopaque, + tj3Destroy: *const fn (?*anyopaque) callconv(.c) void, + tj3Set: *const fn (?*anyopaque, c_int, c_int) callconv(.c) c_int, + tj3Compress8: *const fn (?*anyopaque, [*]const u8, c_int, c_int, c_int, c_int, *?[*]u8, *usize) callconv(.c) c_int, + tj3Free: *const fn (?*anyopaque) callconv(.c) void, +}; + +pub const Error = error{ OutOfMemory, EncodeFailed }; + +const Loaded = union(enum) { unknown, missing, ready: Fns }; + +var state: Loaded = .unknown; +var mutex: std.Io.Mutex = .init; + +/// Looks up every function field of `T` in optional library `id`, optional ones allowed to miss; PDFium uses it too. +pub fn loadDynLib(comptime T: type, comptime id: sys.optional_libraries.Id, a: Allocator) ?T { + const exe = std.process.executableDirPathAlloc(log.io(), a) catch null; + defer if (exe) |p| a.free(p); + var dl = sys.optional_libraries.open(sys.optional_libraries.spec(id), a, sys, exe) orelse return null; + var f: T = undefined; + const info = @typeInfo(T).@"struct"; + inline for (info.field_names, info.field_types) |name, FT| { + if (@typeInfo(FT) == .optional) { + @field(f, name) = dl.lookup(@typeInfo(FT).optional.child, name); + } else { + @field(f, name) = dl.lookup(FT, name) orelse { + dl.close(); + return null; + }; + } + } + return f; +} + +/// libjpeg-turbo's entry points, or null when the library is missing; the first call loads it. +pub fn get(a: Allocator) ?Fns { + const io = log.io(); + mutex.lockUncancelable(io); + defer mutex.unlock(io); + if (state == .unknown) state = if (loadDynLib(Fns, .jpeg, a)) |f| .{ .ready = f } else .missing; + return if (state == .ready) state.ready else null; +} + +/// Encodes PDFium's BGRx bitmap (rows `stride` bytes apart) as a 4:2:0 JPEG at quality 85. +pub fn encodeBgra(a: Allocator, jf: Fns, bgra: []const u8, width: usize, height: usize, stride: usize) Error![]const u8 { + return encode(a, jf, bgra, width, height, stride, 8, 85); // 8: TJPF_BGRA, PDFium's BGRx bitmap +} + +/// Encodes packed RGB rows as a 4:2:0 JPEG at `quality`; NoLibrary when libjpeg-turbo is missing. +pub fn encodeRgb(a: Allocator, rgb: []const u8, width: usize, height: usize, quality: u8) (Error || error{NoLibrary})![]const u8 { + const jf = get(a) orelse return error.NoLibrary; + const pitch = std.math.mul(usize, width, 3) catch return error.EncodeFailed; + if (pitch == 0 or rgb.len / pitch < height) return error.EncodeFailed; + return encode(a, jf, rgb, width, height, pitch, 0, quality); // 0: TJPF_RGB +} + +fn encode(a: Allocator, jf: Fns, pixels: []const u8, width: usize, height: usize, pitch: usize, format: c_int, quality: u8) Error![]const u8 { + const max: usize = std.math.maxInt(c_int); + if (width == 0 or height == 0 or width > max or height > max or pitch > max) return error.EncodeFailed; + const handle = jf.tj3Init(0) orelse return error.EncodeFailed; // 0 = TJINIT_COMPRESS + defer jf.tj3Destroy(handle); + _ = jf.tj3Set(handle, 3, quality); // TJPARAM_QUALITY = 3 + _ = jf.tj3Set(handle, 4, 2); // TJPARAM_SUBSAMP = 4, TJSAMP_420 = 2 + var out_buf: ?[*]u8 = null; + var out_size: usize = 0; + const rc = jf.tj3Compress8(handle, pixels.ptr, @intCast(width), @intCast(pitch), @intCast(height), format, &out_buf, &out_size); + defer if (out_buf) |p| jf.tj3Free(p); + if (rc != 0 or out_buf == null or out_size == 0) return error.EncodeFailed; + return a.dupe(u8, out_buf.?[0..out_size]); +} + +test "jpeg out: encodeRgb round-trips a small image through core.picture" { + const picture = @import("core").picture; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const w = 24; + const h = 10; + var rgb: [w * h * 3]u8 = undefined; + for (&rgb, 0..) |*px, i| px.* = @truncate(i * 7); + const bytes = encodeRgb(a, &rgb, w, h, 90) catch |err| switch (err) { + error.NoLibrary => return error.SkipZigTest, + else => return err, + }; + try std.testing.expectEqual(picture.Format.jpeg, picture.sniff(bytes)); + const img = picture.decode(std.testing.io, a, bytes) catch |err| switch (err) { + error.NoJpegLibrary, error.JpegLibraryVersion => return error.SkipZigTest, + else => return err, + }; + try std.testing.expectEqual(@as(usize, w), img.width); + try std.testing.expectEqual(@as(usize, h), img.height); + try std.testing.expectError(error.EncodeFailed, encodeRgb(a, rgb[0 .. rgb.len - 1], w, h, 90)); +} diff --git a/zig/src/server/pdf.zig b/zig/src/server/pdf.zig index a134564a3ee..837bbc02661 100644 --- a/zig/src/server/pdf.zig +++ b/zig/src/server/pdf.zig @@ -4,6 +4,7 @@ const sys = @import("sys"); const json = @import("json"); const errors = @import("errors.zig"); const images = @import("images.zig"); +const jpeg_out = @import("jpeg_out.zig"); const log = @import("log.zig"); const worker = @import("pdf_worker.zig"); const Value = json.Value; @@ -40,65 +41,23 @@ pub const PdfiumFns = struct { FPDF_RenderPage_Close: ?*const fn (?*anyopaque) callconv(.c) void = null, }; -pub const JpegFns = struct { - tj3Init: *const fn (c_int) callconv(.c) ?*anyopaque, - tj3Destroy: *const fn (?*anyopaque) callconv(.c) void, - tj3Set: *const fn (?*anyopaque, c_int, c_int) callconv(.c) c_int, - tj3Compress8: *const fn (?*anyopaque, [*]const u8, c_int, c_int, c_int, c_int, *?[*]u8, *usize) callconv(.c) c_int, - tj3Free: *const fn (?*anyopaque) callconv(.c) void, -}; - const DlInfo = extern struct { fname: ?[*:0]const u8, fbase: ?*anyopaque, sname: ?[*:0]const u8, saddr: ?*anyopaque }; extern "c" fn dladdr(addr: *const anyopaque, info: *DlInfo) c_int; const PdfiumLoaded = union(enum) { unknown, missing, ready: PdfiumFns }; -const JpegLoaded = union(enum) { unknown, missing, ready: JpegFns }; var pdfium_state: PdfiumLoaded = .unknown; -var jpeg_state: JpegLoaded = .unknown; var pdfium_mutex: std.Io.Mutex = .init; -var jpeg_mutex: std.Io.Mutex = .init; - -pub fn loadDynLib(comptime T: type, comptime id: sys.optional_libraries.Id, a: Allocator) ?T { - const exe = std.process.executableDirPathAlloc(log.io(), a) catch null; - defer if (exe) |p| a.free(p); - var dl = sys.optional_libraries.open(sys.optional_libraries.spec(id), a, sys, exe) orelse return null; - var f: T = undefined; - const info = @typeInfo(T).@"struct"; - inline for (info.field_names, info.field_types) |name, FT| { - if (@typeInfo(FT) == .optional) { - @field(f, name) = dl.lookup(@typeInfo(FT).optional.child, name); - } else { - @field(f, name) = dl.lookup(FT, name) orelse { - dl.close(); - return null; - }; - } - } - return f; -} fn findPdfium(a: Allocator) PdfiumLoaded { - const f = loadDynLib(PdfiumFns, .pdfium, a) orelse return .missing; + const f = jpeg_out.loadDynLib(PdfiumFns, .pdfium, a) orelse return .missing; if (sys.windows) f.FPDF_InitLibrary(); return .{ .ready = f }; } -fn findJpeg(a: Allocator) JpegLoaded { - return if (loadDynLib(JpegFns, .jpeg, a)) |f| .{ .ready = f } else .missing; -} - -fn getJpeg(a: Allocator) ?JpegFns { - const io = log.io(); - jpeg_mutex.lockUncancelable(io); - defer jpeg_mutex.unlock(io); - if (jpeg_state == .unknown) jpeg_state = findJpeg(a); - return if (jpeg_state == .ready) jpeg_state.ready else null; -} - /// Logs the loaded PDFium path and SHA-256 at startup when --vision-pdf is enabled, or warns if a library is missing. pub fn logStartup(a: Allocator, io: std.Io) void { - if (getJpeg(a) == null) log.line("--vision-pdf: {s}", .{sys.optional_libraries.spec(.jpeg).message}); + if (jpeg_out.get(a) == null) log.line("--vision-pdf: {s}", .{sys.optional_libraries.spec(.jpeg).message}); pdfium_mutex.lockUncancelable(io); defer pdfium_mutex.unlock(io); if (pdfium_state == .unknown) pdfium_state = findPdfium(a); @@ -207,7 +166,7 @@ pub fn decodeFileData(cx: *Cx, raw: []const u8) errors.Refused![]const u8 { /// Renders all pages of `pdf_bytes` within the request-wide `budget` into `image_url` parts appended to `parts`. pub fn expandBytes(cx: *Cx, pdf_bytes: []const u8, image_type: []const u8, max_images: u32, max_pages: u32, budget: *Budget, parts: *std.ArrayList(Value)) errors.Refused!void { - const jf = getJpeg(cx.a) orelse return cx.fail(.unimplemented, "{s}", .{sys.optional_libraries.spec(.jpeg).message}); + const jf = jpeg_out.get(cx.a) orelse return cx.fail(.unimplemented, "{s}", .{sys.optional_libraries.spec(.jpeg).message}); const io = log.io(); pdfium_mutex.lockUncancelable(io); defer pdfium_mutex.unlock(io); @@ -255,7 +214,7 @@ fn renderPageWithDeadline(cx: *Cx, f: PdfiumFns, bmp: ?*anyopaque, page: ?*anyop if (std.Io.Clock.awake.now(log.io()).toNanoseconds() >= deadline_ns) return cx.refuse("PDF rendering timed out"); } -pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: JpegFns, pdf_bytes: []const u8, max_images: u32, max_pages: u32, deadline_ns: i96, budget: *Budget, out_jpegs: *std.ArrayList([]const u8)) errors.Refused!void { +pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: jpeg_out.Fns, pdf_bytes: []const u8, max_images: u32, max_pages: u32, deadline_ns: i96, budget: *Budget, out_jpegs: *std.ArrayList([]const u8)) errors.Refused!void { if (pdf_bytes.len > std.math.maxInt(c_int)) return cx.refuse("PDF exceeds the encoded byte limit"); if (std.Io.Clock.awake.now(log.io()).toNanoseconds() >= deadline_ns) return cx.refuse("PDF rendering timed out"); const doc = f.FPDF_LoadMemDocument(pdf_bytes.ptr, @intCast(pdf_bytes.len), null) orelse @@ -287,7 +246,7 @@ pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: JpegFns, pdf_bytes: []const u8 const total_bytes = @as(usize, @intCast(stride)) * @as(usize, @intCast(dims[1])); @memset(bgra[0..total_bytes], 0xFF); try renderPageWithDeadline(cx, f, bmp, page, dims[0], dims[1], idx, deadline_ns); - const jpeg_bytes = encodeJpegBgra(cx.a, jf, bgra[0..total_bytes], @intCast(dims[0]), @intCast(dims[1]), @intCast(stride)) catch |err| switch (err) { + const jpeg_bytes = jpeg_out.encodeBgra(cx.a, jf, bgra[0..total_bytes], @intCast(dims[0]), @intCast(dims[1]), @intCast(stride)) catch |err| switch (err) { error.OutOfMemory => return error.OutOfMemory, error.EncodeFailed => return cx.refuse("rendered PDF page could not be encoded as JPEG"), }; @@ -315,7 +274,8 @@ fn pageDimensions(pw_raw: f32, ph_raw: f32) [2]c_int { return .{ w, h }; } -pub fn jpegPixels(bytes: []const u8) ?u64 { +/// The width and height a baseline or progressive JPEG's frame header names, or null when it has none. +pub fn jpegSize(bytes: []const u8) ?[2]u32 { if (bytes.len < 4 or bytes[0] != 0xFF or bytes[1] != 0xD8) return null; var i: usize = 2; while (i + 4 <= bytes.len) { @@ -330,28 +290,15 @@ pub fn jpegPixels(bytes: []const u8) ?u64 { if (seg_len < 2 or i + seg_len > bytes.len) return null; if (marker == 0xC0 or marker == 0xC1 or marker == 0xC2) { if (seg_len < 7) return null; - const h: u64 = std.mem.readInt(u16, bytes[i + 3 ..][0..2], .big); - const w: u64 = std.mem.readInt(u16, bytes[i + 5 ..][0..2], .big); - return if (w > 0 and h > 0) w * h else null; + const h: u32 = std.mem.readInt(u16, bytes[i + 3 ..][0..2], .big); + const w: u32 = std.mem.readInt(u16, bytes[i + 5 ..][0..2], .big); + return if (w > 0 and h > 0) .{ w, h } else null; } i += seg_len; } return null; } -fn encodeJpegBgra(a: Allocator, jf: JpegFns, bgra: []const u8, width: usize, height: usize, stride: usize) error{ OutOfMemory, EncodeFailed }![]const u8 { - const handle = jf.tj3Init(0) orelse return error.EncodeFailed; // 0 = TJINIT_COMPRESS - defer jf.tj3Destroy(handle); - _ = jf.tj3Set(handle, 3, 85); // TJPARAM_QUALITY = 3 - _ = jf.tj3Set(handle, 4, 2); // TJPARAM_SUBSAMP = 4, TJSAMP_420 = 2 - var out_buf: ?[*]u8 = null; - var out_size: usize = 0; - const rc = jf.tj3Compress8(handle, bgra.ptr, @intCast(width), @intCast(stride), @intCast(height), 8, &out_buf, &out_size); // 8: TJPF_BGRA, PDFium's BGRx bitmap - defer if (out_buf) |p| jf.tj3Free(p); - if (rc != 0 or out_buf == null or out_size == 0) return error.EncodeFailed; - return a.dupe(u8, out_buf.?[0..out_size]); -} - pub fn makeDataUrl(a: Allocator, mime: []const u8, bytes: []const u8) Allocator.Error![]const u8 { const prefix = try std.fmt.allocPrint(a, "data:{s};base64,", .{mime}); const enc = std.base64.standard.Encoder; diff --git a/zig/src/server/pdf_worker.zig b/zig/src/server/pdf_worker.zig index 248cd387a9a..099b751a447 100644 --- a/zig/src/server/pdf_worker.zig +++ b/zig/src/server/pdf_worker.zig @@ -6,6 +6,7 @@ const errors = @import("errors.zig"); const images = @import("images.zig"); const log = @import("log.zig"); const pdf = @import("pdf.zig"); +const jpeg_out = @import("jpeg_out.zig"); const Cx = errors.Cx; const Budget = pdf.Budget; const Allocator = std.mem.Allocator; @@ -117,7 +118,8 @@ pub fn parseChildReply(cx: *Cx, frame: []const u8, max_images: u32, max_pages: u if (len > images.max_image_bytes) return cx.refuse("rendered PDF page exceeds the 10 MB image byte limit"); if (budget.encoded_bytes + len > images.max_request_bytes) return cx.refuse("the request's images exceed 20 MB encoded"); const jb = frame[pos .. pos + len]; - const px = pdf.jpegPixels(jb) orelse return cx.refuse("rendered PDF page could not be encoded as JPEG"); + const size = pdf.jpegSize(jb) orelse return cx.refuse("rendered PDF page could not be encoded as JPEG"); + const px = @as(u64, size[0]) * size[1]; if (budget.pixels_used + px > pdf.max_request_pixels) return cx.refuse("the request's images exceed the 32-megapixel limit"); budget.pdf_pages_used += 1; budget.images_used += 1; @@ -173,8 +175,8 @@ fn work(cx: *Cx, jpegs: *std.ArrayList([]const u8)) errors.Refused!void { if (!readAll(0, pdf_bytes)) return cx.refuse(failed); applyChildLimits(); const libs = sys.optional_libraries; - const f = pdf.loadDynLib(pdf.PdfiumFns, .pdfium, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.pdfium).message}); - const jf = pdf.loadDynLib(pdf.JpegFns, .jpeg, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.jpeg).message}); + const f = jpeg_out.loadDynLib(pdf.PdfiumFns, .pdfium, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.pdfium).message}); + const jf = jpeg_out.loadDynLib(jpeg_out.Fns, .jpeg, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.jpeg).message}); f.FPDF_InitLibrary(); var budget: Budget = .{ .images_used = req.images_used, .pdf_pages_used = req.pdf_pages_used, .encoded_bytes = @intCast(req.encoded_bytes), .pixels_used = req.pixels_used }; return pdf.renderInProcess(cx, f, jf, pdf_bytes, req.max_images, req.max_pages, deadline_ns, &budget, jpegs); diff --git a/zig/src/server/pdf_worker_test.zig b/zig/src/server/pdf_worker_test.zig index 28a56dd3722..446f093a483 100644 --- a/zig/src/server/pdf_worker_test.zig +++ b/zig/src/server/pdf_worker_test.zig @@ -37,7 +37,7 @@ test "tensorfold-native PDF render worker: renders a text PDF in its spawned wor }; try std.testing.expectEqual(@as(usize, 1), jpegs.items.len); try std.testing.expectEqual(pdf.Budget{ .images_used = 1, .pdf_pages_used = 1, .encoded_bytes = jpegs.items[0].len, .pixels_used = 400 * 400 }, budget); - try std.testing.expectEqual(@as(?u64, 400 * 400), pdf.jpegPixels(jpegs.items[0])); + try std.testing.expectEqual([2]u32{ 400, 400 }, pdf.jpegSize(jpegs.items[0]).?); } test "tensorfold-native PDF render worker: libpdfium renders in the spawned worker and enforces request-wide page, byte and pixel caps" { @@ -59,7 +59,7 @@ test "tensorfold-native PDF render worker: libpdfium renders in the spawned work for (parts.items) |p| { const img_bytes = try images.dataBytes(&cx, p.get("image_url").?.get("url").?.string); try std.testing.expectEqual(images.Format.jpeg, images.sniff(img_bytes)); - try std.testing.expectEqual(@as(?u64, 400 * 400), pdf.jpegPixels(img_bytes)); + try std.testing.expectEqual([2]u32{ 400, 400 }, pdf.jpegSize(img_bytes).?); } var second_cx: Cx = .{ .a = a }; try std.testing.expectError(error.Refused, pdf.expandBytes(&second_cx, pdf.minimal_two_page_pdf, "image_url", 4, 3, &budget, &parts)); diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 66d0da36073..d3e4954a94f 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -112,6 +112,7 @@ test { _ = @import("messages.zig"); _ = @import("pdf.zig"); _ = @import("pdf_flags_test.zig"); + _ = @import("jpeg_out.zig"); _ = @import("images.zig"); _ = @import("images_video_marker_test.zig"); _ = @import("image_key.zig"); From 2a0dd7947302e7148a92c6076629df64b7657fa0 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:11 +0000 Subject: [PATCH 140/186] feat(core): add OCR reader types and the image frame size An engine now declares an optional OCR reader in Info.ocr: the prompt that follows the page image, the geometry it gives and a parse from generated text to Markdown blocks with quads in original pixels. PreparedImage gains the tower-input frame so readers can map boxes back, the HIP offer carries the reader through, and a bounded reader for the HTML table subset dots.mocr writes turns tables into GFM pipes, clean spanned HTML or text rows, escaping literal markup so Markdown never reinterprets recognized text. --- zig/src/core/engine_api.zig | 3 + zig/src/core/engine_api_test.zig | 1 + zig/src/core/engine_ocr.zig | 54 ++++ zig/src/core/engine_vision.zig | 5 + zig/src/core/ocr_html.zig | 437 +++++++++++++++++++++++++++++++ zig/src/core/ocr_html_test.zig | 269 +++++++++++++++++++ zig/src/hip_serve/offer.zig | 3 + zig/src/native/hip.zig | 1 + 8 files changed, 773 insertions(+) create mode 100644 zig/src/core/engine_ocr.zig create mode 100644 zig/src/core/ocr_html.zig create mode 100644 zig/src/core/ocr_html_test.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 920af8cd006..7bb5f6fafa4 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -116,6 +116,8 @@ pub const FrameSink = vision_types.FrameSink; pub const PreparedVideo = vision_types.PreparedVideo; pub const VideoOffer = vision_types.VideoOffer; pub const VisionOpen = vision_types.VisionOpen; +pub const Size = vision_types.Size; +pub const ocr = @import("engine_ocr.zig"); pub const SpeechRequest = @import("speech_api.zig").SpeechRequest; pub const SpeechResult = @import("speech_api.zig").SpeechResult; @@ -223,6 +225,7 @@ pub const Info = struct { pictures: bool = false, /// Image input (--vision); null: text only. vision: ?Vision = null, + ocr: ?ocr.Ocr = null, // OCR reader for /v1/ocr and /ocr; null: no OCR routes /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, /// The learner reads conversation prompts at any length through taps. diff --git a/zig/src/core/engine_api_test.zig b/zig/src/core/engine_api_test.zig index 74a4aadc51d..efadf7dc452 100644 --- a/zig/src/core/engine_api_test.zig +++ b/zig/src/core/engine_api_test.zig @@ -15,4 +15,5 @@ test { _ = @import("keepalive.zig"); _ = @import("originals_test.zig"); _ = @import("image_api_test.zig"); + _ = @import("engine_ocr.zig"); } diff --git a/zig/src/core/engine_ocr.zig b/zig/src/core/engine_ocr.zig new file mode 100644 index 00000000000..07d1c507a57 --- /dev/null +++ b/zig/src/core/engine_ocr.zig @@ -0,0 +1,54 @@ +//! An engine's OCR reader (``Info.ocr``): its prompt and the parse from generated text to Markdown blocks. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Size = @import("engine_vision.zig").Size; + +/// The HTML table subset readers and the server share (ocr_html.zig). +pub const html = @import("ocr_html.zig"); + +/// TL, TR, BR, BL of the text in its upright reading frame, in original-image pixels (rotated text keeps its order). +pub const Quad = [4][2]f32; + +/// What a block holds; ``page`` only for a ``.none`` reader (the whole page's Markdown in one block). +pub const Kind = enum { + text, + title, + section_header, + list, + table, + formula, + code, + caption, + footnote, + page_header, + page_footer, + picture, + page, +}; + +/// The geometry a reader gives: none, one quad a block, or one quad a line. +pub const Geometry = enum { none, block, line }; + +/// One block of a page in reading order. +pub const OcrBlock = struct { + content: []const u8, // normalized Markdown (docs/api.md OCR: Reader Markdown); empty only for .picture + quad: ?Quad, // null only for a reader whose geometry is .none + kind: Kind, +}; + +/// The generated text does not follow the model's grammar. +pub const Error = error{Unreadable}; + +/// OCR an engine serves: the prompt after the page image, and the reader of the generated text. +pub const Ocr = struct { + /// Text parts that follow the image in the single user turn (may be empty). + prompt: []const u8, + /// What geometry `parse` gives. + geometry: Geometry, + /// Text to blocks in pixels of ``original``; ``frame`` is the tower input size the prepare step used. + parse: *const fn (a: Allocator, text: []const u8, original: Size, frame: Size) anyerror![]OcrBlock, +}; + +test { + _ = @import("ocr_html_test.zig"); +} diff --git a/zig/src/core/engine_vision.zig b/zig/src/core/engine_vision.zig index a4d4e819da5..285b697bbe0 100644 --- a/zig/src/core/engine_vision.zig +++ b/zig/src/core/engine_vision.zig @@ -13,8 +13,13 @@ pub const PreparedImage = struct { hash: u64, /// The image as tower inputs in order, each image tokens between `before` and `after`; empty: `pixels` alone. pieces: []const Piece = &.{}, + /// Pixel size of the tower input the image was resized to; set by every family with an OCR reader. + frame: ?Size = null, }; +/// An image's pixel size. +pub const Size = struct { width: u32, height: u32 }; + /// One tower input of an image, and the tokens the placeholder expands to around its rows. pub const Piece = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, before: []const u32 = &.{}, after: []const u32 = &.{} }; diff --git a/zig/src/core/ocr_html.zig b/zig/src/core/ocr_html.zig new file mode 100644 index 00000000000..83bacb8397d --- /dev/null +++ b/zig/src/core/ocr_html.zig @@ -0,0 +1,437 @@ +//! The HTML table subset OCR models write: cells and spans, and the table as GFM pipes, clean HTML or text rows. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// Most HTML bytes, grid rows and grid columns ``parse`` reads; a table past them is ``error.Malformed``. +pub const max_bytes = 64 * 1024; +pub const max_rows = 512; +pub const max_cols = 64; + +/// One cell at its top-left grid position, covering ``rows`` x ``cols``; ``text`` holds '\n' where the HTML had
. +pub const Cell = struct { text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool }; + +/// A table's grid and its cells in row-major order; ``spanned``: some cell covers more than one position. +pub const Table = struct { + rows: u32, + cols: u32, + cells: []Cell, + spanned: bool, + + /// Frees what ``parse`` allocated in ``a``. + pub fn deinit(t: Table, a: Allocator) void { + for (t.cells) |c| a.free(c.text); + a.free(t.cells); + } +}; + +pub const ParseError = error{ Malformed, OutOfMemory }; + +/// The one in ``html``, with only whitespace and comments around it; cells and texts allocated in ``a``. +pub fn parse(a: Allocator, html: []const u8) ParseError!Table { + if (html.len > max_bytes) return error.Malformed; + var p: Parser = .{ .a = a, .src = html }; + defer p.text.deinit(a); + errdefer { + for (p.cells.items) |c| a.free(c.text); + p.cells.deinit(a); + } + try p.run(); + var cols: u32 = 0; + var spanned = false; + for (p.cells.items) |*c| { + c.rows = @min(c.rows, p.rows - c.row); // a rowspan past the last row ends there, as in HTML + cols = @max(cols, c.col + c.cols); + spanned = spanned or c.rows > 1 or c.cols > 1; + } + if (cols == 0) return error.Malformed; + return .{ .rows = p.rows, .cols = cols, .cells = try p.cells.toOwnedSlice(a), .spanned = spanned }; +} + +/// The table as a GFM pipe table: first row the header, cell text escaped to read literally (``cellText``); spans lost. +pub fn pipe(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + var k: usize = 0; + for (0..t.rows) |r| { + for (0..t.cols) |c| { + try out.appendSlice(a, if (c == 0) "| " else " | "); + try cellText(a, &out, if (topLeft(t, &k, r, c)) |cell| cell.text else ""); + } + try out.appendSlice(a, " |"); + if (r == 0) { + try out.append(a, '\n'); + for (0..t.cols) |_| try out.appendSlice(a, "| --- "); + try out.append(a, '|'); + } + if (r + 1 < t.rows) try out.append(a, '\n'); + } + return out.toOwnedSlice(a); +} + +/// The table as one clean
: rows of
/, rowspan and colspan only above 1, text entity-escaped. +pub fn normalized(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + try out.appendSlice(a, ""); + var k: usize = 0; + for (0..t.rows) |r| { + try out.appendSlice(a, ""); + while (k < t.cells.len and t.cells[k].row == r) : (k += 1) { + const cell = t.cells[k]; + const name = if (cell.head) "th" else "td"; + try out.print(a, "<{s}", .{name}); + if (cell.rows > 1) try out.print(a, " rowspan=\"{d}\"", .{cell.rows}); + if (cell.cols > 1) try out.print(a, " colspan=\"{d}\"", .{cell.cols}); + try out.append(a, '>'); + for (cell.text) |ch| switch (ch) { + '&' => try out.appendSlice(a, "&"), + '<' => try out.appendSlice(a, "<"), + '>' => try out.appendSlice(a, ">"), + '"' => try out.appendSlice(a, """), + '\n' => try out.appendSlice(a, "
"), + else => try out.append(a, ch), + }; + try out.print(a, "", .{name}); + } + try out.appendSlice(a, ""); + } + try out.appendSlice(a, "
"); + return out.toOwnedSlice(a); +} + +/// One line a grid row, its cells joined by '\t'; a spanned cell's text at its top-left only, trailing cells kept. +pub fn rowsText(a: Allocator, t: Table) Allocator.Error![][]const u8 { + const lines = try a.alloc([]const u8, t.rows); + var done: usize = 0; + errdefer { + for (lines[0..done]) |l| a.free(l); + a.free(lines); + } + var line: std.ArrayList(u8) = .empty; + defer line.deinit(a); + var k: usize = 0; + for (lines, 0..) |*l, r| { + line.clearRetainingCapacity(); + for (0..t.cols) |c| { + if (c > 0) try line.append(a, '\t'); + const text = if (topLeft(t, &k, r, c)) |cell| cell.text else ""; + for (text) |ch| try line.append(a, if (ch == '\n') ' ' else ch); + } + l.* = try a.dupe(u8, line.items); + done += 1; + } + return lines; +} + +/// A cell's text in a pipe row: math as written (only `|` escaped), markup escaped elsewhere, to read back literally. +fn cellText(a: Allocator, out: *std.ArrayList(u8), text: []const u8) Allocator.Error!void { + var no_close = false; // a `$` found no closer, so no later `$` in the cell can (markdown.zig's rule) + var i: usize = 0; + while (i < text.len) { + const ch = text[i]; + if (ch == '$') { + if (mathEnd(text, i, &no_close)) |end| { + try mathText(a, out, text[i..end]); + i = end; + continue; + } + // a `$` that opens no math is escaped, so a reader never pairs it with another + const lone: usize = if (i + 1 < text.len and text[i + 1] == '$') 2 else 1; + for (0..lone) |_| try out.appendSlice(a, "\\$"); + i += lone; + continue; + } + switch (ch) { + '\n' => try out.append(a, ' '), + '|', '`', '*', '[', ']' => try out.appendSlice(a, &[_]u8{ '\\', ch }), + // CommonMark reads a backslash as an escape only before ASCII punctuation; elsewhere it is literal + '\\' => try out.appendSlice(a, if (i + 1 < text.len and isPunct(text[i + 1])) "\\\\" else "\\"), + // `_` between two word characters can neither open nor close emphasis, so snake_case stays readable + '_' => try out.appendSlice(a, if (i > 0 and i + 1 < text.len and isWord(text[i - 1]) and isWord(text[i + 1])) "_" else "\\_"), + '<' => try out.appendSlice(a, if (opensTag(text[i + 1 ..])) "<" else "<"), + '&' => try out.appendSlice(a, if (opensEntity(text[i + 1 ..])) "&" else "&"), + else => try out.append(a, ch), + } + i += 1; + } +} + +/// Math as written, each `|` as `\|`: GFM's table step removes exactly that backslash, whatever precedes it. +fn mathText(a: Allocator, out: *std.ArrayList(u8), m: []const u8) Allocator.Error!void { + for (m) |ch| switch (ch) { + '|' => try out.appendSlice(a, "\\|"), + '\n' => try out.append(a, ' '), + else => try out.append(a, ch), + }; +} + +/// The end of the `$$…$$` or `$…$` math opening at ``s[i]``, by markdown.zig's ``mathSpan`` rule; null when none. +fn mathEnd(s: []const u8, i: usize, no_close: *bool) ?usize { + if (i + 1 < s.len and s[i + 1] == '$') return if (std.mem.indexOfPos(u8, s, i + 2, "$$")) |j| j + 2 else null; + if (i + 1 >= s.len or isSpace(s[i + 1]) or no_close.*) return null; + var j = i + 2; + while (j < s.len) : (j += 1) { + if (s[j] != '$' or isSpace(s[j - 1]) or escaped(s, j)) continue; + if (j + 1 < s.len and std.ascii.isDigit(s[j + 1])) continue; + return j + 1; + } + no_close.* = true; + return null; +} + +/// Whether ``s[i]`` follows an odd run of backslashes. +fn escaped(s: []const u8, i: usize) bool { + var n: usize = 0; + while (n < i and s[i - 1 - n] == '\\') n += 1; + return n % 2 == 1; +} + +fn isSpace(c: u8) bool { + return c == ' ' or c == '\t' or c == '\n' or c == '\r'; +} + +fn isPunct(c: u8) bool { + return c < 0x80 and std.ascii.isPrint(c) and !std.ascii.isAlphanumeric(c) and c != ' '; +} + +/// A byte of a word for emphasis rules: neither whitespace nor ASCII punctuation (UTF-8 letters count). +fn isWord(c: u8) bool { + return !isSpace(c) and !isPunct(c); +} + +/// Whether text after a literal `<` would start raw HTML in Markdown: a tag, an end tag, a comment or a declaration. +fn opensTag(rest: []const u8) bool { + return rest.len > 0 and (std.ascii.isAlphabetic(rest[0]) or std.mem.indexOfScalar(u8, "/!?", rest[0]) != null); +} + +/// Whether text after a literal `&` reads as an entity reference in Markdown (`name;`, `#digits;` or `#xhex;`). +fn opensEntity(rest: []const u8) bool { + var n: usize = 0; + while (n < rest.len and n < 33 and (std.ascii.isAlphanumeric(rest[n]) or (n == 0 and rest[n] == '#'))) n += 1; + return n > 0 and n < rest.len and rest[n] == ';'; +} + +/// The cell whose top-left is (r, c), walking ``k`` through the row-major cells; null: covered or empty. +fn topLeft(t: Table, k: *usize, r: usize, c: usize) ?Cell { + if (k.* >= t.cells.len or t.cells[k.*].row != r or t.cells[k.*].col != c) return null; + k.* += 1; + return t.cells[k.* - 1]; +} + +const Tag = enum { table, tr, td, th, br, other }; + +fn tagKind(name: []const u8) Tag { + inline for (.{ Tag.table, Tag.tr, Tag.td, Tag.th, Tag.br }) |t| { + if (std.ascii.eqlIgnoreCase(name, @tagName(t))) return t; + } + return .other; +} + +/// A span attribute's value as HTML reads it: leading digits, and 1 when there are none or they read 0. +fn spanValue(v: []const u8) u32 { + var i: usize = 0; + while (i < v.len and std.ascii.isWhitespace(v[i])) i += 1; + var n: u32 = 0; + while (i < v.len and std.ascii.isDigit(v[i])) : (i += 1) n = @min(n * 10 + (v[i] - '0'), 1_000_000); + return @max(n, 1); +} + +/// The code point a named or numeric entity body (between '&' and ';') stands for; null: not one this reader knows. +fn entityPoint(name: []const u8) ?u21 { + const named = .{ .{ "amp", '&' }, .{ "lt", '<' }, .{ "gt", '>' }, .{ "quot", '"' }, .{ "apos", '\'' } }; + if (std.mem.eql(u8, name, "nbsp")) return 0xa0; + inline for (named) |e| if (std.mem.eql(u8, name, e[0])) return e[1]; + if (name.len < 2 or name[0] != '#') return null; + const hex = name[1] == 'x' or name[1] == 'X'; + const base: u8 = if (hex) 16 else 10; + const digits = name[if (hex) 2 else 1..]; + if (digits.len == 0 or digits.len > 7) return null; + var v: u32 = 0; + for (digits) |d| v = v * base + (std.fmt.charToDigit(d, base) catch return null); + return if (v == 0 or v > 0x10ffff or (v >= 0xd800 and v <= 0xdfff)) 0xfffd else @intCast(v); +} + +/// Whitespace or one of ``ends``: where an attribute's name or unquoted value stops. +fn stops(ch: u8, ends: []const u8) bool { + return std.ascii.isWhitespace(ch) or std.mem.indexOfScalar(u8, ends, ch) != null; +} + +const Parser = struct { + a: Allocator, + src: []const u8, + i: usize = 0, + where: enum { before, inside, after } = .before, + cells: std.ArrayList(Cell) = .empty, + cell: ?Cell = null, // the open cell; its text gathers in ``text`` until it closes + text: std.ArrayList(u8) = .empty, + space: bool = false, // whitespace met in the open cell since its last character + row_open: bool = false, + rows: u32 = 0, // rows opened so far; the open row is the last + col: u32 = 0, // the open row's next column to try + until: [max_cols]u32 = @splat(0), // each column is covered by earlier cells through row until - 1 + + fn run(p: *Parser) ParseError!void { + while (p.i < p.src.len) { + const ch = p.src[p.i]; + const next = if (p.i + 1 < p.src.len) p.src[p.i + 1] else 0; + const after = if (p.i + 2 < p.src.len) p.src[p.i + 2] else 0; + if (ch == '<' and (next == '!' or next == '?')) { + try p.skipMarkup(); + } else if (ch == '<' and std.ascii.isAlphabetic(if (next == '/') after else next)) { + try p.tag(); + } else if (ch == '&') { + try p.entity(); + } else { + try p.char(ch); + p.i += 1; + } + } + if (p.where != .after) return error.Malformed; + } + + /// A comment, a doctype or a processing instruction: skipped whole. + fn skipMarkup(p: *Parser) ParseError!void { + if (std.mem.startsWith(u8, p.src[p.i..], "") orelse return error.Malformed; + p.i = end + 3; + } else { + const end = std.mem.indexOfScalarPos(u8, p.src, p.i, '>') orelse return error.Malformed; + p.i = end + 1; + } + } + + fn entity(p: *Parser) ParseError!void { + const rest = p.src[p.i..]; + const semi = std.mem.indexOfScalar(u8, rest[0..@min(rest.len, 12)], ';'); + const point = if (semi) |s| entityPoint(rest[1..s]) else null; + const cp = point orelse { + p.i += 1; + return p.char('&'); + }; + p.i += semi.? + 1; + if (cp == 0xa0) return p.char(' '); + var buf: [4]u8 = undefined; + const n = std.unicode.utf8Encode(cp, &buf) catch unreachable; + for (buf[0..n]) |b| try p.char(b); + } + + /// One text byte: whitespace runs fold to one space between words; outside a cell only whitespace may stand. + fn char(p: *Parser, ch: u8) ParseError!void { + const blank = std.ascii.isWhitespace(ch); + if (p.cell == null) { + if (!blank) return error.Malformed; + return; + } + if (blank) { + p.space = true; + return; + } + if (p.space and p.text.items.len > 0 and p.text.getLast() != '\n') try p.text.append(p.a, ' '); + p.space = false; + try p.text.append(p.a, ch); + } + + /// A start or end tag at ``i``: its name, its span attributes in any quoting, then its effect. + fn tag(p: *Parser) ParseError!void { + const src = p.src; + var i = p.i + 1; + const close = src[i] == '/'; + if (close) i += 1; + const start = i; + while (i < src.len and (std.ascii.isAlphanumeric(src[i]) or std.mem.indexOfScalar(u8, "-_:.", src[i]) != null)) i += 1; + const kind = tagKind(src[start..i]); + var span = [2]u32{ 1, 1 }; // rowspan, colspan + while (true) { + while (i < src.len and (std.ascii.isWhitespace(src[i]) or src[i] == '/')) i += 1; + if (i >= src.len) return error.Malformed; + if (src[i] == '>') break; + const name_start = i; + while (i < src.len and !stops(src[i], "=>/")) i += 1; + const name = src[name_start..i]; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len or src[i] != '=') continue; + i += 1; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len) return error.Malformed; + var value: []const u8 = undefined; + if (src[i] == '"' or src[i] == '\'') { + const end = std.mem.indexOfScalarPos(u8, src, i + 1, src[i]) orelse return error.Malformed; + value = src[i + 1 .. end]; + i = end + 1; + } else { + const value_start = i; + while (i < src.len and !stops(src[i], ">")) i += 1; + value = src[value_start..i]; + } + if (std.ascii.eqlIgnoreCase(name, "rowspan")) span[0] = spanValue(value); + if (std.ascii.eqlIgnoreCase(name, "colspan")) span[1] = spanValue(value); + } + p.i = i + 1; + try p.apply(kind, close, span); + } + + fn apply(p: *Parser, kind: Tag, close: bool, span: [2]u32) ParseError!void { + switch (p.where) { + .before => { + if (kind != .table or close) return error.Malformed; + p.where = .inside; + return; + }, + .after => return error.Malformed, + .inside => {}, + } + switch (kind) { + .table => { + if (!close) return error.Malformed; // a table inside a cell + try p.endCell(); + p.row_open = false; + p.where = .after; + }, + .tr => { + try p.endCell(); + p.row_open = false; + if (!close) try p.openRow(); + }, + .td, .th => { + try p.endCell(); + if (!close) try p.openCell(kind == .th, span); + }, + .br => if (p.cell != null) { + p.space = false; + try p.text.append(p.a, '\n'); + }, + .other => {}, + } + } + + fn openRow(p: *Parser) ParseError!void { + if (p.rows == max_rows) return error.Malformed; + p.rows += 1; + p.row_open = true; + p.col = 0; + } + + /// A cell at the open row's first free column; its spans clipped to the grid and to columns earlier rows cover. + fn openCell(p: *Parser, head: bool, span: [2]u32) ParseError!void { + if (!p.row_open) return error.Malformed; + const row = p.rows - 1; + while (p.col < max_cols and p.until[p.col] > row) p.col += 1; + if (p.col == max_cols) return error.Malformed; + var cols: u32 = 0; + while (cols < span[1] and p.col + cols < max_cols and p.until[p.col + cols] <= row) cols += 1; + const rows = @min(span[0], max_rows - row); + for (p.until[p.col..][0..cols]) |*u| u.* = row + rows; + p.cell = .{ .text = "", .row = row, .col = p.col, .rows = rows, .cols = cols, .head = head }; + p.col += cols; + } + + fn endCell(p: *Parser) ParseError!void { + var cell = p.cell orelse return; + p.cell = null; + p.space = false; + defer p.text.clearRetainingCapacity(); + try p.cells.ensureUnusedCapacity(p.a, 1); + cell.text = try p.a.dupe(u8, std.mem.trim(u8, p.text.items, " \n")); + p.cells.appendAssumeCapacity(cell); + } +}; diff --git a/zig/src/core/ocr_html_test.zig b/zig/src/core/ocr_html_test.zig new file mode 100644 index 00000000000..dafaf11bf34 --- /dev/null +++ b/zig/src/core/ocr_html_test.zig @@ -0,0 +1,269 @@ +//! ocr_html.zig against tables dots.mocr and GLM-OCR wrote, and the subset's edges. +const std = @import("std"); +const html = @import("ocr_html.zig"); +const t = std.testing; + +// dots.mocr, TABLE II of the layout oracle (run008), first body rows kept. +const dots_table = "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "
ControlsHodgkin lymphomaNon-Hodgkin lymphoma
CasesOR95% CICasesOR95% CI
Never Exposed (reference group)2,2733151.00—1,8231.00—
>1 year9070.840.37-1.91861.170.87-1.59
"; + +// GLM-OCR, oracle case table-007: no header row. +const glm_table = "" ++ + "
ItemQuantity
Apples3
Pears5
"; + +fn expectCell(got: html.Cell, text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool) !void { + const want: html.Cell = .{ .text = text, .row = row, .col = col, .rows = rows, .cols = cols, .head = head }; + try t.expectEqualDeep(want, got); +} + +fn expectRows(a: std.mem.Allocator, table: html.Table, want: []const []const u8) !void { + const got = try html.rowsText(a, table); + try t.expectEqual(want.len, got.len); + for (want, got) |w, g| try t.expectEqualStrings(w, g); +} + +test "ocr html: the dots.mocr table keeps its spans" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, dots_table); + try t.expectEqual(@as(u32, 4), table.rows); + try t.expectEqual(@as(u32, 8), table.cols); + try t.expect(table.spanned); + try expectCell(table.cells[0], "", 0, 0, 2, 1, true); + const nhl = html.Cell{ .text = "Non-Hodgkin lymphoma", .row = 0, .col = 5, .rows = 1, .cols = 3, .head = true }; + try t.expectEqualDeep(nhl, table.cells[3]); + try expectCell(table.cells[4], "Cases", 1, 2, 1, 1, true); + try t.expect(!table.cells[10].head); + try expectRows(a, table, &.{ + "\tControls\tHodgkin lymphoma\t\t\tNon-Hodgkin lymphoma\t\t", + "\t\tCases\tOR\t95% CI\tCases\tOR\t95% CI", + "Never Exposed (reference group)\t2,273\t315\t1.00\t—\t1,823\t1.00\t—", + ">1 year\t90\t7\t0.84\t0.37-1.91\t86\t1.17\t0.87-1.59", + }); +} + +test "ocr html: the GLM-OCR table has no header row and becomes a pipe table with its first row the header" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, glm_table); + try t.expect(!table.spanned); + for (table.cells) |c| try t.expect(!c.head); + const want = "| Item | Quantity |\n| --- | --- |\n| Apples | 3 |\n| Pears | 5 |"; + try t.expectEqualStrings(want, try html.pipe(a, table)); + try t.expectEqualStrings("" ++ + "
ItemQuantity
Apples3
Pears5
", try html.normalized(a, table)); +} + +test "ocr html: entities decode and unknown ones stay as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "
a & b <c> "q" ABC x y z" ++ + " &bogus; & � � � � &
"; + const table = try html.parse(a, src); + const want = "a & b \"q\" ABC x y z &bogus; & \u{fffd} \u{fffd} \u{fffd} � &"; + try t.expectEqualStrings(want, table.cells[0].text); +} + +test "ocr html: other tags drop and keep their text; names and attributes in any case and quoting" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "\n " ++ + "
" ++ + "Bold and it2!
xy
\n"; + const table = try html.parse(a, src); + try t.expectEqual(@as(u32, 3), table.cols); + try expectCell(table.cells[0], "Bold and it2!", 0, 0, 1, 2, true); + try expectCell(table.cells[2], "y", 1, 1, 1, 2, false); +} + +test "ocr html:
is a line break in the cell and a space in pipe cells and text rows" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "

line one
line two
\n three
b
"; + const table = try html.parse(a, src); + try t.expectEqualStrings("line one\nline two\nthree", table.cells[0].text); + try t.expectEqualStrings("| line one line two three | b |\n| --- | --- |", try html.pipe(a, table)); + try expectRows(a, table, &.{"line one line two three\tb"}); + const clean = "
line one
line two
three		b
"; + try t.expectEqualStrings(clean, try html.normalized(a, table)); +} + +test "ocr html: pipe cells escape |" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
a|b
|c
"); + try t.expectEqualStrings("| a\\|b | |\n| --- | --- |\n| \\| | c |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal < that would start a tag, and keep other < as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
<em>x</em>a < b
"); + try t.expectEqualStrings("| <em>x</em> | a < b |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal & that would start an entity, and keep other & as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
&lt; &#38; &#x26;AT&T & b; &;
"); + try t.expectEqualStrings("| &lt; &#38; &#x26; | AT&T & b; &; |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: an unknown tag whose name starts like a table tag is dropped" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
xy
"); + try t.expectEqual(@as(usize, 1), table.cells.len); + try t.expectEqualStrings("xy", table.cells[0].text); +} + +test "ocr html: malformed tables and tables past the limits" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bad = [_][]const u8{ + "", // unclosed table + "
a
", // unclosed quote + "
a
a