oMLX 0.6.3rc3
Thank you for the time and care behind the recent wave of contributions to oMLX. The current PR volume is more than I can review at once as the sole maintainer, so reviews may take longer than usual. I still welcome useful contributions. For now, I plan to review cluster and distributed inference PRs after the final 0.6.3 release.
For stacked work, please keep dependent follow-up PRs as drafts until the parent PR is merged. If you have several open PRs, I may ask which one or two you want me to review first. AI-assisted contributions are welcome, but authors are expected to understand, validate, and maintain the changes they submit.
This release adds fused Qwen ANE MLP/down-projection offload, tuner-calibrated tail padding, and an optional persistent ANE compile cache. It also hardens ANE execution and memory handling, restores Lightning MTP after temporary performance parking, improves prefix and SSD cache integrity, and reports prefill-memory failures correctly across the OpenAI, Anthropic, and Responses APIs.
Qwen ANE Prefill Performance
- Added fused MLP and down-projection offload. Compatible Qwen3.5/3.6/3.8 Q4 models on dual-ANE systems can now split the MLP hidden dimension across both ANEs, an optional CPU branch, and the GPU. Each branch performs its matching down projection before the results are combined, reducing GPU synchronization and intermediate transfers. The hardware-adaptive tuner now calibrates the fused split, CPU share, worker count, and GDN balance using real native dispatches and full-model verification. By @onthehub97 in #2935.
On the reported M3 Ultra measurement with a Qwen3.8 27B checkpoint:
| Configuration | Prompt throughput | Improvement |
|---|---|---|
| GPU-only baseline | 349.4 tok/s | baseline |
| Fused ANE/CPU/GPU prefill | 527.4 tok/s | +50.9% |
Decode remains on the normal GPU path. Results depend on the Mac, checkpoint, quantization, prompt shape, and system load.
- Added tuner-calibrated intermediate tail padding. Residual prompt tails that are large enough to benefit can now run through a fixed-shape ANE tile instead of always returning to the GPU. Only intermediate MLP and GDN activations are zero-padded, and synthetic rows are removed before recurrence or later model stages. Prompt tokens, positions, KV entries, and recurrent state are not extended. By @onthehub97 in #2966.
For a reported 4,095-token prefill with a 2,047-token residual tail:
| Tail path | Prompt throughput | Improvement |
|---|---|---|
| GPU residual | 415.3 tok/s | baseline |
| Padded hybrid tile | 535.8 tok/s | +29.0% |
- Added an optional persistent ANE compile cache. Previously compiled Qwen ANE programs can be reused across fresh oMLX processes through Apple's AOT cache. Reported fresh-process model-load reductions ranged from 53% to 66% across M1 Ultra, M1 Max, and M2 Max systems, with byte-identical output. The feature can be enabled in Advanced Cache settings or with
OMLX_QWEN35_ANE_COMPILE_CACHE=1, and takes effect on the next restart. By @joshuaswarren in #2975.
ANE Reliability and Memory Management
-
Prevented ANE driver stalls from hanging the server indefinitely. Evaluation waits now have a bounded timeout, failed programs are latched, and affected modules permanently fall back to GPU execution for the current model load. The default timeout is 30 seconds and can be configured with
OMLX_ANE_WAIT_TIMEOUT_S. By @beaglemoo in #3002. -
Stopped procedure-bank retry ladders before process memory reaches jetsam territory. ANE compilation now observes total process footprint, allows failed driver allocations time to settle, and falls back to per-layer compilation when retrying would risk exhausting system memory. By @alytaphoenix in #2999.
-
Reduced fused-down compilation memory. The compiler now dequantizes only the columns assigned to ANE and CPU branches instead of materializing the complete down-projection matrix and discarding the GPU suffix. The resulting staged weights remain bit-identical to the previous path. By @alytaphoenix in #3110.
-
Made ANE dispatch exception-safe. Single, dual, and fused dispatch paths now retire command buffers and outstanding tickets when an exception occurs between submission and worker-thread ownership. This prevents a single failed dispatch from leaving the ANE program stuck in an overlapping-evaluation state until restart. Initial fix by @alytaphoenix in #3105, with follow-up coverage for the fused path.
-
Recovered memory from ANE banks during long-context pressure. When ordinary buffer reclamation cannot restore enough prefill headroom, oMLX can release the requesting model's ANE procedure banks and continue serving with GPU prefill.
/api/statusreports the shed state, and ANE prefill is rebuilt on the next model load. By @beaglemoo in #3103. -
Improved tuner behavior on unsupported and mixed configurations. Machines without the private ANE compiler now receive a completed GPU-only recommendation instead of a failed tuning run. SpecPrefill is disabled during dense ANE calibration so sparse prompt compression cannot invalidate the measurement. By @beaglemoo and @monroewilliams in #3067 and #3080.
Prefix and SSD Cache Reliability
-
Added optional hot-cache write-through. KV blocks can now remain in RAM for fast same-process reuse while also being written to SSD for restart durability. Clearing the hot cache flushes dirty blocks instead of silently discarding them. The feature is disabled by default and can be enabled through settings,
--hot-cache-write-through, orOMLX_HOT_CACHE_WRITE_THROUGH. By @jonathan308 in #3015. -
Aligned Qwen GDN cache boundaries with wide prefill. On validated classic-Metal Qwen3.8 configurations, prefix caching no longer splits a 4,096-token model forward into two numerically different 2,048-token forwards. This fixes deterministic coding-agent sessions that repeated completed tool calls or forgot recent progress only when prefix caching was enabled. NAX-capable and smaller-memory systems retain the validated 2,048-token geometry. By @hesara in #3066.
-
Hardened recurrent-state and shared-prefix correctness. Cache storage now avoids attaching end-of-prompt recurrent state to an earlier full block when a partial tail was skipped, restores pre-call GDN state before stock fallback, verifies shared-prefix hashes during acquisition, and rejects blocks whose stored KV lengths disagree with their declared token count. By @alytaphoenix in #3085, #3086, #3087, and #3094.
-
Improved SSD cache crash safety and preload stability. SSD blocks, GDN sidecars, boundary snapshots, and vision feature entries are flushed before atomic rename, with parent directories synchronized afterward. SSD preload also keeps
mx.load()on the caller thread to avoid known MLX worker-thread deadlocks. By @alytaphoenix in #3091 and #3095.
Lightning MTP and Sampling
-
Made Lightning MTP performance parking reversible. A request that temporarily falls back to standard decoding can retry Lightning MTP after a bounded cooldown. Successful probes resume MTP, while repeated failures use exponential backoff up to 4,096 tokens. Existing aligned-batch and late-join safety checks remain in place. #3004.
-
Restored MTP prompt priming after long cached prefixes. The prime-window memory limit now applies to the newly folded prompt span rather than the absolute context offset. Long sessions with a small uncached remainder can therefore retain the high initial acceptance rate of prompt priming without removing the configured memory bound. By @beaglemoo in #3005.
-
Added
repetition_context_sizeto OpenAI completion requests. Clients can extend the lookback window used byrepetition_penalty, allowing penalties to cover loops longer than mlx-lm's 20-token default. The setting is forwarded through local, VLM, distributed, and DFlash engines while remaining omitted when unset for compatibility. By @xunlinkx in #3012.
API and Model Compatibility
-
Reported prefill-memory failures correctly. Fast non-streaming rejections now return the real HTTP error status instead of a successful HTTP 200 with an error body. Streaming OpenAI, Anthropic, and Responses API paths preserve structured memory-guard details instead of flattening them into generic server errors. By @alytaphoenix and @beaglemoo in #2992 and #3060.
-
Fixed Responses API tool-result images for VLMs. Images in
function_call_outputlists now travel through the normal multimodal path instead of having their base64 payload tokenized as prompt text. Text-only engines receive a placeholder, and ordinary JSON tool results retain their previous behavior. Addresses #2989. -
Restored oQ calibration for legacy Hy-MT2 checkpoints. Root-level
rope_thetavalues are normalized into the structured RoPE configuration expected by the vendored Hy3 implementation, without modifying the source checkpoint. By @chenqianhe in #3042. -
Synchronized the Korean README with the current English documentation. Installation requirements, memory flags, tool-call streaming behavior, native-kernel guidance, distributed inference, and built-in update documentation now match the primary README. By @stellarsailor in #2982.
Upgrade Notes
-
This is a release candidate. Please report regressions before the final 0.6.3 release.
-
Qwen ANE prefill, fused MLP/down offload, CPU sharing, and intermediate tail padding remain experimental and depend on private Apple runtime interfaces and oMLX native custom kernels.
-
The fused MLP/down path currently requires compatible Qwen3.5/3.6/3.8 dense Q4 layers with group size 128 and dual ANE. Unsupported models and quantizations retain the existing prefill path.
-
CPU sharing still requires a separately prepared FP16 checkpoint clone. The source quantized checkpoint is not dequantized or modified in place.
-
The ANE compile cache and hot-cache write-through are disabled by default. Enabling the ANE compile cache requires a server restart.
-
Under sustained long-context memory pressure, oMLX may release a loaded model's ANE procedure banks and continue with GPU prefill until the model is reloaded.
-
On affected classic-Metal Qwen hosts with at least 64 GiB of memory, cache geometry may change from 2,048 to 4,096 tokens. Existing 2,048-token SSD chains remain safely ignored until normal eviction removes them.
-
repetition_context_sizeis optional. Requests that omit it retain mlx-lm's existing default behavior. -
Distributed inference remains experimental and disabled by default.
-
No manual settings migration is otherwise required.
Thanks to @onthehub97, @joshuaswarren, @beaglemoo, @alytaphoenix, @xunlinkx, @jonathan308, @chenqianhe, @hesara, @monroewilliams, and @stellarsailor for their contributions.
Full Changelog: v0.6.3rc2...v0.6.3rc3