From 1ac2fbac9235131c6b375b3d72b7442f496f3f39 Mon Sep 17 00:00:00 2001 From: Pierre Lamy Date: Thu, 2 Jul 2026 12:28:22 -0400 Subject: [PATCH] Raise qmv batch limit for large matrices on M5-class GPUs On g17 hardware the qmv_wide kernels stay ahead of the qmm path well past the generic non-'d' fallback limit of 10. Measured on M5 Max (g17s) with affine 4-bit/8-bit weights (group_size 64) at D=8192, O in {1024, 8192, 28672, 250112}: qmv_wide is ~18% faster at M=12 (241.5ms vs 294.6ms whole-model forward on a 70B) and still ahead at M=16. Raises the large-matrix limit to 16 for gen>=17 non-'d' GPUs; small and mid shapes keep the existing fallback values. Mitigates the small-M speculative-decoding verification cost discussed in #3553. --- mlx/backend/metal/quantized.cpp | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 62d48714e0..6305ef9d87 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -84,6 +84,21 @@ inline array ensure_row_contiguous_matrix( inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { auto arch_size = d.get_architecture().back(); auto arch_gen = d.get_architecture_gen(); + // M5-generation GPUs run the qmv_wide kernels efficiently at larger row + // counts than the generic fallback assumes. Measured on M5 Max (g17s), + // 4-bit and 8-bit affine weights (group_size 64), D=8192 with + // O in {1024, 8192, 28672, 250112}: qmv_wide is ~18% faster than the qmm + // path at M=12 and still ahead at M=16. This mostly benefits speculative + // decoding verification (see #3553 for the small-M step it mitigates). + if (arch_gen >= 17 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 18; + } else if (D <= 4096 && O <= 4096) { + return 12; + } else { + return 16; + } + } if (arch_gen == 13 || arch_gen == 14) { switch (arch_size) { case 'd':