Raise qmv batch limit for large matrices on M5-class GPUs

[pierre427]

Proposed changes

get_qmv_batch_limit currently sends gen-17 (M5-class) GPUs down the generic
non-'d' fallback, which caps the qmv/qmv_wide path at M<10 for large
matrices. Measured on an M5 Max (applegpu_g17s), the qmv_wide kernels stay
ahead of the qmm path well past that limit on large shapes. This PR raises
the large-matrix limit to 16 for arch_gen >= 17 non-'d' devices; small and
mid shapes keep the existing fallback values (no data to justify changing
them).

This mitigates the small-M quantized-matmul cost step discussed in #3553 for
the M=11–16 band, which matters for speculative decoding: verification of
multi-candidate / tree / long prompt-lookup proposals lands exactly in that
window. #2031 notes these limits are empirical and machine-dependent; gen-17
hardware postdates the tunings there.

Measurements

Apple M5 Max, 128GB (applegpu_g17s), macOS 26. Whole-model forward cost vs
row count M against a warm 400-token KV cache, min of 6 reps. Model: 70B dense
llama-arch (LLM360 K2-V2), affine 4-bit, group_size 64 (D=8192; O spans 1024 /
8192 / 28672 / 250112 across layers).

M	main @ e9463bb (limit 10)	this PR (limit 16)
1	76.4 ms	75.9 ms
3	85.0	84.8
5	108.2	112.7
8	167.9	170.8
12	294.6 (qmm)	241.5 (qmv_wide, −18%)
16	—	296.0 (qmv_wide)

Same model at affine 8-bit: M=12 326.7 → 304.1 (−7%).

For reference, the 0.31.2 release (no qmv_wide kernels) costs 298.5 ms at
M=12 on the same 4-bit model, and end-to-end speculative decoding on this
hardware improves from ~18 tok/s (release) to 24.5 tok/s (main + k=3 chain
with a 0.6B draft), with the M=11–16 window enabling wider verification
schemes this change dispatches to the faster path.

Repro

import time
import mlx.core as mx
from mlx_lm import load
from mlx_lm.models.cache import make_prompt_cache

model, _ = load("<4-bit 70B model>")
caches = make_prompt_cache(model)
model(mx.array([[i % 1000 + 100 for i in range(400)]]), cache=caches)
mx.eval([c.state[0] for c in caches])
for n in (1, 3, 5, 8, 12, 16):
    ids = mx.array([[200 + i for i in range(n)]])
    ts = []
    for _ in range(6):
        t0 = time.perf_counter()
        mx.eval(model(ids, cache=caches))
        ts.append(time.perf_counter() - t0)
        for c in caches:
            c.trim(n)
    print(f"M={n}: {min(ts)*1e3:.1f} ms")

Happy to run additional shapes/configs on this hardware if useful for tuning
the small/mid buckets too.

[pierre427]

Friendly bump — this is a small, low-risk change (gen-17 large-matrix qmv limit 10→16; small/mid buckets untouched) and it's been sitting without a reviewer. Happy to extend the sweep to the small and mid qmv/qmv_wide buckets on the same M5 Max, or add any shapes/quant configs you'd want to see before merging — whatever's most useful for justifying the retune.

[pierre427]

Ran the small/mid bucket sweep I offered above, on the same M5 Max (g17s). TL;DR: I don't think small (18) or mid (12) should move — the large-bucket change here is the only one with a measurable effect.

Method: I added a bench-only runtime override to get_qmv_batch_limit (a positive MLX_QMV_LIMIT forces the qmv/qmm crossover for every shape) so one build can A/B the vector vs qmm path at any M.

1. Isolated per-op microbench (4-bit affine, group_size 64; min of 60 reps, 2 trials): small (1024², 2048²) and mid (4096², 2880×4096, 4096×2880) shapes read vector ≈ qmm within ~1–2% across M=1–24 (median vec/qmm 0.99–1.03). But as a sanity check I included the large 8192×8192 shape as a positive control, and the microbench fails it — it reads M=12 as a tie (vec/qmm ≈ 1.00), not the ~18% win this PR measured whole-model. So an isolated single-matmul benchmark is too insensitive to resolve this; the large-bucket win looks like a whole-model/system effect (split-K launch cost amortized across the layer chain with a warm cache), not an isolated-kernel property.

2. Whole-model A/B (matching this PR's own methodology) on Qwen3-0.6B-4bit — a small/mid-shaped model (q/o projections in the small bucket, mlp in the mid bucket), warm 400-token KV cache, forcing every matmul to vector vs qmm:

M	force-vector	force-qmm	vec/qmm
11	4.67 ms	4.69 ms	0.995
12	5.51	5.53	0.998
14	5.73	5.69	1.007
16	5.88	5.91	0.996
18	6.25	6.28	0.994
20	6.27	6.31	0.993
24	6.40	6.38	1.003

Median force-vec/force-qmm over M≥12 = 0.997 — within noise across the whole speculative-verify window; the shipped default tracks both. So for a real small/mid-dominated model the limit choice doesn't materially affect decode/verify throughput.

Net: keeping small at 18 and mid at 12 (the inherited fallback values) looks right for gen-17 — no data to justify changing them, which is consistent with this PR only touching the large bucket. Happy to share the harness or run other shapes if useful.

[angeloskath]

@pierre427 this looks great. I will try to get some numbers from more architectures to in-a-way redo this heuristic and then merge.