From 902e582d2a2c32071e566d04c5cc2454081082a2 Mon Sep 17 00:00:00 2001 From: Isalia20 Date: Wed, 7 Oct 2026 11:41:03 +0400 Subject: [PATCH 1/2] thin n gemm kernels --- mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/jit/includes.h | 1 + mlx/backend/metal/jit_kernels.cpp | 22 ++ mlx/backend/metal/kernels.h | 7 + mlx/backend/metal/kernels/CMakeLists.txt | 4 +- .../steel/gemm/kernels/steel_gemm_thin_nax.h | 223 ++++++++++++++++++ .../gemm/kernels/steel_gemm_thin_nax.metal | 27 +++ mlx/backend/metal/matmul.cpp | 55 +++++ mlx/backend/metal/nojit_kernels.cpp | 9 + 9 files changed, 348 insertions(+), 1 deletion(-) create mode 100644 mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h create mode 100644 mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index 5871f2daf4..cccaad8a5c 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -97,6 +97,7 @@ if(MLX_METAL_JIT) make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_thin_nax) make_jit_source(quantized_nax kernels/quantized_utils.h) make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index 1f1a8c46bf..8ad0a9ca7c 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -54,6 +54,7 @@ const char* steel_gemm_fused_nax(); const char* steel_gemm_gather_nax(); const char* steel_gemm_splitk_nax(); const char* steel_gemm_segmented_nax(); +const char* steel_gemm_thin_nax(); const char* quantized_nax(); const char* fp_quantized_nax(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 6b104f5885..aad94eba1e 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -30,6 +30,9 @@ const char* steel_gemm_splitk_nax() { const char* steel_gemm_segmented_nax() { return ""; } +const char* steel_gemm_thin_nax() { + return ""; +} const char* quantized_nax() { return ""; } @@ -1263,6 +1266,25 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string& template_def) { + auto lib = d.get_library(kernel_name, [&]() { + std::string kernel_source; + concatenate( + kernel_source, + metal::utils(), + metal::gemm_nax(), + metal::steel_gemm_thin_nax(), + template_def); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 41648b4810..ded6c834ab 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -386,6 +386,13 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( int wm, int wn); +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string& template_def); + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 9b72033453..355c763a3a 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -115,7 +115,8 @@ set(STEEL_NAX_HEADERS steel/gemm/kernels/steel_gemm_fused_nax.h steel/gemm/kernels/steel_gemm_gather_nax.h steel/gemm/kernels/steel_gemm_splitk_nax.h - steel/gemm/kernels/steel_gemm_segmented_nax.h) + steel/gemm/kernels/steel_gemm_segmented_nax.h + steel/gemm/kernels/steel_gemm_thin_nax.h) set(STEEL_NAX_ATTN_HEADERS steel/defines.h @@ -172,6 +173,7 @@ if(NOT MLX_METAL_JIT) build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_segmented_nax ${STEEL_NAX_HEADERS}) + build_kernel(steel/gemm/kernels/steel_gemm_thin_nax ${STEEL_NAX_HEADERS}) build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS}) build_kernel(gated_delta_update_nax gated_delta_update_nax.h diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h new file mode 100644 index 0000000000..28229a03cf --- /dev/null +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal new file mode 100644 index 0000000000..e3f87e8dc4 --- /dev/null +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal @@ -0,0 +1,27 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "mlx/backend/metal/kernels/utils.h" + +#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" + +// clang-format off +#define instantiate_gemm_thin(tname, trans_b, iname, itype, um, un, sm, sn, ks) \ + instantiate_kernel( \ + "steel_gemm_thin_nax_" #tname "_" #iname \ + "_um" #um "_un" #un "_sm" #sm "_sn" #sn "_ks" #ks, \ + gemm_thin_nax, itype, um, un, sm, sn, ks, trans_b) + +#define instantiate_gemm_thin_transpose_helper(iname, itype, um, un, sm, sn, ks) \ + instantiate_gemm_thin(nn, false, iname, itype, um, un, sm, sn, ks) \ + instantiate_gemm_thin(nt, true , iname, itype, um, un, sm, sn, ks) + +#define instantiate_gemm_thin_shapes_helper(iname, itype) \ + instantiate_gemm_thin_transpose_helper(iname, itype, 2, 2, 1, 1, 4) \ + instantiate_gemm_thin_transpose_helper(iname, itype, 2, 2, 1, 2, 2) + +instantiate_gemm_thin_shapes_helper(float16, half); +instantiate_gemm_thin_shapes_helper(bfloat16, bfloat); +// clang-format on diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 5fc2c19361..e3bda318c8 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -7,6 +7,7 @@ #include #include "mlx/backend/common/broadcasting.h" +#include "mlx/backend/common/compiled.h" #include "mlx/backend/common/matmul.h" #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/binary.h" @@ -990,6 +991,60 @@ void steel_matmul_axpby( /* float beta = */ beta); } + // Few output columns: rows stay in registers and B is read once + if (use_nax && batch_size_out == 1 && !transpose_a && N <= 64 && + out.dtype() != float32 && (transpose_b || int64_t(K) * ldb <= INT_MAX)) { + int sn = M < 2048 ? 1 : 2, ks = 4 / sn; + std::ostringstream kname; + kname << "steel_gemm_thin_nax_n" << (transpose_b ? 't' : 'n') << "_" + << type_to_name(out) << "_um2_un2_sm1_sn" << sn << "_ks" << ks; + auto base_name = kname.str(); + bool use_out_source = CHECK_AB && (alpha != 0.0f || beta != 1.0f); + bool do_axpby = use_out_source && (alpha != 1.0f || beta != 1.0f); + bool align_M = M % 32 == 0, align_N = N % (32 * sn) == 0; + bool align_K = K % 64 == 0; + kname << "_" << use_out_source << do_axpby << align_M << align_N << align_K; + auto kernel = get_steel_gemm_thin_nax_kernel( + d, + base_name, + kname.str(), + {{&use_out_source, MTL::DataType::DataTypeBool, 100}, + {&do_axpby, MTL::DataType::DataTypeBool, 110}, + {&align_M, MTL::DataType::DataTypeBool, 200}, + {&align_N, MTL::DataType::DataTypeBool, 201}, + {&align_K, MTL::DataType::DataTypeBool, 202}}, + get_template_definition( + base_name, + "gemm_thin_nax", + get_type_string(out.dtype()), + 2, + 2, + 1, + sn, + ks, + transpose_b)); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + int tm = (M + 31) / 32, tn = (N + 32 * sn - 1) / (32 * sn); + int k_iters = ((K + 63) / 64 + ks - 1) / ks; + steel::GEMMParams params{ + M, N, K, lda, ldb, N, tn, tm, 0, 0, 0, 0, k_iters, 0}; + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(b, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(params, 4); + if (use_out_source) { + int ldc = c.strides()[c.ndim() - 2], fdc = c.strides()[c.ndim() - 1]; + steel::GEMMAddMMParams addmm_params{ldc, fdc, 0, alpha, beta}; + compute_encoder.set_input_array(c, 2); + compute_encoder.set_bytes(addmm_params, 5); + } + compute_encoder.dispatch_threadgroups( + MTL::Size(tm * tn, 1, 1), MTL::Size(32 * sn * ks, 1, 1)); + compute_encoder.add_temporaries(std::move(copies)); + return; + } + // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K if (use_nax && batch_size_out == 1 && (K >= 3 * std::max(M, N) || diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index c6a153c5b5..93f8cf311f 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -459,6 +459,15 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gemm_thin_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const std::string&) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, From 23460fd96e754726576c942ae24eb4cc188a7217 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 8 Oct 2026 18:50:36 +0900 Subject: [PATCH 2/2] nit --- mlx/backend/metal/matmul.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index e3bda318c8..7316000af3 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -991,7 +991,7 @@ void steel_matmul_axpby( /* float beta = */ beta); } - // Few output columns: rows stay in registers and B is read once + // Case 2: Few output columns, rows stay in registers and B is read once if (use_nax && batch_size_out == 1 && !transpose_a && N <= 64 && out.dtype() != float32 && (transpose_b || int64_t(K) * ldb <= INT_MAX)) { int sn = M < 2048 ? 1 : 2, ks = 4 / sn; @@ -1045,7 +1045,7 @@ void steel_matmul_axpby( return; } - // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K + // Case 3: Large K with sufficient M, N, and NAX is available, use NAX split-K if (use_nax && batch_size_out == 1 && (K >= 3 * std::max(M, N) || (std::max(M, N) <= 1024 && K > 2 * std::max(M, N)))) {