diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index e67be9a06d..0cedcfa1e8 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -563,6 +563,15 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { } } +// The loader indexes the scales in one of three ways, by where a thread's +// n_reads packed words fall inside a quantization group: +// 1. group_size >= BCOLS: a tile row sits in one group, so next() advances +// the scales every group_steps = group_size / BCOLS tiles. +// 2. BCOLS > group_size, read inside one group: a tile row covers +// n_groups = BCOLS / group_size groups, group_id picks one, and each +// tile advances the scales by n_groups. +// 3. The read covers whole groups: the thread walks groups_per_read of +// them, taking a scale and bias for each. template < typename T, short BROWS, @@ -574,11 +583,8 @@ template < short bits> struct QuantizedBlockLoader { static_assert( - BCOLS <= group_size, - "The group size should be larger than the columns"); - static_assert( - group_size % BCOLS == 0, - "The group size should be divisible by the columns"); + BCOLS % group_size == 0 || group_size % BCOLS == 0, + "The tile columns and the group size must divide one another"); static_assert( bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || bits == 8, @@ -589,7 +595,28 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor; MLX_MTL_CONST short n_reads = (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; - MLX_MTL_CONST short group_steps = group_size / BCOLS; + + // Groups covered by one tile row. + MLX_MTL_CONST short n_groups = (BCOLS > group_size) ? BCOLS / group_size : 1; + // Tiles that share one group. + MLX_MTL_CONST short group_steps = + (group_size > BCOLS) ? group_size / BCOLS : 1; + + // Values one thread reads. + MLX_MTL_CONST short n_elems = n_reads * pack_factor; + // Groups per read, and packed reads in each. + MLX_MTL_CONST short groups_per_read = + (n_elems > group_size) ? n_elems / group_size : 1; + MLX_MTL_CONST short reads_per_group = n_reads / groups_per_read; + + static_assert( + BCOLS_PACKED % n_reads == 0, + "The packed columns of the block tile must be a multiple of the " + "reads of one thread."); + static_assert( + n_elems % group_size == 0 || group_size % n_elems == 0, + "The read of one thread must cover full quantization groups, or it " + "must divide one quantization group exactly."); const int src_ld; const int tile_stride; @@ -600,6 +627,9 @@ struct QuantizedBlockLoader { const short bi; const short bj; + // First group this thread reads from. + const short group_id; + threadgroup T* dst; const device uint8_t* src; const device T* scales; @@ -622,23 +652,40 @@ struct QuantizedBlockLoader { thread_idx(simd_group_id * 32 + simd_lane_id), bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), + group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size), - biases(biases_ + bi * src_ld / group_size) {} + scales(scales_ + bi * src_ld / group_size + group_id), + biases(biases_ + bi * src_ld / group_size + group_id) {} + + void dequantize_reads() const thread { + if (groups_per_read == 1) { + T scale = *scales; + T bias = *biases; + for (int i = 0; i < n_reads; i++) { + dequantize( + src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); + } + } else { + for (int g = 0; g < groups_per_read; g++) { + T scale = scales[g]; + T bias = biases[g]; + for (int i = 0; i < reads_per_group; i++) { + const int r = g * reads_per_group + i; + dequantize( + src + r * bytes_per_pack, scale, bias, dst + r * pack_factor); + } + } + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = *scales; - T bias = *biases; - for (int i = 0; i < n_reads; i++) { - dequantize( - src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); - } + dequantize_reads(); } void load_safe(short2 src_tile_dim) const thread { @@ -660,15 +707,7 @@ struct QuantizedBlockLoader { return; } - T scale = *scales; - T bias = *biases; - for (int i = 0; i < n_reads; i++) { - dequantize( - (device uint8_t*)(src + i * bytes_per_pack), - scale, - bias, - dst + i * pack_factor); - } + dequantize_reads(); } void next() thread { @@ -682,8 +721,8 @@ struct QuantizedBlockLoader { biases++; } } else { - scales++; - biases++; + scales += n_groups; + biases += n_groups; } } else { scales += group_stride; @@ -692,146 +731,6 @@ struct QuantizedBlockLoader { } }; -template < - typename T, - short BROWS, - short BCOLS, - short dst_ld, - short reduction_dim, - short tgp_size, - short bits> -struct QuantizedBlockLoader< - T, - BROWS, - BCOLS, - dst_ld, - reduction_dim, - tgp_size, - 32, - bits> { - MLX_MTL_CONST short group_size = 32; - - static_assert( - BCOLS % group_size == 0, - "The group size should be divisible by the columns"); - static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); - - MLX_MTL_CONST short pack_factor = get_pack_factor(); - MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); - MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor; - MLX_MTL_CONST short n_reads = - (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; - MLX_MTL_CONST short n_groups = BCOLS / group_size; - - static_assert( - (BCOLS_PACKED / n_reads) == n_groups, - "Other configurations are not yet supported"); - - const int src_ld; - const int tile_stride; - const int group_stride; - - const short thread_idx; - const short bi; - const short bj; - - const short group_id; - - threadgroup T* dst; - const device uint8_t* src; - const device T* scales; - const device T* biases; - - QuantizedBlockLoader( - const device uint8_t* src_, - const device T* scales_, - const device T* biases_, - const int src_ld_, - threadgroup T* dst_, - ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread - : src_ld(src_ld_), - tile_stride( - reduction_dim ? BCOLS_PACKED* bytes_per_pack - : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS* src_ld / group_size), - thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads* thread_idx / BCOLS_PACKED), - bj((n_reads * thread_idx) % BCOLS_PACKED), - group_id((bj * pack_factor) / group_size), - dst(dst_ + bi * dst_ld + bj * pack_factor), - src(src_ + bi * src_ld * bytes_per_pack / pack_factor + - bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id), - biases(biases_ + bi * src_ld / group_size + group_id) {} - - void load_unsafe() const thread { - if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { - return; - } - - T scale = *scales; - T bias = *biases; - for (int i = 0; i < n_reads; i++) { - dequantize( - src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); - } - } - - void load_safe(short2 src_tile_dim) const thread { - if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { - 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) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - T scale = *scales; - T bias = *biases; - for (int i = 0; i < n_reads; i++) { - dequantize( - (device uint8_t*)(src + i * bytes_per_pack), - scale, - bias, - dst + i * pack_factor); - } - } - - void next() thread { - src += tile_stride; - if (reduction_dim == 1) { - // if (group_steps > 1) { - // group_step_cnt++; - // if (group_step_cnt == group_steps) { - // group_step_cnt = 0; - // scales++; - // biases++; - // } - // } else { - scales += n_groups; - biases += n_groups; - // } - } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; - } - } -}; - template METAL_FUNC void adjust_matrix_offsets( const device T*& x, diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 5a9c9fb874..1262c6cb8c 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -76,9 +76,20 @@ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) +#define instantiate_gather_qmm_rhs_tile(type, group_size, bits, bm, bn, bk, wm, wn) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, bm, bn, bk, wm, wn, true) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, bm, bn, bk, wm, wn, false) + +// Block tiles for the gather kernel, default first. The rest are selected +// with MLX_QMM_TILE_NAX. Keep in sync with gather_qmm_rhs_nax_tiles in +// mlx/backend/metal/quantized.cpp. #define instantiate_quantized_all_rhs(type, group_size, bits) \ - instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \ - instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false) + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 64, 64, 64, 2, 2) \ + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 32, 64, 64, 2, 2) \ + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 16, 64, 64, 1, 2) \ + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 128, 64, 64, 2, 2) \ + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 64, 32, 64, 2, 2) \ + instantiate_gather_qmm_rhs_tile(type, group_size, bits, 32, 32, 64, 2, 2) #define instantiate_quantized_funcs(type, group_size, bits) \ instantiate_quantized_all_batched(type, group_size, bits) \ diff --git a/mlx/backend/metal/kernels/steel/gemm/nax.h b/mlx/backend/metal/kernels/steel/gemm/nax.h index 6d978ebfb6..f211ab3bee 100644 --- a/mlx/backend/metal/kernels/steel/gemm/nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/nax.h @@ -469,11 +469,13 @@ struct BaseNAXFrag { metal::bool_constant, const thread dtype_frag_t& B, metal::bool_constant) { - // Create Matmul descriptor + // Create Matmul descriptor. This overload pairs two fragments along M and + // takes one fragment along N, so the problem shape is 32x16x16. The + // overload above pairs along N, and its problem shape is 16x32x16. constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, + 16, transpose_a, transpose_b, true, @@ -526,6 +528,70 @@ struct BaseNAXFrag { Cm1[i] = ct_c[kElemsPerFrag + i]; } } + + // Pairs along K, for tiles with odd TM and odd TN where a simdgroup holds + // one fragment each way. MPP needs 32 in M, N or K when both inputs are + // cooperative tensors, so 16x16x16 is not expressible and only K is left. + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& C, + const thread dtype_frag_t& Ak0, + const thread dtype_frag_t& Ak1, + metal::bool_constant, + const thread dtype_frag_t& Bk0, + const thread dtype_frag_t& Bk1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 16, + 32, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Ak0[i]; + ct_a[kElemsPerFrag + i] = Ak1[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bk0[i]; + ct_b[kElemsPerFrag + i] = Bk1[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } + } }; template < @@ -891,6 +957,30 @@ METAL_FUNC void tile_matmad_nax( } } } + } else { + // TM and TN both odd, so pair along K. Without this branch the matmul + // did not run at all and the tile kept its cleared value. + static_assert( + TK % 2 == 0, + "MXU tile matmul: a block tile with one fragment along M and one " + "fragment along N needs an even number of K fragments."); + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; kk += 2) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm, kk + 1, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk + 1, nn, tb), + metal::bool_constant{}); + } + } + } } } diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 57692dfcd5..edd6d055cf 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -375,6 +375,60 @@ array quantize_dequantize_input( return xhat; } +struct GatherQmmNaxTile { + int bm; + int bn; + int bk; + int wm; + int wn; +}; + +// Block tiles the affine NAX gather kernel is built for, default first. +// Keep in sync with instantiate_quantized_all_rhs in +// mlx/backend/metal/kernels/quantized_nax.metal. +constexpr GatherQmmNaxTile gather_qmm_rhs_nax_tiles[] = { + {64, 64, 64, 2, 2}, + {32, 64, 64, 2, 2}, + {16, 64, 64, 1, 2}, + {128, 64, 64, 2, 2}, + {64, 32, 64, 2, 2}, + {32, 32, 64, 2, 2}, +}; + +std::string tile_to_string(const GatherQmmNaxTile& tile) { + std::string s; + concatenate( + s, tile.bm, ",", tile.bn, ",", tile.bk, ",", tile.wm, ",", tile.wn); + return s; +} + +// Picks the block tile for the NAX gather kernel. MLX_QMM_TILE_NAX names one +// of the tiles above as "BM,BN,BK,WM,WN". Only the affine kernels are built +// for the extra tiles, so other modes always use the default. +GatherQmmNaxTile gather_qmm_rhs_nax_tile(const std::string& mode) { + if (mode != "affine") { + return gather_qmm_rhs_nax_tiles[0]; + } + auto requested = env::get_var("MLX_QMM_TILE_NAX", ""); + if (requested.empty()) { + return gather_qmm_rhs_nax_tiles[0]; + } + for (const auto& tile : gather_qmm_rhs_nax_tiles) { + if (requested == tile_to_string(tile)) { + return tile; + } + } + std::ostringstream msg; + msg << "[gather_qmm] MLX_QMM_TILE_NAX is set to '" << requested + << "', but MLX did not build this block tile. These block " + "tiles are available:"; + for (const auto& tile : gather_qmm_rhs_nax_tiles) { + msg << " " << tile_to_string(tile); + } + msg << "."; + throw std::invalid_argument(msg.str()); +} + } // namespace void qmv_quad( @@ -1454,9 +1508,12 @@ void gather_qmm_rhs_nax( array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); - // TODO: Tune the block sizes - int bm = 64, bn = 64, bk = 64; - int wm = 2, wn = 2; + // TODO: Choose the block tile automatically from the shape of the problem. + // MLX now uses the default block tile, unless MLX_QMM_TILE_NAX selects a + // different one. + auto tile = gather_qmm_rhs_nax_tile(mode); + int bm = tile.bm, bn = tile.bn, bk = tile.bk; + int wm = tile.wm, wn = tile.wn; const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 7057075c9b..0102836abb 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4804,6 +4804,16 @@ void init_ops(nb::module_& m) { Returns: array: The result of the multiplication of ``x`` with ``w`` after gathering using ``lhs_indices`` and ``rhs_indices``. + + .. note:: + On hardware with matrix coprocessors, a sorted affine call runs a + special kernel. The environment variable ``MLX_QMM_TILE_NAX`` + selects the block tile of this kernel, in the form + ``"BM,BN,BK,WM,WN"``. The best block tile depends on the shape of + the problem and on the distribution of the rows over the matrices. + Because of this, MLX keeps the default block tile. If you set a + block tile that MLX did not build, MLX raises an error that lists + the available block tiles. )pbdoc"); m.def( "gather_qqmm", diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..9b9ff0b0a2 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,6 +1,8 @@ # Copyright © 2023 Apple Inc. +import os import platform +import re import subprocess import unittest from itertools import product @@ -17,6 +19,19 @@ def is_m1_mac(): return cpu.startswith("Apple M1") +def has_nax(): + # The kernels for the matrix coprocessors need an Apple GPU of generation + # 17 or later, and macOS 26.2 or later. + if platform.system() != "Darwin" or not mx.metal.is_available(): + return False + arch = mx.device_info(mx.gpu).get("architecture", "") + match = re.search(r"(\d+)[a-z]$", str(arch)) + if match is None or int(match.group(1)) < 17: + return False + version = tuple(int(v) for v in platform.mac_ver()[0].split(".")[:2]) + return version >= (26, 2) + + class TestQuantized(mlx_tests.MLXTestCase): def test_quantize_dequantize(self): w = mx.random.normal(shape=(128, 512)) @@ -1403,6 +1418,86 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y3, atol=tol)) self.assertTrue(mx.allclose(y1, y4, atol=tol)) + @unittest.skipUnless(has_nax(), "requires a GPU with matrix coprocessors") + def test_gather_qmm_rhs_nax_tiles(self): + # A mixture of experts layer uses the gather kernel, which MLX builds + # with several block tiles, and MLX_QMM_TILE_NAX selects one of them. + # Each block tile must give the same result as a plain gather_mm on + # the dequantized weights. + # + # The last two block tiles give each simdgroup one fragment of 16 + # values along N. These block tiles use matmul paths that the default + # block tile does not use. A group size of 32 makes one row of the + # block tile cover more than one quantization group, and the block + # loader must handle this. + tiles = [ + "64,64,64,2,2", + "32,64,64,2,2", + "16,64,64,1,2", + "128,64,64,2,2", + "64,32,64,2,2", + "32,32,64,2,2", + ] + + E, K, N = 4, 512, 512 + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + + for L, group_size in product([256, 133], [32, 64]): + indices = mx.sort(mx.random.randint(0, E, shape=(L,), key=k1)).astype( + mx.uint32 + ) + x = (mx.random.normal((L, 1, K), key=k2) / K**0.5).astype(mx.float16) + w = (mx.random.normal((E, N, K), key=k3) / K**0.5).astype(mx.float16) + wq, s, b = mx.quantize(w, group_size=group_size, bits=4) + w_hat = mx.dequantize(wq, s, b, group_size=group_size, bits=4) + expected = mx.gather_mm( + x, w_hat.swapaxes(-1, -2), rhs_indices=indices, sorted_indices=True + ) + mx.eval(expected) + + for tile in tiles: + with self.subTest(L=L, group_size=group_size, tile=tile): + os.environ["MLX_QMM_TILE_NAX"] = tile + try: + y = mx.gather_qmm( + x, + wq, + s, + b, + rhs_indices=indices, + transpose=True, + group_size=group_size, + bits=4, + sorted_indices=True, + ) + mx.eval(y) + finally: + del os.environ["MLX_QMM_TILE_NAX"] + self.assertEqual(y.shape, expected.shape) + self.assertLess((y - expected).abs().max().item(), 2e-3) + + # MLX must report an unknown block tile with an error, and it must not + # fail later when it looks for a kernel. + os.environ["MLX_QMM_TILE_NAX"] = "8,8,8,1,1" + try: + with self.assertRaises(ValueError): + mx.eval( + mx.gather_qmm( + x, + wq, + s, + b, + rhs_indices=indices, + transpose=True, + group_size=group_size, + bits=4, + sorted_indices=True, + ) + ) + finally: + del os.environ["MLX_QMM_TILE_NAX"] + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: