From 5ce6299b2ccc25e1ca63dfded270fbb11f50e585 Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 29 Oct 2026 15:07:25 +0100 Subject: [PATCH 1/2] hip flashnext: parallel n-gram gather, release weight pages after upload The prompt's n-gram gather copied each id's rows one at a time on the host. Rows now locate the shard by division when the shard sizes allow it and split the ids across up to eight threads into disjoint staging rows, same bytes. After every weight is on the device the checkpoint's pages are dropped from the mapping and the page cache, except the n-gram table spans a gather still reads. Fixes #1606 --- zig/src/core/checkpoint.zig | 29 +- zig/src/core/safetensors.zig | 64 +++- zig/src/families/flashnext/hip/engine.zig | 3 + zig/src/families/flashnext/hip/forward.zig | 8 +- zig/src/families/flashnext/hip/tables.zig | 384 +++++++++++++++++++-- zig/src/families/radio/weights.zig | 2 +- 6 files changed, 452 insertions(+), 38 deletions(-) diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index e2e2b6f0d8..3c370ee75b 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -146,7 +146,12 @@ pub const Checkpoint = struct { /// Every file's pages out of the page cache (safetensors.File.release). pub fn release(self: *const Checkpoint) void { - for (self.files.items) |*f| f.release(); + for (self.files.items) |*f| f.release(self.gpa); + } + + /// `release`, leaving `keep`'s bytes (parts of these files read later) in the pages they live in. + pub fn releaseExcept(self: *const Checkpoint, keep: []const []const u8) void { + for (self.files.items) |*f| f.releaseExcept(keep, self.gpa); } /// Names no `get` took (the loader refuses a checkpoint with leftovers). @@ -311,6 +316,28 @@ test "after release the mapping stays readable: a later read faults the same byt try std.testing.expectEqualSlices(u8, image[image.len - 4096 ..], t.bytes); } +test "releaseExcept keeps the kept tensor's bytes where they are, and a read of them is unchanged" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const image = try testImage(gpa, "x", 1 << 20); + defer gpa.free(image); + for (image[image.len - (1 << 20) ..], 0..) |*b, i| b.* = @truncate(i * 13 + 3); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = image }); + var path: [64]u8 = undefined; + var ck = try Checkpoint.openModel(gpa, io, try std.fmt.bufPrint(&path, ".zig-cache/tmp/{s}", .{tmp.sub_path})); + defer ck.close(); + const t = try ck.get("x"); + // a middle span, so bytes are dropped on both sides of it and the mapping stays readable either way + const keep = t.bytes[4096 .. 4096 + 8192]; + const copy = try gpa.dupe(u8, keep); + defer gpa.free(copy); + ck.releaseExcept(&.{keep}); + try std.testing.expectEqualSlices(u8, copy, keep); + try std.testing.expectEqualSlices(u8, image[image.len - (1 << 20) ..], t.bytes); +} + test "release keeps a map that fell back to reads (no mmap): its heap copy is the only one" { const gpa = std.testing.allocator; var threaded: std.Io.Threaded = .init(gpa, .{ .disable_memory_mapping = true }); diff --git a/zig/src/core/safetensors.zig b/zig/src/core/safetensors.zig index 2214cc4a18..450cb41b2d 100644 --- a/zig/src/core/safetensors.zig +++ b/zig/src/core/safetensors.zig @@ -246,12 +246,20 @@ pub const File = struct { } /// The file's pages out of the mapping and the page cache (unified memory needs the room); Linux only, refaults. - pub fn release(self: *const File) void { + pub fn release(self: *const File, scratch: std.mem.Allocator) void { + self.releaseExcept(&.{}, scratch); + } + + /// `release`, with `keep`'s bytes (parts of this mapping read later) left in the pages they live in. + pub fn releaseExcept(self: *const File, keep: []const []const u8, scratch: std.mem.Allocator) void { if (@import("builtin").os.tag != .linux) return; - const linux = std.os.linux; // a map that fell back to a heap copy (no mmap) holds the only bytes: DONTNEED would zero them - if (self.map.section != null) _ = linux.madvise(self.map.memory.ptr, self.map.memory.len, linux.MADV.DONTNEED); - _ = linux.fadvise(self.file.handle, 0, 0, linux.POSIX_FADV.DONTNEED); + const mem = self.map.memory; + if (self.map.section == null) { + _ = std.os.linux.fadvise(self.file.handle, 0, 0, std.os.linux.POSIX_FADV.DONTNEED); + return; + } + dropExcept(mem.ptr, mem.len, self.file.handle, keep, scratch); } pub fn get(self: *const File, name: []const u8) ?Tensor { @@ -260,6 +268,54 @@ pub const File = struct { } }; +/// A mapped file's pages out of memory and its page cache, except the pages holding `keep`'s bytes. +fn dropExcept( + map: [*]u8, + map_len: usize, + fd: std.posix.fd_t, + keep: []const []const u8, + scratch: std.mem.Allocator, +) void { + const linux = std.os.linux; + const base = @intFromPtr(map); + const page = std.heap.pageSize(); + var spans: [8][2]usize = undefined; + const kept = if (keep.len <= spans.len) spans[0..keep.len] else scratch.alloc([2]usize, keep.len) catch { + _ = linux.madvise(map, map_len, linux.MADV.DONTNEED); + _ = linux.fadvise(fd, 0, 0, linux.POSIX_FADV.DONTNEED); + return; + }; + defer if (kept.len > spans.len) scratch.free(kept); + var n: usize = 0; + for (keep) |k| { + if (k.len == 0) continue; + const lo = @intFromPtr(k.ptr); + if (lo < base or lo + k.len > base + map_len) continue; + kept[n] = .{ + std.mem.alignBackward(usize, lo - base, page), + @min(map_len, std.mem.alignForward(usize, lo - base + k.len, page)), + }; + n += 1; + } + std.mem.sort([2]usize, kept[0..n], {}, struct { + fn less(_: void, x: [2]usize, y: [2]usize) bool { + return x[0] < y[0]; + } + }.less); + var at: usize = 0; + for (kept[0..n]) |sp| { + if (sp[0] > at) gap(map, fd, at, sp[0] - at); + at = @max(at, sp[1]); + } + if (at < map_len) gap(map, fd, at, map_len - at); +} + +/// One run of a mapped shard's bytes out of memory and the page cache: a later read faults them back in. +fn gap(map: [*]u8, fd: std.posix.fd_t, at: usize, len: usize) void { + _ = std.os.linux.madvise(map + at, len, std.os.linux.MADV.DONTNEED); + _ = std.os.linux.fadvise(fd, @intCast(at), @intCast(len), std.os.linux.POSIX_FADV.DONTNEED); +} + test "header entries" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/families/flashnext/hip/engine.zig b/zig/src/families/flashnext/hip/engine.zig index ce3b74fd5e..0812c5115f 100644 --- a/zig/src/families/flashnext/hip/engine.zig +++ b/zig/src/families/flashnext/hip/engine.zig @@ -145,6 +145,9 @@ pub const Engine = struct { }; errdefer e.host.staging.free(); e.eos = e.tables.hasher.ref.eos; + // every weight is on the device now: only the n-gram rows a gather reads stay resident + const live = try e.tables.spans(e.keep.allocator()); + e.ck.releaseExcept(live); e.arena = try hip.Arena.init(&e.driver, try e.scratchBytes()); errdefer e.arena.deinit(); e.drawer = try serve.draw.Drawer.init( diff --git a/zig/src/families/flashnext/hip/forward.zig b/zig/src/families/flashnext/hip/forward.zig index 9af288e27f..3f3d75abbc 100644 --- a/zig/src/families/flashnext/hip/forward.zig +++ b/zig/src/families/flashnext/hip/forward.zig @@ -159,7 +159,7 @@ pub const Inputs = struct { pub fn stage(m: *const model.Model, host: *Host, segs: []const Seg, window: bool) !Staged { const d = m.dims; const t = host.tables; - if (d.ple_dim / t.dims != ngram.heads) return error.UnsupportedNgramHeads; + if (d.ple_dim / t.rows.dims != ngram.heads) return error.UnsupportedNgramHeads; var rows: usize = 0; for (segs) |g| rows += g.tokens.len; const n16 = rows * ngram.heads; @@ -268,9 +268,9 @@ fn ple( @backingInt(Kind.bf16), @ptrFromInt(in.local), @intCast(rows * ngram.heads), - t.bits, - @intCast(t.group), - @intCast(t.dims), + t.rows.bits, + @intCast(t.rows.group), + @intCast(t.rows.dims), @ptrFromInt(emb.ptr), @backingInt(m.act), o.stream, diff --git a/zig/src/families/flashnext/hip/tables.zig b/zig/src/families/flashnext/hip/tables.zig index 1bf51c0c4e..3e0f128e17 100644 --- a/zig/src/families/flashnext/hip/tables.zig +++ b/zig/src/families/flashnext/hip/tables.zig @@ -6,24 +6,151 @@ const tensors = @import("tensors.zig"); const ngram = @import("ngram.zig"); /// One shard's packed rows and their group tables, as mapped. -const Shard = struct { words: []const u8, scales: []const u8, biases: []const u8 }; +pub const Shard = struct { words: []const u8, scales: []const u8, biases: []const u8 }; -pub const Tables = struct { - shards: []Shard, +/// A gather copies the checkpoint's rows in id order: a staging too small, or a row past the last shard. +const Error = error{ StagingTooSmall, NgramRowOutOfRange }; + +/// Threads a gather takes at most, and rows a thread takes at least before the gather spreads. +const max_workers = 8; +const rows_per_worker = 256; + +/// The shards' rows as mapped: what a gather reads, apart from the hasher that names the rows. +pub const Rows = struct { + shards: []const Shard, /// The first row of each shard, and one past the last. - starts: []usize, + starts: []const usize, + /// The first `extra` shards hold `wide` rows and the rest `narrow`; `wide` is 0 when the sizes do not fit. + wide: usize, + narrow: usize, + extra: usize, dims: usize, bits: u8, group: usize, - hasher: ngram.Hasher, /// Bytes of a row's packed words, and of its scales (as of its biases). + pub fn wordBytes(r: *const Rows) usize { + return r.dims * r.bits / 8; + } + + pub fn tableBytes(r: *const Rows) usize { + return r.dims / r.group * 2; + } + + /// The shard sizes the reference splits rows into: `wide` rows for the first `extra` shards, then `narrow`. + pub const Split = struct { wide: usize, narrow: usize, extra: usize }; + + /// The sizes `starts` shows, or zeros when they follow no such run and a row must be searched for. + pub fn split(starts: []const usize) Split { + const n = starts.len - 1; + if (n == 0 or starts[1] - starts[0] == 0) return .{ .wide = 0, .narrow = 0, .extra = 0 }; + const wide = starts[1] - starts[0]; + var extra: usize = 1; + while (extra < n and starts[extra + 1] - starts[extra] == wide) extra += 1; + const narrow = if (extra < n) starts[extra + 1] - starts[extra] else wide; + if (narrow == 0) return .{ .wide = 0, .narrow = 0, .extra = 0 }; + for (extra..n) |i| { + if (starts[i + 1] - starts[i] != narrow) return .{ .wide = 0, .narrow = 0, .extra = 0 }; + } + return .{ .wide = wide, .narrow = narrow, .extra = extra }; + } + + /// Row `id`'s shard and its row in it: a shard's rows are a contiguous run, so the sizes need no search. + pub fn locate(r: *const Rows, id: u32) Error!struct { shard: usize, row: usize } { + const at: usize = id; + if (at >= r.starts[r.shards.len]) return error.NgramRowOutOfRange; + if (r.wide != 0) { + const first = r.wide * r.extra; + if (at < first) return .{ .shard = at / r.wide, .row = at % r.wide }; + return .{ .shard = r.extra + (at - first) / r.narrow, .row = (at - first) % r.narrow }; + } + const s = std.sort.upperBound(usize, r.starts, at, struct { + fn order(x: usize, y: usize) std.math.Order { + return std.math.order(x, y); + } + }.order) - 1; + return .{ .shard = s, .row = at - r.starts[s] }; + } + + /// Every byte a gather reads: the pages a checkpoint release over the weights must leave mapped. + pub fn spans(r: *const Rows, a: std.mem.Allocator) ![]const []const u8 { + const out = try a.alloc([]const u8, r.shards.len * 3); + for (r.shards, 0..) |s, i| { + out[i * 3] = s.words; + out[i * 3 + 1] = s.scales; + out[i * 3 + 2] = s.biases; + } + return out; + } + + /// Each id's packed words, scales and biases copied into the three staging slices, in id order. + pub fn gather(r: *const Rows, ids: []const u32, words: []u8, scales: []u8, biases: []u8) !void { + const wb = r.wordBytes(); + const tb = r.tableBytes(); + if (words.len < ids.len * wb or scales.len < ids.len * tb or biases.len < ids.len * tb) + return error.StagingTooSmall; + const workers = @min(max_workers, @max(1, ids.len / rows_per_worker)); + const per = (ids.len + workers - 1) / workers; + var results: [max_workers]Error!void = @splat({}); + var threads: [max_workers]?std.Thread = @splat(null); + for (1..workers) |w| { + const last = @min(ids.len, (w + 1) * per); + const job = .{ &results[w], r, ids, words, scales, biases, w * per, last }; + threads[w] = std.Thread.spawn(.{}, run, job) catch null; + } + run(&results[0], r, ids, words, scales, biases, 0, @min(ids.len, per)); + for (1..workers) |w| { + const last = @min(ids.len, (w + 1) * per); + if (threads[w]) |th| th.join() else run(&results[w], r, ids, words, scales, biases, w * per, last); + } + for (results[0..workers]) |res| try res; + } +}; + +/// Ids [first, last) into their staging rows: every thread copies its own rows, so no two write one. +fn run( + out: *Error!void, + r: *const Rows, + ids: []const u32, + words: []u8, + scales: []u8, + biases: []u8, + first: usize, + last: usize, +) void { + out.* = gatherRows(r, ids, words, scales, biases, first, last); +} + +fn gatherRows( + r: *const Rows, + ids: []const u32, + words: []u8, + scales: []u8, + biases: []u8, + first: usize, + last: usize, +) Error!void { + const wb = r.wordBytes(); + const tb = r.tableBytes(); + for (ids[first..last], first..) |id, i| { + const at = try r.locate(id); + const sh = r.shards[at.shard]; + @memcpy(words[i * wb ..][0..wb], sh.words[at.row * wb ..][0..wb]); + @memcpy(scales[i * tb ..][0..tb], sh.scales[at.row * tb ..][0..tb]); + @memcpy(biases[i * tb ..][0..tb], sh.biases[at.row * tb ..][0..tb]); + } +} + +pub const Tables = struct { + rows: Rows, + hasher: ngram.Hasher, + pub fn wordBytes(t: Tables) usize { - return t.dims * t.bits / 8; + return t.rows.wordBytes(); } pub fn tableBytes(t: Tables) usize { - return t.dims / t.group * 2; + return t.rows.tableBytes(); } /// The tables of the PLE layer `layer` from the mapped checkpoint; the shards' rows as the reference splits them. @@ -61,34 +188,235 @@ pub const Tables = struct { } const dims = c.ple_dim / heads; if (shards[0].words.len != (starts[1] - starts[0]) * dims * spec.?.bits / 8) return error.UnexpectedNgramTable; + const sizes = Rows.split(starts); return .{ - .shards = shards, - .starts = starts, - .dims = dims, - .bits = spec.?.bits, - .group = spec.?.group, + .rows = .{ + .shards = shards, + .starts = starts, + .wide = sizes.wide, + .narrow = sizes.narrow, + .extra = sizes.extra, + .dims = dims, + .bits = spec.?.bits, + .group = spec.?.group, + }, .hasher = try ngram.Hasher.init(c), }; } /// Each id's packed words, scales and biases copied into the three staging slices, in id order. pub fn gather(t: *const Tables, ids: []const u32, words: []u8, scales: []u8, biases: []u8) !void { - const wb = t.wordBytes(); - const tb = t.tableBytes(); - if (words.len < ids.len * wb or scales.len < ids.len * tb or biases.len < ids.len * tb) - return error.StagingTooSmall; - for (ids, 0..) |id, i| { - const s = std.sort.upperBound(usize, t.starts, @as(usize, id), struct { - fn order(x: usize, y: usize) std.math.Order { - return std.math.order(x, y); - } - }.order) - 1; - if (s >= t.shards.len) return error.NgramRowOutOfRange; - const row = id - t.starts[s]; - const sh = t.shards[s]; - @memcpy(words[i * wb ..][0..wb], sh.words[row * wb ..][0..wb]); - @memcpy(scales[i * tb ..][0..tb], sh.scales[row * tb ..][0..tb]); - @memcpy(biases[i * tb ..][0..tb], sh.biases[row * tb ..][0..tb]); + return t.rows.gather(ids, words, scales, biases); + } + + /// Every byte a gather reads at runtime: the pages a checkpoint release must leave mapped. + pub fn spans(t: *const Tables, a: std.mem.Allocator) ![]const []const u8 { + return t.rows.spans(a); + } +}; + +/// The gather before the sizes were read off `starts`: one row at a time, its shard searched for. +fn searchedGather(r: *const Rows, ids: []const u32, words: []u8, scales: []u8, biases: []u8) !void { + const wb = r.wordBytes(); + const tb = r.tableBytes(); + if (words.len < ids.len * wb or scales.len < ids.len * tb or biases.len < ids.len * tb) + return error.StagingTooSmall; + for (ids, 0..) |id, i| { + const s = std.sort.upperBound(usize, r.starts, @as(usize, id), struct { + fn order(x: usize, y: usize) std.math.Order { + return std.math.order(x, y); + } + }.order) - 1; + if (s >= r.shards.len) return error.NgramRowOutOfRange; + const row = id - r.starts[s]; + const sh = r.shards[s]; + @memcpy(words[i * wb ..][0..wb], sh.words[row * wb ..][0..wb]); + @memcpy(scales[i * tb ..][0..tb], sh.scales[row * tb ..][0..tb]); + @memcpy(biases[i * tb ..][0..tb], sh.biases[row * tb ..][0..tb]); + } +} + +/// Synthetic shards of `sizes` rows (`dims` values a row at `bits`, groups of `group`): bytes hashed from their offset. +const Fake = struct { + shards: []Shard, + starts: []usize, + arrays: [3][]u8, + rows: Rows, + + fn init(a: std.mem.Allocator, sizes: []const usize, dims: usize, bits: u8, group: usize) !Fake { + var f: Fake = .{ + .shards = try a.alloc(Shard, sizes.len), + .starts = try a.alloc(usize, sizes.len + 1), + .arrays = .{ &.{}, &.{}, &.{} }, + .rows = undefined, + }; + const wb = dims * bits / 8; + const tb = dims / group * 2; + f.starts[0] = 0; + for (sizes, 0..) |n, i| f.starts[i + 1] = f.starts[i] + n; + const total = f.starts[sizes.len]; + for (&f.arrays, [_]usize{ wb, tb, tb }, [_]u8{ 7, 3, 11 }) |*arr, width, step| { + arr.* = try a.alloc(u8, total * width); + for (arr.*, 0..) |*b, i| b.* = @truncate(i *% step +% width); } + for (sizes, 0..) |n, i| { + const from = f.starts[i]; + f.shards[i] = .{ + .words = f.arrays[0][from * wb ..][0 .. n * wb], + .scales = f.arrays[1][from * tb ..][0 .. n * tb], + .biases = f.arrays[2][from * tb ..][0 .. n * tb], + }; + } + const split = Rows.split(f.starts); + f.rows = .{ + .shards = f.shards, + .starts = f.starts, + .wide = split.wide, + .narrow = split.narrow, + .extra = split.extra, + .dims = dims, + .bits = bits, + .group = group, + }; + return f; + } + + fn deinit(f: *Fake, a: std.mem.Allocator) void { + for (f.arrays) |arr| a.free(arr); + a.free(f.shards); + a.free(f.starts); + } + + fn gathered(f: *const Fake, a: std.mem.Allocator, ids: []const u32) !Staging { + const s: Staging = .{ + .words = try a.alloc(u8, ids.len * f.rows.wordBytes()), + .scales = try a.alloc(u8, ids.len * f.rows.tableBytes()), + .biases = try a.alloc(u8, ids.len * f.rows.tableBytes()), + }; + @memset(s.words, 0); + @memset(s.scales, 0); + @memset(s.biases, 0); + try f.rows.gather(ids, s.words, s.scales, s.biases); + return s; } + + const Staging = struct { + words: []u8, + scales: []u8, + biases: []u8, + fn deinit(s: *const Staging, a: std.mem.Allocator) void { + a.free(s.words); + a.free(s.scales); + a.free(s.biases); + } + }; }; + +test "a row's shard comes from the sizes alone: every id of a remainder split finds the searched shard" { + const a = std.testing.allocator; + // 128 shards of 250 rows with 3 of 251, as a checkpoint's padded row count splits + var sizes: [128]usize = @splat(250); + for (sizes[0..3]) |*n| n.* = 251; + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(a); + var total: usize = 0; + for (sizes) |n| total += n; + for (0..total) |id| try ids.append(a, @intCast(id)); + var f = try Fake.init(a, &sizes, 160, 4, 32); + defer f.deinit(a); + try std.testing.expectEqual(@as(usize, 251), f.rows.wide); + try std.testing.expectEqual(@as(usize, 250), f.rows.narrow); + try std.testing.expectEqual(@as(usize, 3), f.rows.extra); + var shard: usize = 0; + for (ids.items) |id| { + const at = try f.rows.locate(id); + if (id == f.starts[shard + 1]) shard += 1; + try std.testing.expectEqual(shard, at.shard); + try std.testing.expectEqual(@as(usize, id) - f.starts[shard], at.row); + } + try std.testing.expectError(error.NgramRowOutOfRange, f.rows.locate(@intCast(f.starts[128]))); +} + +test "gather: the sizes' fast path copies the same bytes as a shard search for every id" { + const a = std.testing.allocator; + const cases = [_][]const usize{ + &.{}, + &.{97}, + &.{40, 40, 40, 40}, + &.{41, 41, 40, 40}, + &.{40, 39, 41, 40}, + &.{ 250, 250, 250 }, + }; + for (cases) |sizes| { + var ids: [256]u32 = undefined; + var n: usize = 0; + var total: usize = 0; + for (sizes) |s| total += s; + while (n < ids.len and n < total) : (n += 1) ids[n] = @intCast((n * 37) % total); + var f = try Fake.init(a, sizes, 160, 4, 32); + defer f.deinit(a); + const want = try a.alloc(u8, n * f.rows.wordBytes()); + defer a.free(want); + const want_scales = try a.alloc(u8, n * f.rows.tableBytes()); + defer a.free(want_scales); + const want_biases = try a.alloc(u8, n * f.rows.tableBytes()); + defer a.free(want_biases); + try searchedGather(&f.rows, ids[0..n], want, want_scales, want_biases); + const got = try f.gathered(a, ids[0..n]); + defer got.deinit(a); + try std.testing.expectEqualSlices(u8, want, got.words); + try std.testing.expectEqualSlices(u8, want_scales, got.scales); + try std.testing.expectEqualSlices(u8, want_biases, got.biases); + } +} + +test "gather: rows copied by several threads equal one thread's, and a bad row or a short staging is refused" { + const a = std.testing.allocator; + var sizes: [64]usize = @splat(300); + for (sizes[0..8]) |*s| s.* = 301; + var total: usize = 0; + for (sizes) |n| total += n; + var ids: [8192]u32 = undefined; + for (&ids, 0..) |*id, i| id.* = @intCast((i * 7919) % total); + var f = try Fake.init(a, &sizes, 160, 4, 32); + defer f.deinit(a); + const want = try a.alloc(u8, ids.len * f.rows.wordBytes()); + defer a.free(want); + const want_scales = try a.alloc(u8, ids.len * f.rows.tableBytes()); + defer a.free(want_scales); + const want_biases = try a.alloc(u8, ids.len * f.rows.tableBytes()); + defer a.free(want_biases); + try searchedGather(&f.rows, &ids, want, want_scales, want_biases); + const got = try f.gathered(a, &ids); + defer got.deinit(a); + try std.testing.expectEqualSlices(u8, want, got.words); + try std.testing.expectEqualSlices(u8, want_scales, got.scales); + try std.testing.expectEqualSlices(u8, want_biases, got.biases); + // an id past the last row, and staging that cannot hold the rows the ids ask for + var past = ids; + past[0] = @intCast(f.starts[sizes.len]); + const small = try a.alloc(u8, ids.len * f.rows.wordBytes()); + defer a.free(small); + try std.testing.expectError(error.NgramRowOutOfRange, f.rows.gather(&past, small, want_scales, want_biases)); + const short = small[0 .. small.len - 1]; + try std.testing.expectError(error.StagingTooSmall, f.rows.gather(&ids, short, want_scales, want_biases)); +} + +test "spans: one byte range for each of the three tensors of every shard" { + const a = std.testing.allocator; + var ids: [4]u32 = undefined; + for (&ids, 0..) |*id, i| id.* = @intCast(i); + var f = try Fake.init(a, &.{ 2, 2, 2 }, 160, 4, 32); + defer f.deinit(a); + const ranges = try f.rows.spans(a); + defer a.free(ranges); + try std.testing.expectEqual(@as(usize, 9), ranges.len); + for (ranges, 0..) |r, i| { + const shard = f.shards[i / 3]; + try std.testing.expectEqual(r.ptr, switch (i % 3) { + 0 => shard.words.ptr, + 1 => shard.scales.ptr, + else => shard.biases.ptr, + }); + } +} diff --git a/zig/src/families/radio/weights.zig b/zig/src/families/radio/weights.zig index 2ba0c1d2b4..c2f89fdbc8 100644 --- a/zig/src/families/radio/weights.zig +++ b/zig/src/families/radio/weights.zig @@ -44,7 +44,7 @@ pub const Weights = struct { var files = try open(gpa, io, dir); defer { for (files.items) |*f| { - f.release(); + f.release(gpa); f.close(io); } files.deinit(gpa); From 022017d2974b37c2fdd772486cde1601acac2484 Mon Sep 17 00:00:00 2001 From: gilby Date: Fri, 30 Oct 2026 15:02:28 +0100 Subject: [PATCH 2/2] zig fmt --- zig/src/families/flashnext/hip/tables.zig | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/zig/src/families/flashnext/hip/tables.zig b/zig/src/families/flashnext/hip/tables.zig index 3e0f128e17..21942ab6ed 100644 --- a/zig/src/families/flashnext/hip/tables.zig +++ b/zig/src/families/flashnext/hip/tables.zig @@ -342,9 +342,9 @@ test "gather: the sizes' fast path copies the same bytes as a shard search for e const cases = [_][]const usize{ &.{}, &.{97}, - &.{40, 40, 40, 40}, - &.{41, 41, 40, 40}, - &.{40, 39, 41, 40}, + &.{ 40, 40, 40, 40 }, + &.{ 41, 41, 40, 40 }, + &.{ 40, 39, 41, 40 }, &.{ 250, 250, 250 }, }; for (cases) |sizes| {