[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.