[Gajesh2007]

Why

Our M5 Max providers can only get the GPU Neural-Accelerator speedup (Apple's _nax / Metal Performance Primitives path; ~2.5–4× compute-bound prefill/TTFT) if the MLX C++ core actually contains the nax kernels and their fixes. Our submodule was pinned to MLX 0.31.1 (d5213411), which predates the nax MoE work — so on M5 the MoE expert matmul silently fell off the nax path and we captured ~0× of the win.

This branch is upstream MLX 0.32.0 (b410f6c) + our resource-count trim re-applied (aa480bd8), i.e. it brings in all the nax fixes while preserving the Darkbloom allocator patch our provider depends on.

What this includes (was missing on 0.31.1)

• Fix gather_qmm NAX kernel name mismatch ml-explore/mlx#3632 — gather_qmm NAX bk fix (bk 32 → 64): the gather-qmm NAX kernels are instantiated with BK=64 only; the old value made the kernel-name lookup fail → MoE experts fell back off nax. This is the single most important fix for GPT-OSS / Gemma-4 (both MoE).
• Segmented mm nax kernel ml-explore/mlx#3419 — segmented-mm NAX kernel (steel_gemm_segmented_nax.*): was entirely absent on 0.31.1.
• Fix int16 overflow in NAX qmm edge-tile bounds ml-explore/mlx#3631 — int16 overflow in NAX qmm edge-tile bounds.
• Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K ml-explore/mlx#3361 — int16 overflow in SDPA NAX mask indexing for KV > 32K (long-context correctness).
• Speed up NAX split-K by better tuning and routing and fix NAX addmm ml-explore/mlx#3422 — NAX split-K tuning/routing speedup + NAX addmm fix.
• NAX requires setting MACOSX_DEPLOYMENT_TARGET=26.2 ml-explore/mlx#3622 — NAX build gate now also requires CMAKE_OSX_DEPLOYMENT_TARGET ≥ 26.2 (see build note below).

What's preserved

The "Bound Metal buffer COUNT, not just bytes, in MetalAllocator" patch (orig d5213411, re-applied as aa480bd8). Our provider relies on the public getters it adds — MLX.Memory.numResources / resourceLimit are read by EngineCore.swift and CheckpointCapturePipeline.swift to drive count-aware capacity + checkpoint-pressure decisions. Upstream 0.32.0 independently added the num_resources_/resource_limit_ members and a clear_cache()-path reclaim, so the merge keeps upstream's base count logic and adds our public getters + high-water clear() in malloc + MLX_RESOURCE_LIMIT override with no duplicated members.

git diff origin/main..HEAD is exactly the trim (5 files, +105/−1) — nothing else diverges from upstream 0.32.0.

Conflict resolutions

• mlx/backend/metal/allocator.cpp — includes (kept <algorithm>/<cassert>/<cerrno>); malloc reclaim region (kept the autorelease pool + count-aware comment for the high-water clear() we add).
• mlx/backend/no_gpu/allocator.cpp — kept upstream's real set_cache_limit (delegates to common_allocator) and added the trim's get_num_resources/get_resource_limit stubs.

⚠️ Build note (required to actually get nax)

ml-explore#3622 makes the nax kernels compile only when MLX_METAL_VERSION ≥ 400 AND MACOS_SDK_VERSION ≥ 26.2 AND CMAKE_OSX_DEPLOYMENT_TARGET ≥ 26.2. So a nax-enabled mlx.metallib has a minimum OS of macOS 26.2 and will not launch on older macOS. The M5 provider artifact must be built with the macOS 26.2 SDK + deployment target 26.2 (and avoid the Xcode 26.5 NAX-shader miscompile — use 26.4 / a validated toolchain). M4 / pre-26.2 providers keep building the standard (non-nax) artifact, unchanged.

Verification

• git grep for conflict markers: none.
• mlx/version.h = 0.32.0; gather_qmm_nax bk = 64; steel_gemm_segmented_nax.metal present; getters present in memory.h + allocator.h.
• ☐ CI build on macOS 26.2 SDK (nax variant) — required before merge.
• ☐ Runtime verify on an M5 Max: is_nax_available()==true, nax kernels in metallib, gather_qmm_nax dispatches (tracked in DAR-351).

Before / After

flowchart LR
  subgraph Before["Before — submodule @ 0.31.1 + trim"]
    A1[M5 request: prefill] --> B1[gather_qmm bk=32]
    B1 --> C1[NAX kernel-name lookup FAILS]
    C1 --> D1[fallback simdgroup_matrix path]
    D1 --> E1[~0x of nax speedup]
  end
  subgraph After["After — 0.32.0 + trim, built w/ 26.2 SDK"]
    A2[M5 request: prefill] --> B2[gather_qmm bk=64 + segmented nax]
    B2 --> C2[NAX matmul2d / cooperative tensors]
    C2 --> D2[Neural Accelerator]
    D2 --> E2[~2.5-4x prefill/TTFT]
  end

Loading

Pairs with the mlx-swift submodule-bump PR (points Source/Cmlx/mlx at this branch's tip) and DAR-351 (nax telemetry).