From 0316c8c06b89f816a5e3539661028edac85fd72a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 7 Nov 2026 14:24:43 +0100 Subject: [PATCH 1/2] core: checkpoint.direct.Layers gains layerAt and moved, so a loader that maps every tensor first can read direct layerAt enters a layer's prefix by index and starts the next one's read. moved does the same and repoints the tensor views inside a value (fields, unions, optionals, slices) onto the layer's copies; a view outside the layer keeps the mapping. --- zig/src/core/checkpoint_direct.zig | 110 ++++++++++++++++++++++++++++- 1 file changed, 109 insertions(+), 1 deletion(-) diff --git a/zig/src/core/checkpoint_direct.zig b/zig/src/core/checkpoint_direct.zig index d59687d1b8..9aba378a23 100644 --- a/zig/src/core/checkpoint_direct.zig +++ b/zig/src/core/checkpoint_direct.zig @@ -4,6 +4,7 @@ const std = @import("std"); const builtin = @import("builtin"); const gather = @import("direct_gather.zig"); const Checkpoint = @import("checkpoint.zig").Checkpoint; +const Tensor = @import("checkpoint.zig").Tensor; /// One prefix's tensors: their names, their views (repointed at the copies), and the copies by name. const Slot = struct { @@ -13,8 +14,12 @@ const Slot = struct { views: std.ArrayList([]const u8) = .empty, ptrs: std.ArrayList(*[]const u8) = .empty, copies: std.StringHashMapUnmanaged([]const u8) = .empty, + mapped: std.ArrayList(usize) = .empty, // each view's address in the mapping, before the read repointed it + moved: std.AutoHashMapUnmanaged(usize, []const u8) = .empty, // the copies by mapped address (repoint) fn deinit(s: *Slot, gpa: std.mem.Allocator) void { + s.mapped.deinit(gpa); + s.moved.deinit(gpa); s.names.deinit(gpa); s.views.deinit(gpa); s.ptrs.deinit(gpa); @@ -33,6 +38,7 @@ pub const Layers = struct { slots: [2]Slot = .{ .{}, .{} }, at: usize = 0, // the slot the next read fills ahead: bool = false, // slots[at ^ 1] is being read by the pair's thread + current: ?*Slot = null, // the slot `enter` handed out last pub fn init(gpa: std.mem.Allocator, ck: *Checkpoint, readers: usize) Layers { var l: Layers = .{ .gpa = gpa, .pair = .init(gpa, readers) }; @@ -63,13 +69,36 @@ pub const Layers = struct { l.ahead = false; try s.copies.ensureTotalCapacity(l.gpa, @intCast(s.names.items.len)); for (s.names.items, s.views.items) |k, v| s.copies.putAssumeCapacity(k, v); + try s.moved.ensureTotalCapacity(l.gpa, @intCast(s.views.items.len)); + for (s.mapped.items, s.views.items) |a, v| s.moved.putAssumeCapacity(a, v); ck.copies = &s.copies; + l.current = s; // the slot read ahead into held the layer before this one: its bytes are on the device const t = try l.collect(next orelse return); try l.pair.start(l.sources[0..l.n], t.ptrs.items, null); l.ahead = true; } + /// `enter` on `fmt`'s prefix for layer `i`, with layer `next`'s read begun (none on the last). + pub fn layerAt(l: *Layers, comptime fmt: []const u8, i: usize, next: ?usize) !void { + if (l.ck == null) return; + var b: [2][160]u8 = undefined; + const n = if (next) |j| try std.fmt.bufPrint(&b[1], fmt, .{j}) else null; + try l.enter(try std.fmt.bufPrint(&b[0], fmt, .{i}), n); + } + + /// The tensors in `v` (a pointer: fields, unions, optionals, arrays, slices) moved onto the current layer's copies. + pub fn repoint(l: *const Layers, v: anytype) void { + if (l.current) |s| walk(&s.moved, v); + } + + /// `v.*` (layer `i` of `n`, mapped before the loop) on its copies: layerAt, then repoint. + pub fn moved(l: *Layers, v: anytype, comptime fmt: []const u8, i: usize, n: usize) !@TypeOf(v.*) { + try l.layerAt(fmt, i, if (i + 1 < n) i + 1 else null); + l.repoint(v); + return v.*; + } + pub fn deinit(l: *Layers) void { l.pair.deinit(); if (l.ck) |ck| ck.copies = null; @@ -86,6 +115,8 @@ pub const Layers = struct { s.names.clearRetainingCapacity(); s.views.clearRetainingCapacity(); s.copies.clearRetainingCapacity(); + s.mapped.clearRetainingCapacity(); + s.moved.clearRetainingCapacity(); const full = s.prefix[0..s.prefix_len]; for (ck.files.items) |*f| for (f.names.keys()) |k| { if (!std.mem.startsWith(u8, k, full)) continue; @@ -93,11 +124,45 @@ pub const Layers = struct { try s.views.append(gpa, f.get(k).?.bytes); }; try s.ptrs.resize(gpa, s.views.items.len); - for (s.views.items, s.ptrs.items) |*v, *p| p.* = v; + try s.mapped.resize(gpa, s.views.items.len); + for (s.views.items, s.ptrs.items, s.mapped.items) |*v, *p, *a| { + p.* = v; + a.* = @intFromPtr(v.ptr); + } return s; } }; +/// Each Tensor reachable from pointer `v` takes its copy's bytes; one not in the layer keeps the mapping. +fn walk(moved: *const std.AutoHashMapUnmanaged(usize, []const u8), v: anytype) void { + const T = @typeInfo(@TypeOf(v)).pointer.child; + if (T == Tensor) { + if (v.bytes.len != 0) if (moved.get(@intFromPtr(v.bytes.ptr))) |c| { + v.bytes = c; + }; + return; + } + switch (@typeInfo(T)) { + .@"struct" => |st| inline for (st.field_names) |f| walk(moved, &@field(v.*, f)), + .@"union" => switch (v.*) { + inline else => |*x| walk(moved, x), + }, + .optional => if (v.*) |*x| walk(moved, x), + .array => |a| if (holds(a.child)) for (v) |*x| walk(moved, x), + .pointer => |p| if (p.size == .slice and !p.attrs.@"const" and holds(p.child)) for (v.*) |*x| walk(moved, x), + else => {}, + } +} + +/// Whether a value of `T` can hold a Tensor (walk skips byte and number arrays). +fn holds(comptime T: type) bool { + return switch (@typeInfo(T)) { + .@"struct", .@"union", .optional => true, + .array => |a| holds(a.child), + else => false, + }; +} + /// A safetensors image of U8 tensors, one per name, each `bytes` long and filled from its index. fn testImage(a: std.mem.Allocator, names: []const []const u8, bytes: usize) ![]u8 { var header: std.ArrayList(u8) = .empty; @@ -164,6 +229,45 @@ test "each layer's tensors come back as copies off the mapping, the same bytes, try std.testing.expectEqual(@as(usize, 0), ck.unused()); } +test "layerAt names a layer and the next; repoint moves a map-first loader's views onto the same copies" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const names = [_][]const u8{ "layers.0.a", "layers.0.b", "layers.1.a", "layers.1.b", "head" }; + const image = try testImage(gpa, &names, 2 * 4096 + 5); + defer gpa.free(image); + 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 Held = struct { a: Tensor, more: []Tensor, maybe: ?Tensor }; + var more: [2][1]Tensor = undefined; + var map: [2]Held = undefined; + for (&map, &more, 0..) |*m, *x, i| { + var b: [32]u8 = undefined; + x[0] = try ck.get(try std.fmt.bufPrint(&b, "layers.{d}.b", .{i})); + const a = try ck.get(try std.fmt.bufPrint(&b, "layers.{d}.a", .{i})); + m.* = .{ .a = a, .more = x, .maybe = try ck.get("head") }; + } + const mapped = [2][2][]const u8{ .{ map[0].a.bytes, more[0][0].bytes }, .{ map[1].a.bytes, more[1][0].bytes } }; + var l = Layers.init(gpa, &ck, 2); + defer l.deinit(); + if (!l.on()) return error.SkipZigTest; + for (&map, 0..) |*m, i| { + try l.layerAt("layers.{d}.", i, if (i == 0) 1 else null); + l.repoint(m); + try std.testing.expect(m.a.bytes.ptr != mapped[i][0].ptr and m.more[0].bytes.ptr != mapped[i][1].ptr); + try std.testing.expectEqualSlices(u8, mapped[i][0], m.a.bytes); + try std.testing.expectEqualSlices(u8, mapped[i][1], m.more[0].bytes); + // the same copies `get` hands out, and a tensor outside the layer stays on the mapping + var b: [32]u8 = undefined; + const got = try ck.get(try std.fmt.bufPrint(&b, "layers.{d}.a", .{i})); + try std.testing.expectEqual(got.bytes.ptr, m.a.bytes.ptr); + try std.testing.expectEqual((try ck.get("head")).bytes.ptr, m.maybe.?.bytes.ptr); + } +} + test "TF_LOAD_DIRECT off or no readers leaves every call a no-op" { const gpa = std.testing.allocator; var ck: Checkpoint = .{ .gpa = gpa, .io = std.testing.io }; @@ -171,5 +275,9 @@ test "TF_LOAD_DIRECT off or no readers leaves every call a no-op" { defer l.deinit(); try std.testing.expect(!l.on()); try l.enter("layers.0.", "layers.1."); + try l.layerAt("layers.{d}.", 0, 1); + var t: Tensor = .{ .dtype = .u8, .rank = 0, .shape = undefined, .bytes = "x" }; + l.repoint(&t); + try std.testing.expectEqualStrings("x", t.bytes); try std.testing.expectEqual(@as(?*const std.StringHashMapUnmanaged([]const u8), null), ck.copies); } From 151cb02599ad9a1a8bc67328b74968b5471090c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 7 Nov 2026 14:24:43 +0100 Subject: [PATCH 2/2] spark25, dory cuda: load each layer with direct reads, the next one while it packs One checkpoint.direct call per layer: layerAt in Spark 2.5's get-based loader, moved over Dory's tensor map. While direct reads are on, Spark's opt-in read-ahead gets no readers. TF_LOAD_DIRECT=0 keeps the mapped reads. --- CHANGELOG.md | 2 ++ docs/optimization-ledger.md | 5 ++++- zig/src/families/dory/cuda_weights.zig | 6 +++++- zig/src/families/spark25/weights.zig | 6 +++++- 4 files changed, 16 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 264e405603..10f9e0ecf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -266,6 +266,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Gemma 4 E4B (NVIDIA Nemotron 3.5 Bolt E4B) runs on CUDA from `mlx_lm.convert` 4-bit checkpoints: per-layer inputs in their own launch and layers 24 to 41 on the caches of 22 and 23, with replies equal across streams and resumes (#1222, #2230). Thanks to @jschmied. - Qwen3-Reranker-8B on HIP computes a rerank request's shared start (the system prompt, the instruction and the query) once a pass instead of once a document, with the same responses byte for byte: 16 documents of 323 tokens, 87 of them shared, take 1.98 s instead of 2.63 s on the dot2 tile and 0.73 s instead of 0.99 s with the gfx12 WMMA tile (opt-in, `--policy matrix=on`), on a Radeon AI PRO R9700. The decide GPU checks now run on the tile and attention path that `TF_QWEN3_MATRIX` and `TF_QWEN3_ATTENTION` choose (#2262, part of #2245). Thanks to @laudney. +- Spark 2.5 and Dory on CUDA now load with direct reads by default: cold loads on a GB10 go from 5.0 to 2.5 s (Spark X2.6 4B 8-bit) and from 4.5 to 3.8 s (Dory NVFP4), with the same replies. `TF_LOAD_DIRECT=0` keeps the old mapped reads (#2449). + - Gemma 4 and Qwen 3.5 on CUDA now load with direct reads by default, 28-31% faster cold on a GB10 with the same weights and replies. A load uses about 0.6 GB more host memory; `TF_LOAD_DIRECT=0` keeps the old read-ahead (#2282). ## 1.2.3 diff --git a/docs/optimization-ledger.md b/docs/optimization-ledger.md index ddf3cf2873..2687546675 100644 --- a/docs/optimization-ledger.md +++ b/docs/optimization-ledger.md @@ -96,7 +96,7 @@ last two hold the families with only a Metal, HIP or scoring backend. | 19. Shared drafts | no (copies) | no (copies) | no (copies) | open [#1426] | done¹ (opt-in) | | 20. Solo chunks | yes | yes | yes | yes | yes | | 21. Read-ahead | done [#2033] | done [#2033] | yes | done [#2033] | done [#2033] | -| 22. Direct load | yes | yes | yes | yes | open [#1924] | +| 22. Direct load | yes | open [#2449] | open [#2449] | yes | open [#1924] | | 23. Page drop | done [#1884] | done [#1884] | done [#1884] | done [#1884] | done [#1884] | | 24. L2 prefetch | yes | yes | yes | yes | done¹ | | 25. TP ranks | yes | no (4B) | yes | done [#1388] | done¹ | @@ -393,6 +393,8 @@ below; these are estimates for planning a port, not results. | --- | --- | --- | | flashnext | cold load 74 → 48 s | [#1901] | | flashnext | 50 → 43.5 s, next chunk gathered while one packs | [#1960] | + | spark25 (4B 8-bit) | cold load 5.0 → 2.5 s | [#2449] | + | dory (NVFP4) | cold load 4.5 → 3.8 s | [#2449] | 23. **Page-cache drop after load.** Every CUDA loader drops its checkpoint's pages once the weights are up, so the memory goes back to KV and the prompt cache on unified-memory boxes ([#1884]). TP loaders wait for every rank @@ -549,3 +551,4 @@ below; these are estimates for planning a port, not results. [#2283]: https://github.com/ashhart/TensorFold/pull/2283 [#2288]: https://github.com/ashhart/TensorFold/pull/2288 [#2306]: https://github.com/ashhart/TensorFold/pull/2306 +[#2449]: https://github.com/ashhart/TensorFold/pull/2449 diff --git a/zig/src/families/dory/cuda_weights.zig b/zig/src/families/dory/cuda_weights.zig index 734846178d..fd91e47661 100644 --- a/zig/src/families/dory/cuda_weights.zig +++ b/zig/src/families/dory/cuda_weights.zig @@ -121,7 +121,11 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []co const f = c.factors(); w.subs = try gpa.alloc(Sub, c.sublayers); w.embed = try w.upload(d, m.embed.bytes); - for (w.subs, m.subs) |*out, s| out.* = switch (s) { + // each layer read past the page cache, the next one while it packs (TF_LOAD_DIRECT=0: main's mapped reads) + var direct = core.checkpoint.direct.Layers.init(gpa, &ck, try core.checkpoint.loadReaders()); + defer direct.deinit(); + const count = c.sublayers; + for (w.subs, m.subs, 0..) |*out, *s, i| out.* = switch (try direct.moved(s, "backbone.layers.{d}.", i, count)) { .attn => |a| .{ .attn = .{ .alpha = try w.upload(d, a.alpha.bytes), .q = try w.proj(d, a.q), diff --git a/zig/src/families/spark25/weights.zig b/zig/src/families/spark25/weights.zig index a6a6536ecb..2f28673887 100644 --- a/zig/src/families/spark25/weights.zig +++ b/zig/src/families/spark25/weights.zig @@ -168,8 +168,12 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.S var b1: [128]u8 = undefined; var b2: [128]u8 = undefined; const readers = try core.checkpoint_read_ahead.optInReaders(); + // each layer read past the page cache, the next one while it packs (TF_LOAD_DIRECT=0: main's mapped reads) + var direct = core.checkpoint.direct.Layers.init(gpa, &ck, try core.checkpoint.loadReaders()); + defer direct.deinit(); for (w.layers, 0..) |*l, i| { - var reading = try @import("cuda_load_prefetch.zig").begin(&ck, i, readers); + try direct.layerAt(m ++ "layers.{d}.", i, if (i + 1 < n) i + 1 else null); + var reading = try @import("cuda_load_prefetch.zig").begin(&ck, i, if (direct.on()) 0 else readers); defer reading.deinit(); const name = struct { fn of(b: []u8, layer: usize, suffix: []const u8) ![]const u8 {