Add MLX_DISABLE_NAX option to skip Metal-4 nax kernels at build time

[1R053]

Summary

Add a MLX_DISABLE_NAX CMake option (default OFF, no behaviour change for existing users) that downstream projects can set to skip the Metal-4 "nax" tensor GEMM / qmm / attention kernels even when the build environment otherwise qualifies for them (MLX_METAL_VERSION >= 400 AND MACOS_SDK_VERSION >= 26.2).

Motivation

On macOS 26 the bundled Metal toolchain miscompiles the nax kernels. Symptom: matmuls with M > ~8 rows silently return wrong results. Short prompts are bit-exact against the reference (numpy / older MLX pip wheel); medium and long prompts diverge at the very first step. We observed this consistently across every architecture in our local rig (qwen3_5, gemma4 dense + MoE + elastic, llama4 Scout, minimax_m2, nemotron_h, mamba2) with metalfe-32023.883 from com.apple.MobileAsset.MetalToolchain-v17.6.42.0 (macOS 26.5 host, SDK 26.5).

A minimal repro using only this repo's headers:

# 3-row probe that brackets the M=8 nax-fragment threshold
import mlx.core as mx
import numpy as np

def check(M, K, N, tag):
    a = np.sin(np.arange(M * K) * 0.013).astype(np.float32).reshape(M, K)
    b = np.cos(np.arange(K * N) * 0.007).astype(np.float32).reshape(K, N)
    mx.eval(c := mx.matmul(mx.array(a), mx.array(b)))
    err = float(np.abs(np.asarray(c) - a @ b).max())
    print(f"  {tag}: M={M:>3}  max abs err = {err:.6f}")
    return err

check( 8, 64, 64, "M<=8  (below nax fragment row-jump)")  # ~1e-4, correct
check(16, 64, 64, "M>8                              ")    # gross divergence on affected toolchains
check(64, 64, 64, "M>8                              ")

Reproduces against any MLX build from current main on macOS 26 with no CMAKE_OSX_DEPLOYMENT_TARGET pin below 26.0. The bug is in the Metal toolchain codegen for subtile_matmad_nax; this PR is not a fix for that — it's an opt-out so users hitting it can continue to consume mlx without producing miscompiled kernels.

Smoking-gun: Apple's own MLX-shipping pipeline already avoids the nax codepath

The mlx.metallib shipped inside Apple's iOS simulator runtimes — at PrivateFrameworks/GESS.framework/mlx.metallib, built and signed by Apple for distribution with Xcode — contains zero nax kernel symbols across every macOS-26-era runtime we checked:

Simulator runtime	metallib mtime	size	strings ... | grep -ciE 'steel_gemm_fused_nax_|_qmm_n_nax_|_qmm_t_nax_|_gather_qmm_n_nax_|_gather_qmm_t_nax_'
iOS 26.4	2026-03-16	2,427,468 B	0
iOS 26.4	2026-04-09	2,427,468 B	0
iOS 26.5	2026-05-07	2,427,468 B	0

A control build from main on the same host with CMAKE_OSX_DEPLOYMENT_TARGET unset yields a libmlx.a with 7 such symbols and the divergent-matmul behaviour above.

Corroborating evidence from mlx-swift: that package's Package.swift defines METAL_PATH = "default.metallib", meaning it consumes a pre-built default.metallib rather than compiling .metal sources during SwiftPM build. We have three cached mlx-swift_Cmlx.bundle/default.metallib files from our own iOS / macOS app archives built between 2026-03-15 and 2026-03-28 — both after #2772 ("Add Neural Accelerator Support", 2025-11-19) and the #3271 nax refactor (2026-03-18). All three contain 0 strict nax kernel matches at ~3.82 MB each. The same nax-disabled outcome shows up in mlx-swift consumers' artifacts as in Apple's own iOS-sim runtime metallibs, even though the two artifacts have different content (different sizes — one is mlx-swift's full kernel set, the other is GESS.framework's subset).

This implies Apple's MLX-shipping pipeline (across both the iOS-sim runtime and the mlx-swift package) applies the equivalent workaround. There appears to be no Metal Toolchain version where these kernels have ever shipped correctly in production — the bug is not a regression from a known-good state. The proposed MLX_DISABLE_NAX flag is, in effect, the explicit form of what Apple's existing pre-built metallib distribution channels are already doing implicitly. Downstream projects that compile MLX from source via cmake (e.g. anyone using the C++ library or a Rust FFI wrapper) currently have no such shield and are the ones hit by the bug.

Why an explicit flag (vs the existing implicit workaround)

The existing implicit workaround (-DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 to make __METAL_VERSION__ report 310 and trip the existing version gate to else()) works today, but:

