[jschmied]

Closes #1873. Flash-Next's CUDA cold load on a GB10 goes from about 74 s to 48 s (−35 %), and the served replies stay identical.

A/B (GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, --parallel 1 --context 16384 --prompt-cache-gib 0; every shard dropped from the page cache before each start; order main, branch, branch, main; the branch was then rebased onto main 58375fe00):

	start to ready	disk read	CPU user / sys
main	75.1 s, 73.1 s	76.2 GB	46 s / 27–30 s
this branch	48.0 s, 48.0 s	74.4 GB	36 s / 18 s

Four greedy replies per run hash the same in all four runs. TF_LOAD_DIRECT=0 (the mapped path, with the store fix below) takes 58 s.

What the time was. Direct reads alone were slower than main (76 vs 70 s), because main's MADV_WILLNEED read-ahead overlaps with its views loop. Phase timers then found two other costs:

• Reopened shards: the weight store kept only two released shards open, so the expert loader reopened shards 685 times and parsed each header again (about 27 ms each, 18.6 s in all).
• Serial page faults: nvfp4_slice.dimensions runs csf.Plane.init, which read about 7 GB of CSF scale planes through the mapping on the main thread before any pack thread started.

Changes

• core/direct_gather.zig (new): it reads the bytes of views a loader borrowed from its mapped files again with O_DIRECT, on a few threads in 8 MiB pieces. Ranges that sit within 1 MiB of each other are joined. The views are then pointed at the copies, and the bytes stay the same. It uses direct_io, so a file system without O_DIRECT falls back to the page cache. TF_LOAD_DIRECT=0 turns it off. Any loader that hands out mapped views can use it.
• core/nvfp4_slice.zig: shapes() makes every metadata check dimensions() makes, without reading plane bytes. dimensions() is now shapes() plus the CSF Plane.init.
• flashnext/cuda_prompt_wt.zig: a load keeps every shard it released open (with its pages dropped) until loaded(), instead of the last two.
• flashnext/cuda_modelopt_threads.zig: with direct reads, a chunk's codes, scales, fixed parts, exceptions and global scales are gathered before it is packed. The CSF planes are still checked by Plane.init during packing, on the gathered bytes.

Tests: host tests for direct_gather cover views across pieces, a 1-byte view, an unaligned tail, an empty view and a refused unmapped view. Existing tests that now run both ways:

• the 36-shard expert loader test packs mapped and direct, and must give the same words;
• the allocation-failure test fails every allocation point in both modes;
• the store test keeps a third shard open and expects nothing held after loaded().

Left out: the gather (about 405 ms a layer) doesn't overlap with packing yet; about 22 s of the load is outside the experts; the store's use() path still reopens shards (116 opens, 4.1 s).

Checks (GB10, on main 58375fe00): lean 0 new, zig build test and the CI set pass, git diff --check clean.

[ashhart]

Thanks @jschmied, this is a nice result: the same device bytes and replies, with a cold start from about 74 s to 48 s. We've queued it on our GB10 with the FP4-Local checkpoint, main against this branch, and will post the numbers. Three questions from reading it:

1. cuda_prompt_wt.zig's release() now drops pages from the cache on every release, for every caller of the store, not only the direct loader. Could that slow a load that reads the same shard again later (the n-gram table, PLE, or a second pass)? If so, keeping the drop to the direct path's own views would avoid it.
2. direct_gather.zig fails the whole load with TooManySources when a layer's views span more than 64 shards. Could it fall back to the mapping for that layer instead, as it does when a file system refuses O_DIRECT?
3. Views from earlier chunks point into gathered memory that's reused for the next chunk. That's safe today because later code only reads .global, but a short comment there, or clearing the old views, would keep the next change from tripping over it.

[ashhart]

A closer read of point 1, @jschmied. release() now drops pages at each borrower's expiry for every caller of the store, and that bypasses the exemption #1884 added. There, dropPageCache keeps the n-gram shards' pages (runtimePages), because the n-gram table is read from them while serving. The bytes stay correct, but those rows go cold, so lookups during decode reread from disk. The MTP head and any use() reload reread too. Could release() keep #1884's exemption, or drop only the views the direct gather replaced?

On the other two: more than 64 held shards in one gather fails the load with no fallback, and the earlier chunks' views point into reused memory but nothing reads them after packing. A guard or a comment would cover the second. The equality is shown by your mapped-against-direct packed-word test on a 36-shard fixture and four equal reply hashes. That's good, and a real-checkpoint byte comparison would make it airtight. Our GB10 A/B will show whether decode moves.

[jschmied]

Thanks, all three are right. Fixed in 750a556ee, a fast-forward on the reviewed head:

1. release() keeps CUDA: every loader drops its checkpoint from the page cache after load (part of #1868) #1884's exemption. It now drops pages only for shards without n-gram rows, the same runtimePages check Shard.close uses. The n-gram table stays as warm as on main. GB10, cold page cache, order main, before the fix, fixed, fixed, before, main:

   	start to ready	n-gram shards resident at ready / after 4 replies	replies
   main 58375fe00	74.1, 74.0 s	2.97 / 3.03 GiB	same 4 hashes
   6eff834ab (before)	49.0, 48.0 s	0.01 / 0.08 GiB	same
   750a556ee (fixed)	50.0, 50.0 s	2.97 / 3.03 GiB	same

   So the cost of keeping those pages is about 1.5 s of load. Shards without n-gram rows still have their pages dropped at release. On main they were dropped later, when the two-slot ring evicted them.

2. More shards than one gather takes. If store.sources() returns TooManySources, that layer reads through the mapping: its globals are parsed from the mapped bytes and its chunks are not gathered. The new host test "a direct layer whose shards don't fit one gather loads through the mapping, the same words" uses a store that always reports too many shards, on plain and CSF layers. With the fallback removed, it fails with TooManySources.

3. Views of packed chunks. After a chunk is uploaded, its views are cleared to their global alone, so nothing can read the reused gather memory through them.

Checks on 750a556ee: lean 0 new, zig build test and the CI set pass, git diff --check clean, test-flashnext-loader-host 25/25.

For the real-checkpoint byte comparison, I'll load every expert layer of the FP4-Local checkpoint mapped and direct on the test driver, compare the packed bytes, and post the result here.

[jschmied]

Here's the real-checkpoint byte comparison. A scratch test (not in the PR) runs loadMode on all 48 expert layers of nvidia/Qwen3.8-Flash-Next-FP4-Local (512 experts, 2560 × 640) at 750a556ee, once mapped and once direct. It uses the test driver, whose device memory is host memory, and hashes each layer's four device buffers with SHA-256: packed gate/up (943,718,400 bytes), packed down (471,859,200), and the two global-scale buffers.

All 48 layers give identical hashes in both modes. Layer 0, for example: up ab0b580199ef…, down e040d5d92b66…, globals 5358ec0b7391… / c19df8eb1516…. Layer 47: 279e528ef55d…, 85cb877d1ff9…, 86236a48c392…, 163cf031a623….

[ashhart]

Thanks @jschmied, the fixes are exactly right. Our GB10 run on nvidia/Qwen3.8-Flash-Next-FP4-Local (main against 750a556ee, warm-up then base, new, new, base): start to ready 105 s and 100 s on main, 55 s and 55 s on this branch, with prompt tokens identical in all 20 rows. Decode is level to slightly up at every prompt size (2K 41.3 -> 42.2, 64K 38.5 -> 39.3 tok/s), and the 8-stream sweeps are flat, so keeping the n-gram pages costs nothing while serving. It's in our gate now and lands right after #1739.

[ashhart]

Thank you @jschmied. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

[ashhart]

Thanks @jschmied! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2