Changed direction. Instead of trying to predict multiple tokens into the future, predicting 8 layers into the future instead. Less room for improvement, but much more readily attainable on limited training capacity and data. Now at +16% TG throughput JIT Loads 4016 -> 1621 (-59.7%) Physical loads 4016 -> 2227 The new model is also applicable to PP, not yet wired up or benchmarked. We should expect greater relative throughput as we have the full prompt ahead of processing. We now have an 83% apparent waste overhead that training should be able to reduce to near-zero. Depth specialised heads and budgets could narrow budget where confidence is higher. Latent state prediction could bring load-level confidence a full 8 layers ahead of time, currently most d=8 predictions are hold confidence.

First notable net-positive results from speculative expert streaming experiment (0.8 hold confidence / 0.9 load confidence): JIT Loads: 4546 Baseline -> 3111 w/ speculative streaming (31.6% reduction) AOT loads: 239, 140-141 used (~59% acceptance) Combined load/hold route coverage: 61.23% TG Throughput +4.48% versus baseline Not much, but it's a positive start. I think tweaking the load confidence might yield slightly better results, but further training of the speculative router should yield several orders of magnitude higher TG throughput. I haven't even started training on the 4090 yet, and the M3 Ultra's training rounds are still lasting <1 minute each. Don't ask why I'm training on an M3 Ultra when I have an Ada card. I don't know why.

Could someone with 128+GB memory on an M1-M4 preferably Ultra chip please test this out, tune and benchmark? It brings ANE offload support to Qwen4 architecture. I was able to achieve 33.8% speedup in GDN projection versus GPU alone. github.com/jundot/omlx/pull/… Loading Qwen3.8-flash on 96GB is just too unstable to benchmark.

I don't know who I was calling a masochist when I was benchmarking on a 0.7GiB/s USB 3 SSD. Expect significantly better, but still bad, performance.

For the low-memory masochists who keep asking for SSD expert streaming in oMLX: github.com/jundot/omlx/pull/… Supports pinning override (SoftREAP), although whether that's a useful feature is yet to be determined. Huge thanks to DwarfStar - a lot of the optimisations here are based on that project. Now, to turn this cursed feature off and never use it again...

Very marginal improvement found by offloading attention to the ANE. Offloading GDN output had negative results but included in the PR as experimental anyway, I think in some lower-end chips it may yield a positive result. github.com/jundot/omlx/pull/…

M5 support is ready (in theory). Could someone with an M5 Pro or M5 Max test? Fused approach only supports Q4 at the moment.

5.1% relative speedup from yesterday's 45.8% isn't the bump I was looking for, but I'll take it. Non-M5 users, give it a try github.com/jundot/omlx/pull/… M5 support WIP

Initially I'd discontinued efforts to offload down-projection to the ANE despite the significant performance gains to be found there (potentially 16.6%) due to accuracy issues. It seems those accuracy issues may now be caused by procedure index. Iterating now, another huge bump may be coming.

Wrapping up on this, we're at a net 45.8% performance increase against GPU alone on M3 Ultra at a staggering 517.9 PP tok/s using just 3 CPU cores + ANE. Now featuring a new and improved tuner which takes just a couple of minutes instead of 10+ - which now balances 5 dimensions in a fraction of the time the old one balanced 2, at a much higher arbitrary precision level. Remaining work is to add options to the tuner, so users can opt-out of specific offloads if they don't want them, then will mark PR as ready for review. Feel free to test it out.