diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6d87dc770f..c6ac98ed1e 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -644,7 +644,31 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { + // `bi` indexes BROWS. When reduction_dim == 1, K is BCOLS (the + // reduction axis) and BROWS is the *output* axis (M for the x loader, + // N for the w loader in qmm_t_impl/affine_gather_qmm_rhs's transpose + // path) -- so the valid-count to guard `bi` against is the caller's + // BROWS bound. Every call site already passes that bound as + // src_tile_dim.y by convention (qmm_t_impl: `short2(BK, num_outs)`; + // affine_gather_qmm_rhs: `short2(k_remain, tgp_bn)`), but this line + // compared `bi` against src_tile_dim.x instead -- the BCOLS/K count, + // not the BROWS/output count. Comparing against the wrong field means + // this early-return under-fires: whenever the true BROWS bound (.y) is + // smaller than the BCOLS one (.x), some `bi` in [.y, .x) fall through + // to the real dequantize-and-load below and read live weight memory + // past the valid output boundary, instead of being zero-filled here. + // In every kernel instantiation shipped to date BROWS == BK == 32, so + // .x (== BK == BROWS) is never smaller than the true bound and the + // bug was unreachable; the resulting rows -- always >= the true bound + // either way -- are outside `num_outs`/`tgp_bn` and get discarded by + // the caller's store_result_safe regardless of what garbage they read, + // so output was correct by luck, not by this check doing anything. + // Fixing the compared field to .y makes the guard actually fire (and + // is a strict no-op for every existing BROWS==BK instantiation, since + // .x == BROWS there too -- the two fields were interchangeable exactly + // in that one case). It matters once BROWS is parametrized + // independently of BK, e.g. qmm_t_impl's new large-M tile below. + if (reduction_dim == 1 && bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1189,7 +1213,9 @@ template < const bool aligned_N, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + const int WM = 2, + const int WN = 2> METAL_FUNC void qmm_t_impl( const device uint32_t* w, const device T* scales, @@ -1211,8 +1237,6 @@ METAL_FUNC void qmm_t_impl( (void)lid; - constexpr int WM = 2; - constexpr int WN = 2; constexpr int pack_factor = get_pack_factor(); constexpr int bytes_per_pack = get_bytes_per_pack(); @@ -1890,7 +1914,9 @@ template < const bool batched, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + const int WM = 2, + const int WN = 2> [[kernel]] void affine_qmm_t( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1937,7 +1963,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 069482cbaf..27e9e788fd 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -42,6 +42,25 @@ aligned, \ batched) +// Like instantiate_quantized_aligned_batched, but also pins BM/BK/BN (the +// qmm_t tile size), for kernels that need a non-default tile ahead-of-time +// compiled. Name must match quantized.cpp's qmm() kname exactly ("_bmX_bnY" +// suffix, no separate "_bk" component since BK is always 32 here). +#define instantiate_quantized_aligned_batched_tile( \ + name, type, group_size, bits, aligned, batched, bm, bk, bn) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned \ + "_batch_" #batched "_bm" #bm "_bn" #bn, \ + name, \ + type, \ + group_size, \ + bits, \ + aligned, \ + batched, \ + bm, \ + bk, \ + bn) + #define instantiate_quantized_quad(name, type, group_size, bits, D, batched) \ instantiate_kernel( \ #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \ @@ -179,4 +198,36 @@ instantiate_quantized_groups(6) \ instantiate_quantized_groups(8) -instantiate_quantized_all() // clang-format on +instantiate_quantized_all() + +// Large-M tile (BM=128, BN=64, BK=32) for affine_qmm_t, ahead-of-time +// compiled so the default (MLX_METAL_JIT=OFF) build -- what `pip install +// mlx` actually ships -- has this kernel available without falling back to +// runtime JIT compilation. See qmm()'s large-M dispatch branch in +// quantized.cpp: it only takes this path for the (group_size, bits) +// combos instantiated here, so a shape that lands outside this set at +// large M safely falls through to the original 32x32x32 tile instead of +// hitting a missing kernel. Scoped to bf16/fp16 x {gs=32,64} x {bits=4,8} +// -- the combo this fork's actual workload (evidence/2026-08-09-mlx-fork/) +// uses -- rather than instantiated across the full type/group_size/bits +// matrix above, to keep this addition's compile-time and binary-size cost +// bounded until a wider need is demonstrated. Naming must exactly match +// the kname built in quantized.cpp's qmm(). +#define instantiate_quantized_large_m_tile(type, group_size, bits) \ + instantiate_quantized_aligned_batched_tile( \ + affine_qmm_t, type, group_size, bits, true, 1, 128, 32, 64) \ + instantiate_quantized_aligned_batched_tile( \ + affine_qmm_t, type, group_size, bits, true, 0, 128, 32, 64) \ + instantiate_quantized_aligned_batched_tile( \ + affine_qmm_t, type, group_size, bits, false, 1, 128, 32, 64) \ + instantiate_quantized_aligned_batched_tile( \ + affine_qmm_t, type, group_size, bits, false, 0, 128, 32, 64) + +#define instantiate_quantized_large_m_tile_types(group_size, bits) \ + instantiate_quantized_large_m_tile(float16_t, group_size, bits) \ + instantiate_quantized_large_m_tile(bfloat16_t, group_size, bits) + +instantiate_quantized_large_m_tile_types(32, 4) +instantiate_quantized_large_m_tile_types(32, 8) +instantiate_quantized_large_m_tile_types(64, 4) +instantiate_quantized_large_m_tile_types(64, 8) // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 7c461bc1b5..e879e48160 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1048,16 +1048,73 @@ void qmm( int B = out.size() / M / N; + // qmm_t's tile is fixed at 32x32x32 regardless of M/N/K/device, which is + // sized for small-M (decode) calls. At large M (prefill-shaped calls) + // this dispatches far more threadgroups than the dense steel GEMM path + // takes at the same shape, and each one re-reads a larger share of both + // operand matrices from device memory -- see + // evidence/2026-08-09-mlx-fork/ in the minimax-h3-mlx repo for the + // measured gap, dispatch trace, and tile-size sweep that motivated this + // branch and this specific 128x64 choice (64x64 got partway there; + // 128x64 landed closest to the dense baseline of the sizes tried; + // 128x128 was a large regression, likely too few threadgroups per core + // to hide memory latency once M/BM * N/BN drops far enough. See + // fix-report.md's config table for the full sweep). + // + // affine_qmm_t / qmm_t_impl (kernels/quantized.h) already accept + // BM/BK/BN (and, after this change, WM/WN) as template parameters with + // defaults matching the old fixed 32x32x32/WM=2/WN=2 config -- this + // dispatch layer simply never varied them. This branch opts large-M, + // transpose=true calls into a 128x64 output tile (BK stays 32: + // QuantizedBlockLoader requires BK <= group_size, and group_size can be + // as small as 32, so BK can't grow without also constraining + // group_size). wm/wn are left at their existing defaults (2, 2) -- + // varying them (tried wm=1,wn=2, mirroring dense's simdgroup config) made + // no measurable difference at this tile size, so there was no reason to + // deviate from the existing default. Every other call -- small M, + // transpose=false, or M below the threshold -- takes the exact same path + // and kernel as before. + // + // kLargeMTileThreshold is intentionally a compile-time constant (not a + // runtime knob) to keep this a pure, auditable tuning branch; bump it + // temporarily to force the old 32x32 tile at a given M for A/B + // benchmarking. + // + // The (type, group_size, bits) guard below matters because this same + // qmm() is linked into both the JIT build (jit_kernels.cpp, compiles any + // template instantiation on demand at runtime -- would work for any + // combo unguarded) and the default, non-JIT build (nojit_kernels.cpp, + // `pip install mlx`'s actual configuration -- looks up a *fixed* set of + // ahead-of-time-compiled kernel names and throws if the name it computes + // isn't in that set). The 128x64 tile is only ahead-of-time instantiated + // for {float16_t, bfloat16_t} x {group_size 32, 64} x {bits 4, 8} (see + // quantized.metal's instantiate_quantized_large_m_tile_types calls) -- + // this fork's actual workload (evidence/2026-08-09-mlx-fork/) -- so this + // guard keeps behavior identical across both builds: every combo outside + // this scoped set falls through to the original 32x32x32 path at every M, + // whether or not the running build happens to have JIT available to + // synthesize it on demand. Widening this list is a matter of adding the + // corresponding instantiate_quantized_large_m_tile_types(...) call. + constexpr int kLargeMTileThreshold = 4096; + constexpr int kLargeMBlockDimM = 128; + constexpr int kLargeMBlockDimN = 64; + bool large_m_tile_dtype_supported = + x.dtype() == float16 || x.dtype() == bfloat16; + bool large_m_tile_quant_supported = + (group_size == 32 || group_size == 64) && (bits == 4 || bits == 8); + bool large_m_tile = transpose && M >= kLargeMTileThreshold && + large_m_tile_dtype_supported && large_m_tile_quant_supported; + int wm = 2; int wn = 2; - int bm = 32; - int bn = 32; + int bm = large_m_tile ? kLargeMBlockDimM : 32; + int bn = large_m_tile ? kLargeMBlockDimN : 32; MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); std::string kname; kname.reserve(64); - bool aligned = N % 32 == 0; + bool aligned = N % bn == 0; bool batched = B > 1; std::string type_string = get_type_string(x.dtype()); concatenate( @@ -1069,20 +1126,38 @@ void qmm( "_b_", bits, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + large_m_tile ? "_bm" + std::to_string(bm) + "_bn" + std::to_string(bn) + : ""); std::string template_def; MTL::ComputePipelineState* kernel; if (transpose) { - kernel = get_quantized_kernel_wrapped( - d, - kname, - "qmm_t", - mode, - type_string, - group_size, - bits, - aligned, - batched); + if (large_m_tile) { + kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + /* BK = */ 32, + bn); + } else { + kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched); + } } else { kernel = get_quantized_kernel_wrapped( d, kname, "qmm_n", mode, type_string, group_size, bits, batched);