Learned lookahead prediction + int2 compensation: 5.5x speedup on 3 GB/s NVMe

[deepak282886]

Hi — I forked colibri and added a single file (learned.py) that replaces
the PILOT=1 heuristic with a trained lookahead predictor and adds int2
quantization with learned compensation.

Results on OLMoE-1B-7B (simulated 3.0 GB/s NVMe):

• Baseline: 1.14 tok/s
• System: 6.24 tok/s
• Speedup: 5.47x

Three components:

1. Autoregressive lookahead predictor — 0.72 recall @ L+3 (vs 0.716 for PILOT=1 heuristic)
2. int2 quantization — 8x smaller experts, faster streaming
3. Per-layer linear compensator — recovers 85% of int2 quality loss

Key finding: predictors must be trained on autoregressive generation data,
not fixed text chunks. Training on chunks gave 0.49 recall.
Training on generation data gave 0.72+.

Repo: https://github.com/deepak282886/colibri-learned

Next steps: validate on GLM-5.2, C runtime, target 50 tok/s on TUF 15 + 4060.

Happy to discuss or contribute back.

[JustVugg]

This is a strong direction and the framing is right. A few reactions:

The predictor. Replacing the PILOT=1 heuristic with a trained autoregressive lookahead is exactly the kind of thing that should live behind a flag. Your headline finding — predictors must be trained on autoregressive generation traces, not fixed text chunks (0.49 recall on chunks vs 0.72 on generation data) — is the part I'd want front and center; it's a non-obvious result and it generalizes beyond your specific model. The 0.72 @ L+3 vs 0.716 for the heuristic is a modest recall delta on its own, so the interesting claim is really the speedup, which depends on how the prediction is spent (prefetch depth, eviction policy), not just recall. That's where the design discussion is.

Coupling. The 5.47× bundles three things — predictor, int2, and the per-layer compensator — measured on OLMoE-1B-7B under a simulated 3 GB/s NVMe. To land any of it I'd want them decoupled: (1) the predictor's speedup at fixed int4 quality, and (2) int2+compensation's quality delta at fixed routing, each measured separately. The int2 compensator recovering 85% of the loss is promising but int2 is a real coherence risk (see #225/#168 on how much scale granularity alone costs), so that half needs a quality ablation on a reasoning prompt set, not just tok/s.

Ask: could you (a) open a Discussion in Ideas (#209) with the generation-vs-chunk training result and the prefetch-depth design, and (b) if you want it merged, a clean PR of just the predictor (learned.py behind an env flag, int2 held back as a separate proposal)? Simulated-NVMe numbers are a fine starting point but the merge bar is a measurement on real streaming. Keeping this open — genuinely want to see it developed.