github jundot/omlx v0.6.3rc2
0.6.3rc2

2 hours ago

oMLX 0.6.3rc2

This is the second release candidate for oMLX 0.6.3. It focuses on performance, memory use, and reliability issues found during rc1 testing. Please report any remaining regressions before the final 0.6.3 release.

This release extends experimental Qwen ANE prefill with optional CPU sharing, sharply reduces ANE compilation memory spikes, and improves first-request latency and tuner reliability. It also adds an M2 Ultra-optimized DeepSeek-V4-Flash indexer kernel, corrects hybrid-model KV memory accounting, makes settings writes resilient to concurrent processes, and hardens SSH and CUDA cluster enrollment.

Qwen ANE Prefill: CPU Sharing and Reliability

CPU Sharing requirement: CPU Sharing requires a separately prepared checkpoint clone whose floating-point tensors use the FP16 dtype, such as Qwen3.8-27B-oQ4e-fp16-mtp. This remains an oQ4e quantized checkpoint: packed integer weights are unchanged. On the stock BF16 checkpoint, CPU Sharing stays disabled and prefill continues on the ANE/GPU path.

  • Added optional CPU sharing to Qwen ANE/GPU prefill. Independent slices of MLP gate/up, MLP down projection, and residual GDN qkv work can now run on the CPU alongside the ANE and GPU branches. The in-app tuner jointly calibrates five workload controls and supports both single- and dual-ANE configurations. By @onthehub97 in #2892.

A fully warmed A/B on an M3 Ultra used 4K and 16K prefills followed by 128 generated tokens. Each configuration ran twice after an ane_2048 warm-up, with the second pass reported:

Configuration Model 4K prefill 16K prefill Versus baseline
ANE off Qwen3.8-27B-oQ4e-mtp 458 tok/s 436 tok/s baseline
ANE/GPU split Qwen3.8-27B-oQ4e-mtp 588 tok/s 555 tok/s +28% / +27%
ANE/CPU/GPU split Qwen3.8-27B-oQ4e-fp16-mtp 625 tok/s 579 tok/s +36% / +33%

CPU Sharing added approximately another 4–6% prefill throughput over the ANE/GPU split on this workload. Decode throughput was unchanged, and the measured peak-memory cost was approximately 7 GB.

  • Reduced ANE bank compilation memory use. Weight slices are converted and staged incrementally instead of retaining every layer in FP32 until both banks are built. On Qwen3.8-27B-oQ4e-mtp, the measured enable-time memory spike fell from 35.8 GB to 4.7 GB, while retained MLX buffer-cache memory fell from 32.1 GB to 0.2 GB with unchanged compile time and parity. This addresses the memory pressure reported in #2781.

  • Moved ANE first-evaluation overhead into model loading. Compiled procedures are now warmed immediately after creation, so the first matching user prompt measures inference rather than private-runtime initialization. On an M5 Pro with a 16K prompt, first-request time improved from 46.1–48.3 seconds to 41.1–41.5 seconds, with approximately 2.9 seconds added to model loading. By @beaglemoo in #2898.

  • Made ANE no-ops visible. Ineligible models, failed dispatch installation, and zero-procedure compilation now produce actionable warnings. /api/status also reports whether ANE prefill was attempted, whether it was successfully configured, and the number of MLP, GDN, dual-ANE, and resident procedures for each loaded model. By @finaltv971 in #2904.

  • Hardened the ANE tuner and memory admission path. Tuning now respects each checkpoint's structural GDN floor, avoids recommending GDN when no GDN operations ran, stages candidates without DFlash wrappers, and releases calibration engines before verification reloads. The memory guard now includes fixed ANE I/O surfaces and CPU-sharing allocations, preventing long prefills from repeatedly crossing the hard watermark on smaller-memory systems. CPU-sharing failures also drain committed GPU work before propagating the error.

  • Preserved model profiles while editing ANE settings. The native app now stages ANE controls and tuner recommendations in the working profile until the user explicitly saves or updates it, instead of clearing the active profile association. By @popfido in #2932.

