hip flashnext: random-access advice on the n-gram table mappings

zig/src/core/checkpoint.zig

@@ -146,7 +146,23 @@ 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);
+    }
+
+    /// `short`'s bytes marked read in no order (safetensors.File.adviseRandom): a table's rows, not its file.
+    pub fn adviseRandom(self: *Checkpoint, short: []const u8) !void {
+        var buf: [320]u8 = undefined;
+        const name = try self.scoped(&buf, short);
+        for (self.files.items) |*f| if (f.get(name)) |t| {
+            f.adviseRandom(t.bytes);
+            return;
+        };
+        return error.MissingTensor;
     }
 
     /// Names no `get` took (the loader refuses a checkpoint with leftovers).
@@ -293,6 +309,61 @@ test "has finds a name without taking it; skip marks a prefix used" {
     try std.testing.expectEqual(@as(usize, 0), ck.unused());
 }
 
+/// Whether the smaps block holding `at` is marked read in no order (`rr`, MADV_RANDOM); null with no such block.
+fn randomAdviceAt(gpa: std.mem.Allocator, io: Io, at: usize) !?bool {
+    // procfs reports size 0: read the stream, not the length
+    var file = try Io.Dir.cwd().openFile(io, "/proc/self/smaps", .{});
+    defer file.close(io);
+    var text: std.ArrayList(u8) = .empty;
+    defer text.deinit(gpa);
+    var buf: [64 << 10]u8 = undefined;
+    while (true) {
+        const got = file.readStreaming(io, &.{&buf}) catch |err| switch (err) {
+            error.EndOfStream => break,
+            else => return err,
+        };
+        if (got == 0) break;
+        try text.appendSlice(gpa, buf[0..got]);
+    }
+    var lines = std.mem.splitScalar(u8, text.items, '\n');
+    var inside = false;
+    while (lines.next()) |line| {
+        if (std.mem.indexOfScalar(u8, line, '-')) |dash| {
+            const sp = std.mem.indexOfScalarPos(u8, line, dash, ' ') orelse continue;
+            const lo = std.fmt.parseInt(usize, line[0..dash], 16) catch continue;
+            const hi = std.fmt.parseInt(usize, line[dash + 1 .. sp], 16) catch continue;
+            inside = at >= lo and at < hi;
+        } else if (inside and std.mem.startsWith(u8, line, "VmFlags:")) {
+            var flags = std.mem.tokenizeAny(u8, line["VmFlags:".len..], " ");
+            while (flags.next()) |flag| if (std.mem.eql(u8, flag, "rr")) return true;
+            return false;
+        }
+    }
+    return null;
+}
+
+test "adviseRandom marks the mapping's pages as read in no order, and leaves the bytes alone" {
+    if (@import("builtin").os.tag != .linux) return error.SkipZigTest;
+    const io = std.testing.io;
+    const gpa = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    const image = try testImage(gpa, "x", 64 << 10);
+    defer gpa.free(image);
+    for (image[image.len - (64 << 10) ..], 0..) |*b, i| b.* = @truncate(i * 7 + 1);
+    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");
+    const at = @intFromPtr(t.bytes.ptr);
+    // a fresh mapping: the kernel has not been told how it will be read
+    try std.testing.expectEqual(@as(?bool, false), try randomAdviceAt(gpa, io, at));
+    try ck.adviseRandom("x");
+    try std.testing.expectEqual(@as(?bool, true), try randomAdviceAt(gpa, io, at));
+    try std.testing.expectEqualSlices(u8, image[image.len - (64 << 10) ..], t.bytes);
+}
+
 test "after release the mapping stays readable: a later read faults the same bytes back in" {
     const io = std.testing.io;
     const gpa = std.testing.allocator;
@@ -311,6 +382,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,31 @@ 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);
+    }
+
+    /// `bytes` (part of this file's mapping) marked read in no order, e.g. a table's rows; Linux only.
+    pub fn adviseRandom(self: *const File, bytes: []const u8) void {
+        if (@import("builtin").os.tag != .linux or bytes.len == 0) return;
+        // a map that fell back to a heap copy has no file pages to advise
+        if (self.map.section == null) return;
+        const page = std.heap.pageSize();
+        const lo = std.mem.alignBackward(usize, @intFromPtr(bytes.ptr), page);
+        const hi = std.mem.alignForward(usize, @intFromPtr(bytes.ptr) + bytes.len, page);
+        _ = std.os.linux.madvise(@ptrFromInt(lo), hi - lo, std.os.linux.MADV.RANDOM);
     }
 
     pub fn get(self: *const File, name: []const u8) ?Tensor {
@@ -260,6 +279,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,