diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index f3aa24bf7a..ab71c37c6c 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -163,9 +163,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -220,14 +225,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -848,9 +846,12 @@ METAL_FUNC void fp_qmm_t_impl( loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -860,7 +861,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -870,9 +871,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -882,7 +892,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -893,6 +903,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -969,13 +988,20 @@ METAL_FUNC void fp_qmm_n_impl( loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -984,9 +1010,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1004,7 +1040,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1013,9 +1053,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1031,8 +1081,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6d87dc770f..6b766b502d 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -589,6 +589,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -644,14 +650,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..6421ef8514 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -715,6 +715,167 @@ def test_fp_qvm(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 2e-3) + FP_TAIL_TOLS = {mx.float32: 1e-3, mx.float16: 5e-3, mx.bfloat16: 4e-2} + + def _fp_tail_dtypes(self): + # The half types only run on the GPU + if mx.default_device() == mx.gpu: + return [mx.float32, mx.float16, mx.bfloat16] + return [mx.float32] + + def test_fp_qmm_non_multiple_of_32(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + + def check(M, K, N, transpose, mode, dtype): + with self.subTest( + M=M, K=K, N=N, transpose=transpose, mode=mode, dtype=dtype + ): + w_shape = (N, K) if transpose else (K, N) + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + + y_q = mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode=mode + ) + # An unbounded store of the N tail spills onto the next output + # row and races the threadgroup that owns it, so two identical + # runs must agree bit for bit. + y_q2 = mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode=mode + ) + y_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat) + + self.assertEqual(y_q.shape, y_hat.shape) + self.assertTrue(mx.array_equal(y_q, y_q2)) + self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype]) + + # (M, K, N, transpose). M is kept clear of the qmv batch limit + # (get_qmv_batch_limit caps at 32) so these route to the tiled kernel. + nvfp4_shapes = [ + (50, 1040, 128, True), # K % 32 == 16 + (33, 528, 64, True), # K % 32 == 16, M just over one BM tile + (50, 512, 1040, False), # N % 32 == 16 + (50, 1040, 1040, False), # both dims carry a tail + (50, 16, 128, True), # K is a single partial tile + ] + control_shapes = [ + (50, 1056, 128, True), # K % 32 == 0 + (50, 512, 1056, False), # N % 32 == 0 + ] + + for dtype in self._fp_tail_dtypes(): + for M, K, N, transpose in nvfp4_shapes: + check(M, K, N, transpose, "nvfp4", dtype) + for mode in ["mxfp4", "mxfp8"]: + for M, K, N, transpose in control_shapes: + check(M, K, N, transpose, mode, dtype) + + def test_fp_qmm_non_multiple_of_32_vjp(self): + # The backward of quantized_matmul runs the same kernels with the + # transpose flipped, so cover both directions at a tail shape. + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + M, K, N = 50, 1040, 128 + + for dtype in self._fp_tail_dtypes(): + for transpose in [True, False]: + with self.subTest(transpose=transpose, dtype=dtype): + w_shape = (N, K) if transpose else (K, N) + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode="nvfp4") + w_hat = mx.dequantize(w_q, scales, mode="nvfp4") + cot = (mx.random.normal(shape=(M, N), key=k3) / N**0.5).astype( + dtype + ) + + def f(x): + return mx.quantized_matmul( + x, w_q, scales, transpose=transpose, mode="nvfp4" + ) + + (out,), (grad,) = mx.vjp(f, (x,), (cot,)) + tol = self.FP_TAIL_TOLS[dtype] + + # The forward runs the tail kernel for this transpose; the + # backward runs it for the other one, we might as well check both. + out_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat) + self.assertEqual(out.shape, out_hat.shape) + self.assertLess((out - out_hat).abs().max(), tol) + + grad_hat = cot @ ( + w_hat if transpose else mx.swapaxes(w_hat, -1, -2) + ) + self.assertEqual(grad.shape, grad_hat.shape) + self.assertLess((grad - grad_hat).abs().max(), tol) + + def test_fp_gather_qmm_non_multiple_of_32(self): + # Both gather_qmm kernels: the sorted-indices (MoE prefill) kernel and + # the tiled gather kernel. + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + mode = "nvfp4" + E = 4 + + def check(K, N, transpose, sorted_indices, dtype): + with self.subTest( + K=K, + N=N, + transpose=transpose, + sorted_indices=sorted_indices, + dtype=dtype, + ): + w_shape = (E, N, K) if transpose else (E, K, N) + w = mx.random.normal(shape=w_shape, key=k2).astype(dtype) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + + if sorted_indices: + T = 64 + x = (mx.random.normal(shape=(T, 1, K), key=k1) / K**0.5).astype( + dtype + ) + rhs = mx.sort(mx.random.randint(0, E, shape=(T,), key=k3)).astype( + mx.uint32 + ) + lhs = None + else: + B, M = 8, 50 + x = (mx.random.normal(shape=(B, M, K), key=k1) / K**0.5).astype( + dtype + ) + rhs = mx.random.randint(0, E, shape=(B,), key=k3).astype(mx.uint32) + lhs = mx.arange(B, dtype=mx.uint32) + + y_q = mx.gather_qmm( + x, + w_q, + scales, + lhs_indices=lhs, + rhs_indices=rhs, + transpose=transpose, + mode=mode, + sorted_indices=sorted_indices, + ) + w_sel = mx.take(w_hat, rhs, axis=0) + y_hat = x @ (mx.swapaxes(w_sel, -1, -2) if transpose else w_sel) + + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype]) + + # (K, N, transpose): transpose puts the tail on K, else on N. + shapes = [ + (1040, 128, True), # K % 32 == 16 + (512, 1040, False), # N % 32 == 16 + ] + for dtype in self._fp_tail_dtypes(): + for K, N, transpose in shapes: + for sorted_indices in [True, False]: + check(K, N, transpose, sorted_indices, dtype) + def test_mode_error_cases(self): w = mx.random.normal(shape=(256, 256)) x = mx.random.normal(shape=(1, 256))