DeepSeek-V4-Flash Performance on M2 Ultra

  • Added an M2 Ultra-tuned MMA indexer score kernel for DeepSeek-V4-Flash. The new simdgroup_matrix path keeps the K tile resident in threadgroup memory and reduces synchronization overhead. On the reported M2 Ultra measurements, the score operation improved by 1.37–1.38x from 32K through 1M context while remaining bit-exact with the existing Steel kernel. Activation is restricted to the validated model and hardware fingerprint; unsupported configurations continue using the existing path. By @nathanodle in #2802.

Memory and Settings Reliability

  • Corrected KV-cache sizing for hybrid models. Memory estimates now use only layers that retain per-token KV state rather than all transformer layers. This removes the 4x overestimate seen on affected Qwen3.5/3.6/3.8 models, preventing adaptive prefill throttling and HTTP 507 rejections caused by phantom KV growth. Rotating-only and fixed-state cache layouts also retain a conservative SSD writer-queue bound. By @saichowdary007 in #2912, addressing #2725.

  • Made settings writes safe across concurrent server processes. Global and per-model settings are written through per-process temporary files, flushed, and atomically replaced instead of allowing two servers to interleave writes. Existing malformed settings.json files are moved to a timestamped .corrupt-* backup and logged clearly before defaults are loaded, preserving evidence for recovery. By @alexferrao in #2926.

Cluster Enrollment and Recovery

  • Improved recovery from SSH pairing failures. Authentication errors now reach the correct pairing guidance, Bonjour .local names and internal FQDNs can resolve to the same local Mac, and managed SSH keys rotate only after an explicit confirmed request. Replacement keys are generated off-path so a failed rotation leaves the working key pair intact. By @ashhart in #2930.

  • Hardened CUDA worker enrollment and activation. One-time join credentials are claimed before slow provisioning, controller key restrictions are updated safely during re-enrollment, and symlinked authorized_keys paths are rejected. CUDA workers rerun bounded NCCL direct-link verification before every activation instead of trusting stale browser state. By @ashhart in #2929.

App and Documentation

  • Reorganized reasoning-effort controls in the native app. Custom mode, value selection, and force-override controls now appear in task order while retaining a readable two-line fallback at narrow widths and accessibility text sizes. By @popfido in #2922.

  • Updated installation documentation for M5 Macs. Apple Silicon requirements now consistently list M1 through M5 across the translated READMEs. By @qdaszx in #2925.

Upgrade Notes

  • This is a release candidate. Please report regressions before the final 0.6.3 release.

  • Qwen ANE/GPU prefill and CPU sharing remain experimental, opt-in, and dependent on private Apple runtime interfaces and the oMLX native custom kernels.

  • CPU sharing requires a separately preprocessed FP16 clone created with tools/clone_mlx_model_fp16.py. The source checkpoint is not modified or dequantized in place.

  • ANE prefill remains an approximate acceleration path because selected weights are requantized to per-output-channel INT8. Use the built-in tuner for the specific Mac, checkpoint, quantization, and prompt shape.

  • The DeepSeek-V4-Flash MMA score kernel activates by default only for the validated M2 Ultra/model pairing. Other hardware and model layouts retain the existing Steel path.

  • If an existing settings.json is malformed, oMLX now moves it aside as settings.json.corrupt-<timestamp> and starts with defaults. Review the log and restore authentication or other settings from the preserved file when applicable.

  • Distributed inference remains experimental and disabled by default.

  • No manual settings migration is otherwise required.

Thanks to @onthehub97, @beaglemoo, @finaltv971, @popfido, @nathanodle, @saichowdary007, @alexferrao, @ashhart, and @qdaszx for their contributions, and to @taozhiyuai, @quochuy, @panuhallfors, @Collinw24, and @waterliu1981 for detailed reports and validation.

Full Changelog: v0.6.3rc1...v0.6.3rc2

Don't miss a new omlx release

NewReleases is sending notifications on new releases.