[gilby]

What this changes

Two parts of the HIP Flash Next host path for the n-gram tables: the gather splits a chunk's ids across up to eight threads (shard located by division when the shard sizes are a remainder run, search otherwise), and after upload the checkpoint's weight pages are dropped from the mapping and page cache except the n-gram spans a gather still reads (Checkpoint.releaseExcept). Same bytes. Fixes #1606. Happy to split the two parts into separate pull requests if preferred.

Receipt

• Environment: TensorFold 7c068895 plus this change, tensorfold-native 1.1.8 built with zig build native-hip -Dhip-arch=gfx1151, Ubuntu archive HIP 7.1 (clang 21), AMD Ryzen AI MAX+ 395 with Radeon 8060S gfx1151 (128 GB unified memory). Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6. --parallel 1 --context 131072 --prompt-cache-gib 0, base and candidate on the same machine.
• Exactness: tools/bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial every cell alone 3/3 and 12/12 equal, serial equal, 0 failed; token_sha256 of bench_openai, bench_concurrent and both bench_cache prompts identical to the base build.
• Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s code / chat: base 50.03 / 43.93, this change 49.90 / 43.86.
• Prompt speed, tools/bench_cache.py --require-cold --tokens 16 on public 7,377- and 29,364-token arrays, first cold request ttft_s: 7,377 tokens 48.44 -> 28.55 s, 29,364 tokens 108.97 -> 105.80 s. Later cold requests in the same process are unchanged (22.45 / 22.33 -> 22.52 / 22.34 s at 7,377): the first-request gain is the load no longer leaving about 29 GiB of weight pages in the page cache for the first prompt to reclaim (page cache after load 31.5 GiB -> 2.8 GiB, device usage unchanged). Host gather alone on synthetic 29.8 GiB shards: 4,096 ids 1.75x, 16,384 ids 2.2-2.5x, 65,536 ids 2.27x faster.
• Tests run: zig build test -Dhip=true --summary all 165/167 passed, 2 skipped, 0 failed (identical on the base); zig build hip-kernel-test 15/15 steps; new tests: shard location over every id of uneven shards, gather equality against the previous row-at-a-time gather (single and multi-threaded, repeated and boundary ids, refusals), span coverage, and releaseExcept keeping its spans readable; mutating the shard location makes them fail. python3 tools/zig/lean_check.py 425 problems on base and change (inherited); git diff --check clean.
• Not run: Metal and CUDA (HIP Flash Next host path; release gains a scratch allocator argument, updated at its one other caller in radio); other AMD chips; overlapping the next chunk's gather with the current chunk (needs a second staging region and graph/arena ordering changes; left out).

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. If the bits change, the description says which sums change and why.
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (No kernel change.)
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks — the gather rewrite reads well, and the byte-for-byte tests against the old search are exactly what we want. Two things before this can leave draft: could you add a short description with before/after resident memory and gather (or prefill) time on a real checkpoint? And #1615 contains this whole diff plus the MADV_RANDOM advice — is 1615 meant to replace this one, or to stack on it? Either is fine, we just want to land one path.

[gilby]

Thanks. #1615 is meant to stack on this one, not replace it: it contains this commit plus one more (the MADV_RANDOM advice). Once this lands I'll rebase #1615 so it carries only its own commit; until then this is the one to review.

The receipt in the description has the real-checkpoint numbers (TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP, gfx1151, same machine for base and change). In short:

• Resident memory after load (read 5 s after the server reports loaded): page cache 31.5 GiB on the base vs 2.8 GiB with this change, device usage unchanged (83.09 GB GTT both). MemAvailable barely moves because the cache was already counted as reclaimable; the change stops the first prompt from paying that reclaim.
• First cold prefill (tools/bench_cache.py --require-cold): 7,377 tokens 48.44 -> 28.55 s, 29,364 tokens 108.97 -> 105.80 s; later cold requests unchanged.
• Gather alone on synthetic 29.8 GiB shards: 4,096 ids 1.75x, 16,384 ids 2.2-2.5x, 65,536 ids 2.27x faster.

