feat(prefetch): async expert prefetcher v4 -- 94.3% cache hit rate, 2.8x speedup, zero RAM overhead

[EgonRuiter]

Summary

This PR implements an asynchronous expert prefetcher (PILOT modes 0–3) for olmoe.c, significantly accelerating streaming inference directly from SSD storage.

By combining multi-layer lookahead prediction, stale request pruning, consolidated expert I/O, routing EMA (momentum logits), and persistent pre-warmed expert pinning, this achieves a 94.3% overall cache hit rate (up from 62.4% baseline) and a 2.8× speedup (1.26 → 3.54 tok/s), with zero additional RAM overhead (Peak RSS stays at 4.80 GB).

All code has been reviewed and hardened for thread safety, numerical stability, and memory safety before submission.

---

Benchmark Results (OLMoE-1B-7B merged model, cache=32 experts/layer, bits=8)

Version	Description	Cache Hit Rate	Speed	SSD reads/token	RAM (Peak RSS)
Baseline	No Prefetch	62.4%	1.26 tok/s	64	4.80 GB
PILOT=1	1-layer lookahead	79.3%	1.83 tok/s	35	4.80 GB
Consolidated I/O	1 disk read/expert instead of 3	86.9%	2.12 tok/s	22	4.80 GB
Prefetcher v3.0	Deduplication, stale pruning, topk limit	90.1%	2.95 tok/s	16	4.80 GB
Prefetcher v3.1	Routing EMA + dynamic pinning	90.5%	3.25 tok/s	16	4.80 GB
Prefetcher v3.2 (this PR)	Persistent pinning + pre-warmup (PILOT=3)	94.3% 🚀	3.54 tok/s 🚀	16	4.80 GB

The hit rate during the decode phase alone reaches 95.2% under the final configuration.

---

Core Enhancements

1. Consolidated Expert I/O

   • Merges gate_proj, up_proj, and down_proj weights into a single contiguous file (merged_weight) plus a .qs scale file per expert.
   • Reduces disk read syscalls from 3 per expert to 1, maximising sequential SSD throughput.

2. Stale Request Pruning & Queue De-duplication

   • Lock-free ring buffer (pilot_q) decouples prediction from loading.
   • On every new token step the queue is flushed so the background thread focuses on the current token's early layers immediately.

3. Routing Momentum (EMA Logits)

   • Tracks an exponential moving average (EMA) of routing logits across tokens (SMOOTH env var, default 0.3).
   • Stabilises prediction for recurring routing patterns and closes the early-token prediction gap.

4. Persistent Hot Pinning & Pre-warmup

   • Saves dynamic pinning masks to hot_pinned.bin in the snapshot folder after the warmup phase.
   • On subsequent runs the engine pre-loads these experts into the cache at startup, eliminating cold-start and prefill misses without extra RAM.

5. Multi-layer Lookahead (PILOT=2/3)

   • PILOT=2 dispatches prefetch at layer i+2 (after MoE residual is available).
   • PILOT=3 adds an additional i+3 speculative pass to hide even deeper I/O latency.

6. Cross-platform Test Suite & Tooling

   • Python helpers to download the model, convert to merged format, run determinism checks, and compare against the HF PyTorch reference.

---

Implementation Quality

• Thread safety: all mutations of is_queued and the prefetch queue are guarded by g_pilot_mx. The worker thread is detached (pthread_detach) immediately after creation; thread creation failures exit cleanly.
• Numerical stability: candidate logits are shifted by max_logit before expf() to prevent overflow to INF in the softmax-style scoring loop.
• Memory safety: malloc return value is checked in slot_ensure_allocated; hot_eids array bounds are clamped to 256 to prevent stack overflow.
• Wide prefetch: max_cand is now correctly derived from topk * g_wide so the WIDE env var has the intended effect.
• Dead code removed: unused variables, the buggy rebalance_cache() function, and its associated REBAL env var have been removed entirely.

---

Verification

