perf(streaming): WILLNEED readahead + in-place staged stack for edge0…

[linux295650-ux]

…-35b

A memory-starved host (16 GiB RAM, 19.7 GB 35b checkpoint) spends each decode step in ~3000 MAJOR faults, i.e. the step degrades into "cold pages x per-fault latency" (217 ms / 3079 = 70 us), while the bandwidth floor is only 26 ms (49 MiB @ 1.88 GB/s). Those demand faults serialize on the VM map lock, so moving expert loads to a background pool does not by itself buy anything.

Changes:

• mmap.SafetensorsMmap.advise_willneed_range(offset, length): page-aligned MADV_WILLNEED over ONE byte range. The existing advise_willneed() only covered a whole shard.
• layer.StreamingMoELayer.warm_experts_willneed(experts): one hint per (expert, tensor) byte range, issued before the build in prefetch() and stage_experts(). Retains no MLX arrays, so it does not displace the page cache. (warm_pages() is pre-existing dead code and demand-faults synchronously via sl.sum(); left untouched.)
• LayerOptions.warm_willneed, default False.
• staged_k4(): enable the existing incr_stack + incr_writeback switches together with warm_willneed. incr_stack replaces the per-step core.stack with persistent in-place tensors and incr_writeback keeps the staged bundles out of the shared LRU, removing duplicate copies of the staged set.

Measured on edge0-35b (16 GiB M2, 60-step sampled REPLAY with 20 untimed warmup steps, cache_slots=64, 5 rounds, rotated arm order, paired per round; replay feeds both arms the SAME token sequence so the work is identical):

                    step        tok/s

new (incr+willneed) 187.58 ms 5.33
shipped default 244.25 ms 4.09
prerouter=None 264.06 ms 3.79

new vs shipped : -56.8 ms median, 4/5 paired -> +31.0%
prerouter gain : +8.0% -> +37.8% (5/5 paired)

Output is unchanged: new == old token-for-token on the consistency control, and 60/60 identical tokens in the replay A/B. Also CPU 143 -> 109 ms/step, peak MLX 2.92 -> 2.60 GiB, minor faults 5162 -> 1925, staged wall 162 -> 27 ms/step, critical-path expert load 105 -> 84 ms/step.

Caveats:

• All numbers come from one 16 GiB M2. incr_writeback keeps staged bundles out of the shared LRU, which is free here (LRU hits ~0 at 64 slots) but needs revalidation on a larger-memory host where the LRU actually hits.
• incr_stack is gated on staged_sync (see _incr_mode); the switches disable themselves silently if staged_sync is turned off.
• Pre-existing and unrelated to this change: scripts/e2e_smoke.py's staged-vs-exact gate already fails on this checkpoint with the OLD default too (old vs exact MISMATCH; new vs old is identical).

[ahrazzle]

Follow-up after merge: PR #102

Scope

Independent A/B measurement taken on head 62dcab2, which has since been merged as c18eb62. This ran on a 64 GiB M1 Max with the same checkpoint. Different protocol: one fixed prompt, 128 greedy tokens. Flags controlled from outside the preset.

Shared host: a local model router was present but served no request during the sweep. Arms were interleaved, so page-cache warmup cannot favour one arm.

Results

regime	arm	ms/step	tok/s
cache_slots=64 (the preset's own size)	main, before the merge	80.23	12.47
cache_slots=64 (the preset's own size)	head 62dcab2 (merged)	64.30	15.55
cache_slots=2048 (LRU hits 75-78%)	main, before the merge	51.12	19.56
cache_slots=2048 (LRU hits 75-78%)	head 62dcab2 (merged)	55.23	18.11

At the preset's own cache size the merged head is about 20% fewer ms/step and about 25% higher tok/s. At the larger size it is about 8% slower. There the two arms emit different token sequences, and that regime is reachable only by overriding cache_slots from outside the preset.

With all three flags off, the head's code reproduces the shipped default (52.02 vs 51.12 ms/step).

Follow-up

The three flags are now forced inside staged_k4 in shipped main. Making them opt-in would let a larger-LRU deployment avoid the regression above.

• D-1: staged_k4 cannot override its newly hardcoded flags.
• D-2: advise_willneed_range clamps only the end of the range.

Nothing tests the new advisory and incremental paths.

Not verified

Output identity reproduced at cache_slots=64 (128/128 tokens). The author's caveat that the numbers needed revalidation on a larger-memory host is now answered for the preset's own cache size. Only unmeasured: the PR's own 16 GiB M2 figures and its sampled-replay protocol.

Automated posting by agentic team with human oversight.