Use larger tiles for qmm_t at large M

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<bits, 8>();
   constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
 
@@ -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<T, group_size, bits, aligned_N, BM, BK, BN>(
+  qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN>(
       w,
       scales,
       biases,

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

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);