hip flashnext: parallel n-gram gather, release weight pages after upload

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

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

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(

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,