fix(metal): gate NAX to gen-18+ (M5/gen-17 miscompute) + regression test

[bri-prism]

What

The flattened Cmlx header Source/Cmlx/include-framework/mlx-backend-metal-device.h (which the SPM build actually compiles) still carries the pre-fix NAX gate:

can_use_nax &= gen >= (arch == 'p' ? 18 : 17);

On M5-class gen-17 GPUs (arch != 'p') that leaves NAX enabled, and the gen-17 NAX steel-gemm/qmm path miscomputes: wrong fp16 GEMM at M>=8, N>=8192 and wrong quantized matmul at M>=64, N>=9216. On device this shows up as garbled or looping generation, and the current workaround is forcing MLX_METAL_GPU_ARCH=applegpu_g15.

This applies the same gen >= 18 gate already merged in the mlx core (PrismML-Eng/mlx#4) to the flattened copy, and adds a regression test.

Why the flattened header and not just the submodule

The build compiles the flattened Cmlx sources, not the Source/Cmlx/mlx submodule directly. That submodule is pinned at d90771c (2026-06-09), which predates mlx#4, so bumping it is the durable fix but needs a full tools/update-mlx.sh regen. This PR is the minimal, immediately-correct edit so the build is NAX-safe on M5 now.

@khosravipasha as part of the 0.31.6 integration, could you bump the mlx submodule past mlx#4 and regen, so this stays in sync? The header edit will otherwise be overwritten on the next regen, but by then the submodule carries the fix.

Test

Tests/MLXTests/NAXGateRegressionTests.swift runs a fp16 matmul at a NAX-threshold shape (M=16, N=16384) and compares it to the fp32 reference by relative Frobenius error. Correct fp16 is about 1e-2; a mis-enabled gen-17 NAX path is O(1e-1) or worse, so the 0.05 bound fails exactly when the gate regresses. It is a no-op on GPUs where NAX is off or correct.

Heads up: I could not run it locally (a fresh mlx-swift 0.31.6 Cmlx build OOMs on my M5 Pro alongside other worktrees), so please let CI or your machine confirm it compiles and passes.

[khosravipasha]

Thanks for the catch.

as part of the 0.31.6 integration, could you bump the mlx submodule past mlx#4 and regen, so this stays in sync? The header edit will otherwise be overwritten on the next regen, but by then the submodule carries the fix.

Not sure about mlx-swift's mlx pin, I think we need it pin it to the same one they did in their release, don't think they updated their mlx pin or maybe ported it incorrectly. Or is it did not pick up the changes?

Heads up: I could not run it locally (a fresh mlx-swift 0.31.6 Cmlx build OOMs on my M5 Pro alongside other worktrees), so please let CI or your machine confirm it compiles and passes.

There should not be an OOM, was it during compilation?

[Copilot]

Pull request overview

This PR fixes Metal NAX availability gating in the flattened Cmlx header that the SwiftPM build actually compiles, ensuring NAX is only enabled on GPU gen-18+ to avoid known gen-17 (M5-class) miscomputations. It also adds a regression test intended to fail specifically if the gen-18 NAX gate regresses and fp16 matmul correctness diverges.

Changes:

• Update is_nax_available() in the flattened Metal backend header to require gen >= 18.
• Add an XCTest regression that compares fp16 matmul against an fp32 reference at a NAX-threshold shape.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
Tests/MLXTests/NAXGateRegressionTests.swift	Adds a regression test to detect incorrect NAX enablement via fp16-vs-fp32 matmul divergence.
Source/Cmlx/include-framework/mlx-backend-metal-device.h	Gates NAX availability to GPU gen-18+ in the flattened header used by SwiftPM builds.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[bri-prism]

Thanks for the quick review and approve.

On the submodule pin: the gen-18 gate is a fork-only change. It is PrismML-Eng/mlx#4 on our mlx fork, not upstream ml-explore/mlx, so pinning our mlx-swift submodule to whatever upstream 0.31.6 used will not carry it. The right target is our own fork: bump Source/Cmlx/mlx from its current d90771c (2026-06-09, which predates #4) up to PrismML-Eng/mlx prism HEAD, then regenerate with tools/update-mlx.sh so the flattened Cmlx picks it up. So to your "did it not pick up the changes" question: correct, the pin is just stale (pre-#4), not a bad port. The header edit in this PR is the interim so the build is NAX-safe now, and it gets superseded cleanly once the submodule is bumped and regenerated.

On the OOM: yes, at compile time, during a from-scratch Cmlx/metal-kernel build. It was in an isolated worktree with a few other worktrees active on the same M5 Pro, so almost certainly local memory pressure on a cold build rather than anything about the branch. An incremental build (or CI) should be fine. That is why I asked you to confirm the test compiles and passes on your side.