MLX native restored-cache primitive — systemic fix for the Mac throughput collapse

[FluffyAIcode]

⛔ FORBIDDEN FOR ARCHITECTURE VALIDATION — product/throughput bypass only (2026-06-13 directive)

This native-cache path's recall comes from Gemma-4's native retained full-attention layers + native sliding-window eviction — it never exercises f_θ or proposer K/V restoration (ADR 0008 §11). It is therefore structurally incapable of failing in a way that tests the restoration architecture, so its "recall" results must not be cited as evidence that the K/V-Restoration architecture works (see the ADR 0012 revision 2026-06-13).

Do not use this primitive (or the Step-1 incremental path) for any architecture-validation attempt. The restoration architecture must be validated on a pure sliding-window model (Qwen3, the K1/K2 path) where recall is mathematically impossible without proposer/f_θ restoration. Gemma-4 may be used as a product model, never as the validation vehicle for architectural integrity.

Status of this PR: kept open as a labelled product bypass; the memory/throughput numbers are real, but it is not architecture evidence. (Close it in the GitHub UI if you'd prefer it withdrawn entirely — I can't close PRs from the agent.)

---

Why

Per the architecture review, the first MLX port made decode native but the cache lifecycle wasn't — and the naive Step-0 was unusable. Root cause (now confirmed against the authoritative mlx_lm 0.31.3 source): build_native_prefill_cache ran the whole ~5k-token NIAH prompt through one forward → O(T_q×T_k) SDPA materialization → Mac OOM (the same failure as the earlier ctx280 runs). mlx_lm's own generate_step avoids this by chunked prefill (prefill_step_size, generate.py ~430-443).

This PR makes the whole cache lifecycle native and fixes the OOM.

What

inference_engine/backends/mlx/native_restored_cache.py:

• native_generate — pure end-to-end path: hands the whole prompt to mlx_lm's own generate_step (chunked prefill + async decode + optional kv_bits) over a native cache. mlx_lm's validated loop verbatim → cannot OOM on prefill. The default Step-0 path (no quantization).
• build_native_prefill_cache — chunked native prefill (mirrors generate_step) → (cache, last_logits); full-attn KVCache (exact own K/V → S5 recall), sliding RotatingKVCache (bounded). For the selective-quantization path.
• quantize_full_attn_layers — full-attn KVCache → native QuantizedKVCache (real memory reduction, native quantized decode); sliding untouched.
• set_kv_cache_state / inject_restored_into_native_cache — write K/V directly into native layout via .state.
• native_restored_decode, cache_resident_bytes.

Verified against real source: multimodal gemma4.Model.make_cache delegates to language_model (KVCache + RotatingKVCache) → make_prompt_cache types are correct.

Key reframe (now qualified by the banner above): recall rides S5's exact full-attention K/V (native prefill gives it free) → no f_θ/drafter/patch/bridge in the loop. This is the collapse fix, not architecture validation.

Files

• inference_engine/backends/mlx/native_restored_cache.py
• scripts/research/k3_integrated_niah_eval_mac.py — --native-cache, --quantize-full-attn-bits, --prefill-step-size
• scripts/review_mlx_port_on_mac.sh — Step 0
• tests/backends/mlx/test_native_restored_cache.py
• docs/mlx-port-lessons.md

Validation

• ✅ Linux: compiles; native_restored_cache.py 100% line-covered; 58 MLX tests pass.
• ⚠️ MLX execution needs Apple Silicon. Note: any Mac "recall" from this path is Gemma-4 native behaviour, not architecture evidence (see banner).