Fix M5 split-K NAX dispatch cliff

mlx/backend/metal/quantized.cpp

@@ -6,6 +6,7 @@
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/quantized_dispatch.h"
 #include "mlx/backend/metal/reduce.h"
 #include "mlx/backend/metal/unary.h"
 #include "mlx/backend/metal/utils.h"
@@ -1020,6 +1021,17 @@ void gather_qmm_nax(
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
+bool qmm_nax_eligible(
+    const array& x,
+    bool transpose,
+    int K,
+    const std::string& mode) {
+  bool has_nax_kernel =
+      metal::is_nax_available() && (transpose || mode == "affine");
+  return has_nax_kernel && transpose && (K % 64 == 0) &&
+      (env::enable_tf32() || x.dtype() != float32);
+}
+
 void qmm(
     const array& x,
     const array& w,
@@ -1035,10 +1047,7 @@ void qmm(
     metal::Device& d,
     const Stream& s,
     const std::string& mode) {
-  bool has_nax_kernel =
-      metal::is_nax_available() && (transpose || mode == "affine");
-  if (has_nax_kernel && transpose && (K % 64 == 0) &&
-      (env::enable_tf32() || x.dtype() != float32)) {
+  if (qmm_nax_eligible(x, transpose, K, mode)) {
     return qmm_nax(
         /* const array& x = */ x,
         /* const array& w = */ w,
@@ -1129,7 +1138,8 @@ void qmm_splitk(
     int K,
     metal::Device& d,
     const Stream& s,
-    const std::string& mode) {
+    const std::string& mode,
+    bool fallback_qmm_is_nax_eligible) {
   // Choose split_k to target ~512 threadgroups
   int bm = 32, bn = 32;
   int n_tiles = (N + bn - 1) / bn;
@@ -1148,6 +1158,22 @@ void qmm_splitk(
   while (split_k > 1 && (K % (split_k * k_align) != 0)) {
     split_k--;
   }
+  if (metal::qmm_t_splitk_should_use_nax(
+          fallback_qmm_is_nax_eligible,
+          d.get_architecture_gen(),
+          d.get_architecture().back(),
+          /* transpose = */ true,
+          /* single_batch = */ true,
+          /* affine = */ mode == "affine",
+          group_size,
+          bits,
+          m_tiles,
+          N,
+          K,
+          split_k)) {
+    return qmm(
+        x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode);
+  }
   if (split_k <= 1) {
     return qmm(
         x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode);
@@ -1807,8 +1833,22 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     // Use split-K qmm for small M with transposed weights (non-batched only)
     int B = out.size() / M / N;
     if (transpose_ && B == 1) {
+      bool fallback_qmm_is_nax_eligible = qmm_nax_eligible(x, true, K, mode);
       qmm_splitk(
-          x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
+          x,
+          w,
+          scales,
+          biases,
+          out,
+          group_size_,
+          bits_,
+          M,
+          N,
+          K,
+          d,
+          s,
+          mode,
+          fallback_qmm_is_nax_eligible);
       return;
     }
     qmm(x,

mlx/backend/metal/quantized_dispatch.h

@@ -0,0 +1,33 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+namespace mlx::core::metal {
+
+// The #4198 measurements cover the applegpu_g17s M5 Max class's two-way split-K
+// region between these output widths. Keep this deliberately narrow until a
+// wider device and shape sweep establishes a better crossover.
+constexpr int kQmmTSplitKNaxMinN = 6656;
+constexpr int kQmmTSplitKNaxMaxN = 8192;
+
+constexpr bool qmm_t_splitk_should_use_nax(
+    bool fallback_qmm_is_nax_eligible,
+    int architecture_generation,
+    char architecture_size,
+    bool transpose,
+    bool single_batch,
+    bool affine,
+    int group_size,
+    int bits,
+    int m_tiles,
+    int N,
+    int K,
+    int split_k) {
+  return fallback_qmm_is_nax_eligible && architecture_generation == 17 &&
+      architecture_size == 's' && transpose && single_batch && affine &&
+      group_size == 64 && bits == 4 && m_tiles == 1 &&
+      N >= kQmmTSplitKNaxMinN && N <= kQmmTSplitKNaxMaxN && K % 128 == 0 &&
+      split_k == 2;
+}
+
+} // namespace mlx::core::metal

python/tests/test_quantized.py

@@ -354,6 +354,39 @@ def test_qmm_large_dims(self):
                 tol = 1e-3 if dtype == mx.float32 else 1.5e-3
                 self.assertLess((y_q - y_hat).abs().max(), tol)
 
+    def test_qmm_t_splitk_nax_boundaries(self):
+        # #4198: on M5, affine 4-bit gs64 qmm_t may bypass the two-way split-K
+        # path only for the measured one-M-tile N=6656..8192 band. This checks
+        # numerical correctness at both affected endpoints and adjacent tails;
+        # dispatch selection itself is covered by the pure C++ policy test.
+        if not self.is_apple_silicon or mx.default_device() != mx.gpu:
+            self.skipTest("requires an Apple Silicon Metal GPU")
+
+        key = mx.random.key(4198)
+        group_size, bits = 64, 4
+        dtype = mx.bfloat16
+        # K=6656 is the reported decoder-shaped case. The surrounding cheap
+        # cases retain the N/M tails without making the regression too heavy.
+        for M, N, K in (
+            (32, 6655, 128),
+            (32, 6656, 6656),
+            (32, 8192, 128),
+            (32, 8193, 128),
+            (33, 6656, 128),
+        ):
+            with self.subTest(shape=(M, N, K)):
+                key, k1, k2 = mx.random.split(key, 3)
+                x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype)
+                w = (mx.random.normal(shape=(N, K), key=k2) / K**0.5).astype(dtype)
+                w_q, scales, biases = mx.quantize(w, group_size, bits)
+                w_hat = mx.dequantize(w_q, scales, biases, group_size, bits)
+                y_q = mx.quantized_matmul(
+                    x, w_q, scales, biases, True, group_size, bits
+                )
+                # The unchanged split-K tail can differ by one bfloat16
+                # quantum from the dequantized reference.
+                self.assertLess((y_q - x @ w_hat.T).abs().max(), 2e-3)
+
     @unittest.skipIf("CI" in os.environ, "too slow in CI")
     def test_qmm_non_transposed(self):
         # The non-transposed matmul (w is [K, N]) is reachable mainly from the

tests/ops_tests.cpp

@@ -8,10 +8,64 @@
 #include "doctest/doctest.h"
 
 #include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/metal/quantized_dispatch.h"
 #include "mlx/mlx.h"
 
 using namespace mlx::core;
 
+TEST_CASE("test qmm_t split-K NAX policy") {
+  auto use_nax = [](int architecture_generation,
+                    char architecture_size,
+                    bool fallback_qmm_is_nax_eligible,
+                    bool transpose,
+                    bool single_batch,
+                    bool affine,
+                    int group_size,
+                    int bits,
+                    int m_tiles,
+                    int N,
+                    int K,
+                    int split_k) {
+    return metal::qmm_t_splitk_should_use_nax(
+        fallback_qmm_is_nax_eligible,
+        architecture_generation,
+        architecture_size,
+        transpose,
+        single_batch,
+        affine,
+        group_size,
+        bits,
+        m_tiles,
+        N,
+        K,
+        split_k);
+  };
+
+  CHECK(use_nax(17, 's', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK(use_nax(17, 's', true, true, true, true, 64, 4, 1, 8192, 128, 2));
+
+  CHECK_FALSE(use_nax(16, 's', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(18, 's', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 'g', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 'd', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 'p', true, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(
+      use_nax(17, 's', false, true, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(
+      use_nax(17, 's', true, false, true, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(
+      use_nax(17, 's', true, true, false, true, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(
+      use_nax(17, 's', true, true, true, false, 64, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 32, 4, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 8, 1, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 4, 2, 6656, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 4, 1, 6655, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 4, 1, 8193, 128, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 4, 1, 6656, 64, 2));
+  CHECK_FALSE(use_nax(17, 's', true, true, true, true, 64, 4, 1, 6656, 128, 3));
+}
+
 TEST_CASE("test copy") {
   array x(1.0);
   auto y = copy(x);