feat(turboquant): add opt-in mid-prefill KV conversion

[Collinw24]

Summary

Completes the opt-in TurboQuant mid-prefill path from #1843 and resolves the long-context validation blocker.

• Adds turboquant_mid_prefill, gated by normal TurboQuant enablement and off by default.
• On sustained physical-memory pressure, gives the existing one-shot idle-LRU eviction path first chance, then converts eligible dense KV layers once per request instead of shrinking straight to rejection.
• Supports external and chunked prefill. Conversion is layer-by-layer in bounded 8,192-token slices; rotating, recurrent, and pass-through caches keep their native form.
• Keeps stream_experts: false inert. Streaming code does not use @mx.compile.
• Exposes the setting in the web admin UI and macOS app without erasing the saved child value when the parent TurboQuant toggle is off.
• Adds a fail-closed, restart-safe validation runner for pinned MRCR quality, throughput, memory, host headroom, and source-tensor integrity.

Safety contract

• mx.eval(router_logits) completes before any expert-slot mutation.
• Patched production resolve() / __call__ hot paths add no dynamic allocation.
• Mid-prefill starts only after the scheduler confirms cache eligibility; unsupported or partial caches fail closed.
• A request gets one conversion attempt and one LRU-eviction retry budget. A partially advanced external prefix is discarded before retry, so the suffix cannot be replayed twice.
• Mid-prefill owns a process-exclusive Metal lane for the engine lifetime. Claim and later access checks include batched, VLM, embedding, reranker, audio, and DFlash fallback engines. Conflicts remain retryable and do not poison model-load state.
• DFlash fallbacks inherit eviction reporting, return route preflight results, claim the same process ownership, and release partial state on claim/start failure.
• Paged SSD writer bounds price the full hybrid boundary payload, including rotating-window and measured recurrent state.
• Validation hashes source .safetensors before and after every run. Darwin reads must successfully set and verify F_NOCACHE through fcntl; there is no cached-read fallback.

Independent review fixes

An independent high-risk review found four blockers; all are fixed and re-reviewed:

1. Hybrid PagedSSD writer queue undercounted rotating/recurrent boundary state.
2. External non-chunked eviction after partial work retained a mutated restored prefix.
3. DFlash fallback omitted eviction/preflight propagation and mid-prefill process ownership.
4. Sole-model ownership did not count non-EngineCore pool engines.

The second review found no P0-P3 issue.

Audited real-model validation

Model: Justanut/Qwen3.6-27B-uncensored-heretic-v2-Native-MTP-Preserved-oQ6e-mtp

Quality matrix

Clean commit c40dc069; 30/30 cells passed across six pinned MRCR rows and five modes: dense, ordinary Q8, ordinary Q4, mid-prefill Q8, and mid-prefill Q4.

• Every mode produced the same greedy response hash and official retrieval score as dense for each row.
• Scores were 1.0 on five rows and 0.999418942475305 on row 311 in every mode.
• Maximum absolute FP32 mean-NLL delta from dense: ordinary Q8 2.064e-6, ordinary Q4 3.144e-6, mid Q8 1.840e-6, mid Q4 3.770e-6.
• Five source shards (23,716,288,094 bytes) matched exactly before and after.

Controlled prefill performance

t/s is total prefill throughput; suffix t/s measures the same post-midpoint suffix boundary in every mode. Peak is supervised child phys_footprint.

Prompt	Mode	t/s	Suffix t/s	Peak GiB
32,045	dense	153.40	145.29	30.79
32,045	ordinary Q8	171.94	163.05	30.80
32,045	ordinary Q4	173.27	163.22	30.80
32,045	mid Q8	156.13	135.68	29.05
32,045	mid Q4	156.48	135.39	28.12
63,118	dense	153.04	136.29	36.32
63,118	ordinary Q8	151.41	134.53	36.33
63,118	ordinary Q4	154.09	138.45	36.32
63,118	mid Q8	131.30	105.87	36.22
63,118	mid Q4	130.34	104.21	30.90

Mid-prefill is a memory-safety option, not a universal speed win. Q4 cut the supervised 64K peak by 5.42 GiB, with a 14.8% total-prefill and 23.5% same-boundary suffix cost. Q8 saved 1.74 GiB at 32K, but only 0.10 GiB at the supervised 64K process peak.

Organic 131,071-token run

Clean commit 0bef73a3; status ok; forced trigger disabled; verified no-cache source reads; source shards matched exactly.

