From 87a1b58b8434ced5ef8de0651071d8fb4affe1c2 Mon Sep 17 00:00:00 2001 From: gordofreemo Date: Fri, 7 Aug 2026 12:59:45 +0200 Subject: [PATCH] Fix and enable non-transposed NAX qmm qmm() only lets NAX through when transpose is true, so affine_qmm_n_nax never actually runs and nothing in CI touches it. Two things are broken in it. The weight addressing is set up for [N, K]: it offsets by the K_w/K_g row strides and passes K as the leading dim. With transpose=false the operand is [K, N], grouped along N. The loader instantiation is already the same one the generic qmm_n_impl uses, so the offsets just need to match it. Same impl backs affine_gather_qmm_n_nax, so that gets fixed as well. It also ignores M entirely. The bounds line is commented out and the tile load/store are unguarded, but the grid is ceil(M/BM), so partial tiles do get dispatched and anything not a multiple of 64 writes past the end of y. Pulled sgp_sm + dispatch_bool + load_safe/store_safe over from qmm_t_nax_tgp_impl. The split is compile time so the aligned path doesn't change. Used min(int(SM), ...) rather than truncating the distance to short, same as the int16 fix over there. The guard now also needs N % 64 == 0 for the non-transposed case since the weight tile isn't clamped in N, but any group_size >= 64 gives that anyway. Left gather_qmm() gated since I haven't tested that dispatch path. Test covers unaligned M, mainly 33..63 where a whole simdgroup ends up past the end of the matrix, plus batched x and a large-M partial tile. --- mlx/backend/metal/kernels/quantized_nax.h | 84 +++++++++++++---------- mlx/backend/metal/quantized.cpp | 8 ++- python/tests/test_quantized.py | 47 +++++++++++++ 3 files changed, 102 insertions(+), 37 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index e67be9a06d..5918c34336 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1094,7 +1094,6 @@ METAL_FUNC void qmm_n_nax_tgp_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { (void)lid; - (void)M; static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -1115,23 +1114,20 @@ METAL_FUNC void qmm_n_nax_tgp_impl( bits>; // Set the block - const int K_w = K * bytes_per_pack / pack_factor; - const int K_g = K / group_size; const int y_row = tid.y * BM; const int y_col = tid.x * BN; auto wl = (const device uint8_t*)w; + // Here w is [K, N]: packed and group-quantized along N, with row stride N. x += y_row * static_cast(K); - wl += y_col * K_w; - scales += y_col * K_g; - biases += y_col * K_g; + wl += y_col * bytes_per_pack / pack_factor; + scales += y_col / group_size; + biases += y_col / group_size; y += y_row * static_cast(N) + y_col; - // Make the x loader and mma operation - // const short num_els = min(BM, M - y_row); - // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + // Make the weight loader + loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -1144,6 +1140,13 @@ METAL_FUNC void qmm_n_nax_tgp_impl( const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); + // Rows of this simdgroup's SM-row slice that actually exist. This can go + // <= 0 when a whole simdgroup sits past the end of the matrix, which + // load_safe/store_safe handle, as they already do for the transposed + // kernel. N needs no equivalent: the dispatch guard keeps N % BN == 0. + const short sgp_sm = min(int(SM), M - (y_row + tm)); + const bool is_unaligned_sm = (sgp_sm != SM); + const short ldb_tgp = BN_padded; constexpr bool transpose_a = false; @@ -1156,39 +1159,50 @@ METAL_FUNC void qmm_n_nax_tgp_impl( x += tm * K; - for (int k = 0; k < K; k += BK) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_unsafe(); - threadgroup_barrier(mem_flags::mem_threadgroup); + dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_unsafe(); + threadgroup_barrier(mem_flags::mem_threadgroup); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; - volatile int compiler_barrier; + volatile int compiler_barrier; - Atile.load(x + kk1, K); - Btile.template load(Ws + tn + kk1 * ldb_tgp); + if constexpr (kAlignedM.value) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + Btile.template load(Ws + tn + kk1 * ldb_tgp); - (void)compiler_barrier; - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - x += BK; - loader_w.next(); - } + (void)compiler_barrier; + } - // Store results to device memory - threadgroup_barrier(mem_flags::mem_threadgroup); + x += BK; + loader_w.next(); + } + + // Store results to device memory + threadgroup_barrier(mem_flags::mem_threadgroup); - Dtile.store(y + tm * N + tn, N); + if constexpr (kAlignedM.value) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm)); + } + }); } template < diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index c9f6a7e00e..ff2688afba 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1011,8 +1011,12 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { - if (metal::is_nax_available() && transpose && (K % 64 == 0) && - (env::enable_tf32() || x.dtype() != float32)) { + // The non-transposed kernel loads a BK x BN tile of w = [K, N] without + // clamping in N, so it additionally requires N % 64 == 0. Affine + // quantization with group_size >= 64 satisfies that structurally, since + // the groups run along N. + if (metal::is_nax_available() && (transpose || (N % 64 == 0)) && + (K % 64 == 0) && (env::enable_tf32() || x.dtype() != float32)) { return qmm_nax( /* const array& x = */ x, /* const array& w = */ w, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..83c47befc5 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -353,6 +353,53 @@ 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_non_transposed(self): + # The non-transposed matmul (w is [K, N]) is reachable mainly from the + # vjp of a quantized linear layer, so it gets much less coverage than + # the transposed one. Sweep it over transformer-sized K/N and over M + # values that leave a partial M-tile. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 + + def check(M, K, N, group_size, bits, batch=()): + x = mx.random.normal(shape=(*batch, M, K), key=k1) / K**0.5 + w = mx.random.normal(shape=(K, N), key=k2) / K**0.5 + x = x.astype(dtype) + w = w.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, False, group_size, bits) + y_hat = x @ w_hat + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), tol) + + # M sweep. 33..63 is the interesting range: a whole simdgroup of the + # threadgroup's M-tile falls past the end of the matrix. + for M in [1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129]: + for group_size, bits in [(64, 4), (128, 4), (64, 8)]: + with self.subTest(M=M, group_size=group_size, bits=bits): + check(M, 512, 1024, group_size, bits) + + # Transformer-sized K/N, aligned and unaligned M. + for K, N in [(2048, 2048), (512, 2048), (2048, 512), (11008, 2048)]: + for M in [100, 256]: + with self.subTest(shape=(M, K, N)): + check(M, K, N, 64, 4) + + # Batched x, unaligned M. + for batch in [(2,), (2, 3)]: + for M in [33, 250]: + with self.subTest(batch=batch, M=M): + check(M, 512, 1024, 64, 4, batch=batch) + + # M > 2**15 with a partial M-tile, so the per-simdgroup row count is a + # distance that does not fit in an int16. Same failure mode as the one + # test_qmm_large_dims covers for the transposed kernel. + with self.subTest(shape=(33000, 128, 64)): + check(33000, 128, 64, 64, 4) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key)