• Compiled and tested on Windows/x86-64 with GCC (-O2 -march=native -fopenmp).
• Determinism confirmed: two consecutive runs produce identical token outputs (Matching tokens: 12/12 vs HF oracle).
• Peak RSS confirmed: stays within the original 4.80 GB budget (no memory regression).
• Hit rate confirmed: 94.3% measured via c/tools/test_olmoe_real.py.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[JustVugg]

First: thank you — and not as a formality. This PR gets right, on its own, several things this repo learned the hard way over the past week:

• The prefetcher changes residency, never values. I read the EMA path carefully: the momentum buffer feeds only pilot_prefetch's ranking, while the real routing keeps using the fresh gate logits. That invariant (a predictor must never alter what gets computed) is exactly what separates a safe optimization from a quality bug here, and you respected it without anyone asking.
• Your headline metric is falsifiable. SSD reads/token: 64 → 10 is a counter that can be checked against reality, not a counter that measures its own intentions. After what happened with ~335 GB I/O saved in EXPERT_BUDGET=N: miss-aware cap on distinct experts/layer (+75% decode tok/s, 6x prefill on low-RAM hosts) #254/EXPERT_BUDGET (merged in #254, now on main): measured empty operating window — needs a quality gate and an MTP-contract decision #303, this matters more than the speedup number itself.
• You shipped an independent oracle (transformers-based, not self-referential) alongside the feature. That's the standard we're trying to make normal.

Housekeeping: I've retargeted the PR to dev — everything lands there first now; main is release-only.

The ask: would you run the same protocol on GLM-5.2?

Here's the honest reason, explained calmly rather than as an objection. Your OLMoE numbers are credible in OLMoE's regime: with cache=32 of 64 experts per layer, the cache covers 50% of the expert pool, so there's real headroom for prefetch to land in, and 62→94% hit is physically plausible. GLM-5.2 — the model this project is named for — lives in a different regime on the low-RAM machines colibrì targets:

• the cache holds 3-4 of 256 experts per layer (~1.5% coverage);
• GLM's routing is famously flat (see the measurements in EXPERT_BUDGET (merged in #254, now on main): measured empty operating window — needs a quality gate and an MTP-contract decision #303 and the LOOKA numbers in [Feature]: Pre-fetching MoE experts #333: ~144 of 256 experts per layer carry 90% of traffic);
• during decode the disk is already saturated (~85-90% of wall time is expert-disk wait), so prefetch there can't create bandwidth — it can only keep the request queue deep and hide latency.

None of that says your design won't help GLM — multi-layer lookahead and the EMA predictor are exactly the kind of thing that pays on a latency-bound queue, and glm.c already has a PILOT ring your ideas could slot into. It says the numbers won't transfer: on GLM the win, if it's there, will look like +15-30% from queue depth, not 94% hit — and a measured +15-30% on the flagship would honestly be a bigger deal for this project than 2.8× on OLMoE.

If you take the GLM run, the protocol that will make it land without a second round:

• TEMP=0 TOKENS=1, one run PILOT=0 and one with your best mode: token streams must be identical (your invariant makes this a formality, but it's the gate);
• cold-vs-cold (AUTOPIN=0, same or deleted .coli_usage, paste the [RAM_GB=...] cap= line) — warm/cold mixes have manufactured fake speedups here before;
• report experts loaded/token + hit rate as primary, wall clock as color (on VHDX-class disks we've measured the same binary varying 3× on wall time between runs).

And if OLMoE-only is the scope you want for this PR, that's completely fine too — say so and we'll review it as an OLMoE improvement on its own merits, with the oracle you already provided. Either way: genuinely good work, and welcome.

[JustVugg]

Merged into dev — thank you, this is careful work and the ablation table is exactly the kind of evidence we like.

One framing note for future readers: this lands in the OLMoE testbed (olmoe.c), and that's where the 62.4%→94.3% / 2.8× numbers live. The production GLM engine (glm.c) already ships the equivalent techniques it was missing there — coalesced single-pread expert slabs, PILOT router-lookahead, persistent pinning — so those gains don't transfer 1:1. The genuinely new ideas here (multi-layer lookahead depth, routing EMA momentum, stale-request pruning) are worth A/B-ing on glm.c next; the testbed now makes that experiment cheap, which is the real value of this PR.