• 131,071 prompt tokens, two prefill attempts, one actual scheduler pause/retry, one natural mid-prefill trigger at 36,864 tokens.
• Converted 15 layers in 75 bounded slices in 1.779s; one dense layer was intentionally skipped.
• 94,206 post-trigger tokens at 73.75 t/s; full prefill 1,720.66s at 76.18 t/s.
• 256-token replay: mean NLL 2.338879, perplexity 10.369607, decode 6.81 t/s.
• MLX peak 29.03 GiB; supervised child peak 36.32 GiB under the 37 GiB limit; minimum host headroom 6.72 GiB over 6,738 samples.

The first final-head repeat stopped fail closed when unrelated host pressure briefly reduced headroom to 5.76 GiB, below the unchanged 6 GiB floor. Source hashes still matched. The unchanged-command retry started after memory_pressure -Q reported 90% free and passed without lowering either safety limit.

Automated checks

• Full non-slow Python suite: 8,041 passed, 67 skipped, 71 deselected (0bef73a3).
• Affected runtime/cache/DFlash/pool suites: 634 passed.
• Selected macOS XCTest settings/UI checks: 22 passed.
• Browser UI check covered parent/child toggle state, disabled-state clarity, and saved child-value preservation.
• Independent final review: correct, no P0-P3 findings.

The PR remains draft because the feature is experimental and the 64K tradeoff is material.

[jundot]

Thanks for taking this on. I reviewed the memory math, conversion path, and tests again. I think the feature has a valid use case: it can let a long-context request finish when dense prefill would otherwise hit the memory guard.

The KV memory reduction is real. With the Qwen geometry used in this PR, the full-attention KV payload at 131,071 tokens is about 8.00 GiB dense, 4.28 GiB with Q8, and 2.40 GiB with Q4. That saves 3.72 GiB or 5.60 GiB respectively. Model weights and the 48 recurrent cache states do not shrink, so these percentages do not apply to the total process footprint.

I also ran a synthetic MLX check on an M3 Ultra. Dense K/V storage fell from 32.0 MiB to 16.1 MiB with Q8 and 8.1 MiB with Q4. In the same check, quantized attention was about 1.85 times slower than dense attention. The relative attention-output difference was about 1.3% for Q8 and 10.0% for Q4. This is not an end-to-end model benchmark, but it confirms that the memory, speed, and precision tradeoff is real.

I found two validation gaps and two runtime issues that should be addressed before another review.

1. Add real quality validation

Mid-prefill conversion changes more than the cache format. After the trigger, the rest of the prompt attends to quantized history, so later hidden states and K/V values can diverge from dense prefill.

The existing equality test uses _AppendModel, which appends fixed tensors without letting attention affect later K/V values. It proves that sliced conversion matches ordinary conversion for the same tensors, but not that a real model preserves quality. Eight coherent output tokens are also only a smoke test. The long runs use Q8 even though the configured default is Q4.

Please compare dense prefill, ordinary post-prefill TurboQuant, and mid-prefill TurboQuant on the same real model and prompts. Q8 and Q4 both need coverage. A token-level KL or perplexity result plus a long-context retrieval test would be enough to show whether the change is safe for the intended workload. The TurboQuant paper and a later vLLM study report different results depending on bit width and workload, so oMLX needs its own measurements.

2. Add a controlled performance comparison

The reported conversion pause of about 1.1 seconds is small. The main cost is the quantized attention path used for the rest of the prefill.

The 131K run completed, but it used a forced trigger, a 64-token chunk, and took 2,753.20 seconds. There is no matching dense or ordinary TurboQuant baseline, so the slowdown caused by mid-prefill is still unknown.

Please add an A/B comparison at 32K and 64K using the same model, prompt, chunk size, and memory limits. Report peak process footprint, total prefill time, prefill throughput after the trigger, and decode throughput. The 131K case should also be repeated with a normal pressure trigger rather than a forced one.

3. Keep exclusivity conflicts retryable

If another engine prevents the process-exclusive claim, _load_engine() cleans up the new engine but then records the conflict as a cached model-load failure. Unloading the other model does not help because the next request is rejected before another load is attempted.

Please handle this as a retryable busy or capability error instead of calling _mark_load_failure(). Add a test that fails to load model B while model A is present, unloads A, and then successfully retries B.

4. Require confirmed eligibility before deferring preflight

Route-level preflight runs before a prompt cache exists. _can_defer_mid_prefill_preflight() currently treats unknown conversion eligibility as eligible. An unsupported cache layout can therefore skip the pre-streaming HTTP 400 and fail only after the streaming response has started.

Please defer preflight only when conversion eligibility is confirmed. Add a streaming test with an unsupported cache layout and verify that the request receives the normal HTTP 400 before the response starts.

The branch also needs a rebase against current main, since the engine, memory, and preflight paths currently conflict.

Once these items are addressed, I can review this as an experimental, default-off path for requests that would otherwise fail under memory pressure. I would keep the PR in draft until then.