From b3af83a20583bf105ebcd3c7f78ed79b97fe931f Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 29 Oct 2026 19:37:41 +0100 Subject: [PATCH] hip flashnext: random-access advice on the n-gram table mappings A prompt's n-gram gather reads one 100-byte row per id at random rows, but the table shards were mapped without advice, so each faulted page pulled a readahead window: about 1.4 GB read from disk for a first short request. The table tensors' spans now get MADV_RANDOM after mapping. Fixes #1602 --- zig/src/core/checkpoint.zig | 66 +++++++++++++++++++++++ zig/src/core/safetensors.zig | 11 ++++ zig/src/families/flashnext/hip/tables.zig | 3 ++ 3 files changed, 80 insertions(+) diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index f16b8a454a..d888d8b9f7 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -164,6 +164,17 @@ pub const Checkpoint = struct { return prefetching.start(self, self.scoped(&buf, short) catch return .{}, readers); } + /// `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). pub fn unused(self: *const Checkpoint) usize { var n: usize = 0; @@ -308,6 +319,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; diff --git a/zig/src/core/safetensors.zig b/zig/src/core/safetensors.zig index cbb5e2ea4d..bd94227aee 100644 --- a/zig/src/core/safetensors.zig +++ b/zig/src/core/safetensors.zig @@ -272,6 +272,17 @@ pub const File = struct { 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 { const e = self.names.get(name) orelse return null; return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.data + e.begin .. self.data + e.end] }; diff --git a/zig/src/families/flashnext/hip/tables.zig b/zig/src/families/flashnext/hip/tables.zig index 21942ab6ed..1d92b527cf 100644 --- a/zig/src/families/flashnext/hip/tables.zig +++ b/zig/src/families/flashnext/hip/tables.zig @@ -184,6 +184,9 @@ pub const Tables = struct { if (spec == null) spec = (try c.declared(stem, core.quant.keys.canonical)) orelse return error.UnsupportedQuantization; if (w.shape[0] != n or sc.dtype != .bf16 or bi.dtype != .bf16) return error.UnexpectedNgramTable; + // a gather reads these rows in no order: readahead per row would fetch a window for it + for ([_][]const u8{ "weight", "scales", "biases" }) |part| + try ck.adviseRandom(try std.fmt.allocPrint(a, "{s}.{s}", .{ stem, part })); s.* = .{ .words = w.bytes, .scales = sc.bytes, .biases = bi.bytes }; } const dims = c.ple_dim / heads;