• It couples the workaround to the deployment-target value, which downstream projects often have their own reasons to set independently.
• It's fragile against the metal vs metallib link-step target leakage that mlx: fix macOS 26 target leakage in v3 metallib ollama/ollama#16053 had to add a separate workaround for — when the final link step stamps the metallib with the SDK's macOS 26 target despite AIRs being compiled at 14.0, runtime dispatch can still look for nax kernels that aren't in the library and fail to load (the symptom in MLX models not working on M5 MacBook Pro lmstudio-ai/lmstudio-bug-tracker#1356).
• It's not searchable in MLX's CMake interface — a new consumer hitting the bug has no obvious flag to find.

An explicit, intent-named flag is friendlier to consumers who hit the bug, easier to remove cleanly when the toolchain is fixed, and discoverable via cmake -LA.

Related reports

• mx.fast.scaled_dot_product_attention produces incorrect attention scores (MLX 0.31.2) #3585 — mx.fast.scaled_dot_product_attention produces incorrect attention scores on M4 / MLX 0.31.2; total model collapse on the fused MPS path, manual attention works. Same failure-mode class as the bug this PR's flag defends against; an explicit MLX_DISABLE_NAX=ON build closes the symptom cleanly.
• MLX models not working on M5 MacBook Pro lmstudio-ai/lmstudio-bug-tracker#1356 — Unable to load function steel_gemm_fused_nax_* on M5 MacBook Pro / macOS 26.2. Different surface (runtime kernel-load failure rather than miscompiled inference output), same kernel family. With MLX_DISABLE_NAX=ON those symbols are never compiled and the load is never attempted.
• mlx: fix macOS 26 target leakage in v3 metallib ollama/ollama#16053 — mlx: fix macOS 26 target leakage in v3 metallib. Demonstrates the failure mode where the final metal vs metallib link step leaks the macOS 26 deployment target into the library and defeats the -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 workaround. MLX_DISABLE_NAX=ON is robust against that leak — the kernels are never compiled, so metadata leakage can't resurrect them.
• Fix nax condition for iphone #3083 — A18 hardware-detection gate added because nax kernels were "incorrectly detected as having Neural Accelerator support, causing garbage output." Establishes that the MLX team is already aware nax codegen has correctness traps on undertested toolchain × hardware combinations and is willing to gate via the build system.

Implementation

• New top-level option MLX_DISABLE_NAX (default OFF) in CMakeLists.txt, sitting next to MLX_METAL_DEBUG.
• Existing gate at mlx/backend/metal/kernels/CMakeLists.txt:158 extended with AND (NOT MLX_DISABLE_NAX).
• The else-branch (which already defines MLX_METAL_NO_NAX) is reached on ON, so the existing C++ dispatcher fallback is reused — no new source code paths, no new C++ surface.

Net: 8 lines added, 1 modified. Reverting the workaround is a one-line revert.

Test plan

• Diff is additive: default OFF preserves the existing version/SDK gate exactly, so default-off behaviour is unchanged by construction.
• On-path verification of the underlying bug (in our local rig on metalfe-32023.883, default-off build with the version/SDK gate qualifying): strings build/.../mlx.metallib | grep -ciE 'steel_gemm_fused_nax_|_qmm_n_nax_|_qmm_t_nax_|_gather_qmm_n_nax_|_gather_qmm_t_nax_' returns 7 nax kernel matches, and oracle_run against qwen3_5-0.8B-4bit produces 1/3 greedy-match (short bit-exact, medium err 21.88 at step 0, long_prefill err 30.16 at step 0).
• -DMLX_DISABLE_NAX=ON end-to-end build verified on this fork at the commit in this PR:
  • cmake configure passed (macOS SDK 26.5, Metal found, deployment target deliberately not set — so the existing version/SDK gate would otherwise qualify, isolating MLX_DISABLE_NAX as the suppression).
  • cmake --build . --target mlx-metallib completes in 40s.
  • Resulting mlx.metallib: 0 strict nax kernel matches (pass) AND 0 broader _nax_ substring matches (pass), with the non-nax steel-GEMM workaround kernels present (steel_gemm_fused_(nn|nt|tn|tt)_*: 384 instantiations) — confirming the flag suppresses nax compilation without dropping the workaround kernels.
• iOS / iOS-simulator cross-compile: untested. Those branches already set CMAKE_OSX_DEPLOYMENT_TARGET independently and have been failing the version/SDK gate (and hitting the else() branch) all along, so this PR is a no-op for them — not expected to regress.

[zcbenz]

The issue has nothing to do with NAX as you get same results with versions without NAX support.