core: Spark 2.5 and Dory on CUDA load with direct reads (one checkpoint.direct call each), 17-50 % faster cold

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

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

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,18 +115,54 @@ 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;
             try s.names.append(gpa, k);
             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,12 +229,55 @@ 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 };
     var l = Layers.init(gpa, &ck, 0);
     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);
 }

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),

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 {