I also pushed a zig fmt fix for the failing check.

[ashhart]

Thanks @gilby. We reviewed this independently: every ask from our last review is answered, by commit and receipt, and other GPUs and families are unchanged. We have no AMD GPU, so your gfx1151 receipts carry the device side. It lands now. We marked it ready, since GitHub won't merge a draft. It's gated and queued, and we've written the changelog line.

[gilby]

Yeah, I shouldn't have the strix halo either and should have bought another dgx spark instead. Hopefully we can help catch it up to what its real potential should be. It's where my models go to die or as a back up node while testing on the rest of my fleet.

My agent is finishing up this and collecting all the receipts.

This post by a real human.

[gilby]

Thanks @ashhart. #1611 conflicted with main in engine.zig, where #1651's scratch fallback and the page release sit side by side: both stay, in that order (the release first, then the scratch counted for the chunk). It is rebased on 089a18df (0a6d74f7d is the original change).

It also needed a fix that main's newer CUDA code forced: File.release takes an allocator here, and three callers that arrived after I wrote it, the Flash Next CUDA prompt store (Shard.close, Store.release) and the LFM2 CUDA loader, called the old zero-argument form, so the CUDA host build did not compile. 3df310759 passes them the allocator their file was opened with, and updates the two host tests that pin those calls by their text. I have no CUDA GPU, so that commit is covered by the host build and tests only.

On the rebased head, gfx1151: zig build test -Dhip=true --summary all 507/507 steps, 0 failed, the same as main's 507/507; python3 tools/zig/lean_check.py 424 as on main. Fresh servers in the order main, this, this, main, TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, --parallel 1 --context 131072 --prompt-cache-gib 0: token_sha256 identical to main across bench_openai, bench_concurrent and both bench_cache --require-cold prompts, bench_concurrent --alone --serial every cell equal; the first cold request after start 40.8 / 41.7 -> 23.5 / 23.5 s at 7,377 tokens and 91.9 / 91.2 -> 83.8 / 83.9 s at 29,364, steady cold 18.5 / 18.2 -> 18.2 / 18.3 s and 83.0 / 81.9 -> 81.6 / 81.7 s, decode median tok/s code / chat 51.3 / 45.2 and 51.3 / 45.2 for this change against main's 50.6 / 44.6 and 51.3 / 45.2 (level, within the run spread), MTP acceptance 2.753 / 2.37 in all four.

[gilby]

A correction to my last comment: I wrote that I have no CUDA GPU, which is wrong for this project, and the CUDA commit (3df310759) has now had a compile check against real CUDA tooling. On an x86_64 Linux machine with zig 0.17.0 and nvcc 12.8.93, compile-only (no GPU, kernel or model step was run), upstream main 089a18df and this branch's head give the same result on every step:

step	main	this branch
check-native-cuda-host	7/7 steps succeeded	7/7 steps succeeded
test-cuda-host	287/287 steps succeeded; 96/96 tests passed	287/287 steps succeeded; 96/96 tests passed
compile-cuda-tests	10/10 steps succeeded	10/10 steps succeeded

That covers the two CUDA callers of File.release and the host tests that pin them. It is a compile and host-test check, not a run on a GPU.

[ashhart]

Thanks @gilby. Rebasing this on current main hit a semantic conflict: this PR gives safetensors.File.release() a scratch allocator (for releaseExcept), and since then main gained three callers of the old zero-argument form (flashnext/cuda_prompt_wt.zig twice, lfm2/weights.zig once). I've added one commit passing the scratch at those three (05d69abf7; with an empty keep list it's never touched), and the Linux CUDA build is clean again. Pushed to your branch so CI runs; it's in the landing queue.

[ashhart]

Thank you @gilby. 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]

Landed as 2b8d2aa. Thanks @gilby!

[ashhart]

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