From 3a9d8a7c7a79fbeaa5bc74edf2f084c667217682 Mon Sep 17 00:00:00 2001
From: Brian <288398250+bri-prism@users.noreply.github.com>
Date: Mon, 6 Jul 2026 13:11:28 -0700
Subject: [PATCH 1/2] fix(metal): gate NAX to gen-18+ in flattened device
 header (M5/gen-17 miscompute)

The flattened Cmlx copy the SPM build compiles carried the pre-fix gate
(gen >= arch=='p'?18:17), leaving NAX enabled on gen-17 M5-class GPUs where the
steel-gemm/qmm path miscomputes. Apply the same gen>=18 gate as the mlx core
(PrismML-Eng/mlx#4), and add a regression test.

Durable fix: bump the mlx submodule past #4 and regenerate via tools/update-mlx.sh.
---
 .../mlx-backend-metal-device.h                |  7 ++-
 Tests/MLXTests/NAXGateRegressionTests.swift   | 47 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 Tests/MLXTests/NAXGateRegressionTests.swift

diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-device.h b/Source/Cmlx/include-framework/mlx-backend-metal-device.h
index bfb07dea..3e91d116 100644
--- a/Source/Cmlx/include-framework/mlx-backend-metal-device.h
+++ b/Source/Cmlx/include-framework/mlx-backend-metal-device.h
@@ -277,9 +277,12 @@ inline bool is_nax_available() {
       can_use_nax = true;
     }
     auto& d = metal::device(mlx::core::Device::gpu);
-    auto arch = d.get_architecture().back();
     auto gen = d.get_architecture_gen();
-    can_use_nax &= gen >= (arch == 'p' ? 18 : 17);
+    // NAX steel-gemm/qmm miscomputes on M5-class gen-17 GPUs (wrong fp16 GEMM for
+    // M>=8,N>=8192 and quantized matmul for M>=64,N>=9216). Gate NAX to gen-18+ (A19+).
+    // Mirrors PrismML-Eng/mlx#4 in the mlx core; this flattened copy needs the same fix
+    // until the mlx submodule is bumped past #4 and regenerated via tools/update-mlx.sh.
+    can_use_nax &= gen >= 18;
     return can_use_nax;
   };
   static bool is_nax_available_ = _check_nax();
diff --git a/Tests/MLXTests/NAXGateRegressionTests.swift b/Tests/MLXTests/NAXGateRegressionTests.swift
new file mode 100644
index 00000000..1cdef27d
--- /dev/null
+++ b/Tests/MLXTests/NAXGateRegressionTests.swift
@@ -0,0 +1,47 @@
+// Regression guard for the M5-class (gen-17) NAX steel-gemm/qmm miscompute.
+//
+// The prism fork gates NAX to GPU gen-18+ (A19+) because the gen-17 NAX matmul
+// path returns wrong results for fp16 GEMM at M>=8, N>=8192 (and quantized matmul
+// at M>=64, N>=9216). If that gate regresses and NAX is (wrongly) enabled on a
+// gen-17 device, a fp16 matmul crossing the NAX dispatch threshold diverges hard
+// from the fp32 reference. This test fails in exactly that case.
+//
+// It is a no-op elsewhere: on GPUs where NAX is off or correct, the fp16 matmul
+// tracks fp32 within normal half-precision accumulation error.
+
+import Foundation
+import MLX
+import MLXRandom
+import XCTest
+
+class NAXGateRegressionTests: XCTestCase {
+
+    override func setUp() {
+        setDefaultDevice()
+    }
+
+    /// fp16 matmul over a NAX-threshold-crossing shape must match the fp32 result
+    /// within half-precision accumulation error. Broken gen-17 NAX yields
+    /// O(1e-1)+ relative error (vs ~1e-2 correct), so a 0.05 relative-Frobenius
+    /// bound cleanly separates correct from mis-enabled NAX.
+    func testNAXThresholdMatmulMatchesReference() {
+        MLXRandom.seed(0)
+        let m = 16, k = 512, n = 16384  // M>=8, N>=8192 -> NAX-eligible steel-gemm shape
+
+        let a = MLXRandom.normal([m, k])
+        let b = MLXRandom.normal([k, n])
+
+        let ref = a.matmul(b)  // fp32 reference
+        let got = a.asType(.float16).matmul(b.asType(.float16)).asType(.float32)
+
+        let diffFro = ((got - ref) * (got - ref)).sum().sqrt().item(Float.self)
+        let refFro = (ref * ref).sum().sqrt().item(Float.self)
+        let relErr = diffFro / refFro
+
+        XCTAssertLessThan(
+            relErr, 0.05,
+            "fp16 NAX-threshold matmul rel-Frobenius error \(relErr) exceeds 0.05 "
+                + "(~1e-2 expected for correct fp16). NAX is likely mis-enabled on this GPU "
+                + "generation — check the gen-18 gate in is_nax_available().")
+    }
+}

From 947ebf2581555bda112207b3ca43b525651437ce Mon Sep 17 00:00:00 2001
From: Brian <288398250+bri-prism@users.noreply.github.com>
Date: Mon, 6 Jul 2026 14:03:22 -0700
Subject: [PATCH 2/2] test: use class setUp() for default-device config (match
 MLXTests suite convention)

---
 Tests/MLXTests/NAXGateRegressionTests.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tests/MLXTests/NAXGateRegressionTests.swift b/Tests/MLXTests/NAXGateRegressionTests.swift
index 1cdef27d..c8961a9e 100644
--- a/Tests/MLXTests/NAXGateRegressionTests.swift
+++ b/Tests/MLXTests/NAXGateRegressionTests.swift
@@ -16,7 +16,7 @@ import XCTest
 
 class NAXGateRegressionTests: XCTestCase {
 
-    override func setUp() {
+    override class func setUp() {
         setDefaultDevice()
     }