SIGN IN SIGN UP

fix(cache,mtp): bf16 stores broken by MLX 0.32.1, plus a runtime cost gate

Three related fixes, each live-proven in the app:

1. prefix_cache: MLX 0.32.1 removed the PEP-3118 buffer shim for bfloat16, so
   np.array(bf16) raises "'bfloat16' is not a valid PEP 3118 buffer format
   string" and EVERY prefix store on every model silently failed since the
   upgrade (blocks_on_disk stayed 0; each turn re-prefilled the full history).
   The generic-KV and ZAYA branches were already guarded; the state-tree paths
   (_to_numpy_tree, cumulative np_state) were not. Cast through float32 -
   value-exact for bf16, loaders already restore the original dtype.
   Proven: dots3 store blocks_on_disk 0 -> 1 and the next identical prompt
   restores ("20 block-disk+mixed_swa cached").

2. Runtime cost gate (default ON): acceptance gates cannot catch a request
   whose CYCLE COST explodes while acceptance stays healthy - measured live on
   dots3, a restored prefix kept 60-90% acceptance while MTP decoded at
   11.7-12.2 t/s and plain AR on the SAME restored cache did 35.1. The seed
   forward is timed as a true AR baseline; if MTP's measured ms-per-emitted-
   token exceeds it by 1.25x over 48+ cycles the request falls back. Runtime
   measurement choosing the faster path. Proven live: the slow hit run demoted
   at cycle 48 with "runtime_cost mtp_ms_per_tok=24.7>..." and recovered to
   34.6 t/s.

3. The seed AR baseline is measured at MATERIALIZATION (mx.eval before the
   draft), not graph build - the first version read 2.3ms for a ~30ms forward,
   a baseline that would demote healthy MTP. The added seed eval also flushes
   the restored lazy graph before cycles begin; with it, hit-path MTP now runs
   25.8 t/s (was 12.2) WITHOUT tripping the gate, and cold runs are unchanged
   (32.7, no false demote).

MTP suite 347 green; prefix/block-disk suite 572 green.
J
Jinho Jang committed
f53897259afc44293f8d721bc7d19a6438f58ea7
Parent: 24706cb