Prefill OOM: long-context requests fail mid-stream with "Memory limit exceeded during prefill"

[sdiamanEXUS]

Describe the bug

When oMLX is used as the backend for a coding agent (e.g. OpenCode, Claude Code), a single long‑context prefill can spike GPU memory past the hard ceiling and the scheduler aborts the entire request mid‑stream with:

omlx.scheduler - WARNING - Prefill force-stopped at 1441 tokens: memory 42.8GB exceeds ceiling 42.0GB
omlx.scheduler - ERROR   - Prefill failed for <id>: Memory limit exceeded during prefill
omlx.server    - ERROR   - Error during chat streaming: Memory limit exceeded during prefill

The client receives a finish_reason="error" / SSE error event and the agent stops working. Notably, memory drops back to a low baseline immediately after the failure:

omlx.process_memory_enforcer - INFO - Memory pressure level: hard -> ok (current=24.6GB, soft=35.7GB, hard=39.9GB, ceiling=42.0GB)

i.e. the 42.8 GB was a transient prefill activation spike, not resident KV — the request is killed for a spike that is gone microseconds later.

To Reproduce

1. Serve a model whose attention uses head_dim > 128 (MLX takes the float32 SDPA fallback that materializes the full [B, n_q, query_len, kv_len] attention matrix). MoE models are especially affected — measured transient was ~18 MB/token.
2. Drive it with a coding agent so the prompt/context grows large (tens of thousands of tokens), keeping the model resident near the ceiling.
3. Send a turn whose prefill chunk is large. The per‑chunk attention transient overshoots the hard ceiling and the request fails with Memory limit exceeded during prefill.

Observed on a 48 GB Mac with a ~23 GB model (ceiling ~42 GB from iogpu.wired_limit_mb), failing at only ~1.4k processed tokens.

Expected behavior

The server should keep memory bounded and the request should keep working — by sizing prefill chunks to fit the available headroom — rather than hard‑failing a request mid‑stream for a transient spike. If a turn genuinely cannot fit, it should fail gracefully/retriably rather than crashing the stream.

Root cause (analysis)

• Scheduler._adaptive_chunk_size only shrinks the chunk once current memory crosses the soft watermark. But a single large chunk's transient can blow the ceiling from a low baseline (current well below the watermark), so the throttle never engages before the overshoot.
• The per‑scheduler PrefillTransientTracker.predict() (an EWMA of measured bytes/token, whose docstring is literally "size the next chunk so its predicted transient stays under the remaining headroom") is never called.
• The static _preflight_memory_check is wired but under‑estimates for MoE / head_dim>128 models (its SDPA term doesn't capture the real transient), so the request is admitted and then fails.
• On failure the request is terminated; there is no reclaim‑and‑retry. The ProcessMemoryEnforcer "all loaded models are pinned and no loads in progress" branch is a no‑op (it only logs), so nothing is reclaimed between turns.

Desktop (please complete the following information):

• macOS Version: 26.5
• oMLX Version: 0.3.12

Additional context

A fix is proposed in an accompanying pull request: a predictive per‑chunk throttle (sizes each chunk so current + predicted_transient(n) ≤ ceiling, using the measured EWMA with a static first‑chunk fallback), a reclaim‑before‑fail step, a bounded requeue‑then‑retry fallback, and turning the dead‑end enforcer branch into an idle reclaim. Verified live: a 23,313‑token prompt that previously crashed now prefills and completes with no OVER_HARD / force‑stop.