diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index f659e16c93..5024fceee3 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/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& 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, diff --git a/mlx/backend/metal/quantized_dispatch.h b/mlx/backend/metal/quantized_dispatch.h new file mode 100644 index 0000000000..47e2c4d992 --- /dev/null +++ b/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 diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 15bc892bd8..e837533d18 100644 --- a/python/tests/test_quantized.py +++ b/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 diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index f7a2b8ab92..38bac1a742 100644 --- a/tests/ops_tests.cpp +++ b/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);