Use larger tiles for qmm_t at large M

[katnoria]

This branch also includes the one-line QuantizedBlockLoader::load_safe
bounds fix submitted standalone as #4203, since the tile change below
depends on it (see "Why this was never observed" in #4203 for why).

Summary

mx.quantized_matmul's transpose=True path (qmm_t) dispatches a
fixed 32x32x32 output tile with 2x2 simdgroups regardless of
M/N/K/device — sized for small-M (decode) calls. At large M
(prefill-shaped calls, e.g. M >= 4096), this is ~1.7x slower than the
equivalent dense bf16 matmul at the same shape, despite the quantized
weight being ~4x smaller on disk.

Root cause

qmm()'s dispatch (mlx/backend/metal/quantized.cpp, around the tile
selection before get_quantized_kernel_wrapped) hardcodes
bm=32, bn=32, wm=2, wn=2 unconditionally. affine_qmm_t / qmm_t_impl
(mlx/backend/metal/kernels/quantized.h) already accept BM/BK/BN as
template parameters with defaults matching that fixed config — the
dispatch layer simply never varies them. (WM/WN were hardcoded
constexpr inside qmm_t_impl itself, not even template parameters,
until this change.) A sibling kernel, get_gather_qmm_kernel
(jit_kernels.cpp), already threads bm/bn/bk/wm/wn through as JIT
template params — this change brings qmm_t's dispatch up to the same
pattern.

At the fixed 32x32x32 tile, qmm_t dispatches ~4x more threadgroups
than the dense path's adaptive 64x64 tile at the same shape (528,192
vs. 132,160 at the measured production shape below), and each
threadgroup does 8x less useful output work per thread — concretely,
each of the N/bn column-tile-groups re-reads the same row of x from
device memory, and each of the M/bm row-tile-groups re-reads (and
re-dequantizes) the same weight columns, roughly doubling redundant
reads of both operands versus the dense tile.

The fix

Add a dispatch branch in qmm() that selects a 128x64 output tile
(BM=128, BN=64, BK unchanged at 32) for transpose=true calls with
M >= 4096, scoped to {float16, bfloat16} x {group_size 32, 64} x {bits 4, 8}. BK stays 32 in every config tried:
QuantizedBlockLoader requires BK <= group_size, and group_size can
be as small as 32, so BK can't grow without also constraining which
group_sizes the fast path supports. WM/WN are promoted to template
parameters (still defaulting to 2, 2 — an wm=1,wn=2 variant, mirroring
the dense path's own simdgroup split, showed 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 (qmm_n), or a
dtype/group_size/bits combination outside the scoped set — takes the
exact original path and kernel, unconditionally.

Also fixes aligned = N % bn == 0 (was hardcoded % 32), which would
otherwise silently mis-detect alignment for any bn != 32 — latent until
this change since bn was always 32 before.

Why AOT instantiation matters

mlx/backend/metal/CMakeLists.txt sets MLX_METAL_JIT OFF by default —
the configuration pip install mlx actually ships — which links
nojit_kernels.cpp and looks up kernels by name from a fixed,
ahead-of-time-compiled set; there is no runtime JIT fallback in that
build. This PR adds a scoped AOT instantiation
(instantiate_quantized_large_m_tile_types in
mlx/backend/metal/kernels/quantized.metal) for exactly the
(dtype, group_size, bits) combinations the dispatch guard allows, kept
separate from the existing exhaustive instantiate_quantized_all() sweep
to bound this addition's compile-time/binary-size cost. Both
MLX_METAL_JIT=ON and the default OFF build were verified to produce
identical dispatch and numbers — a tile change that only touched
quantized.cpp would work under MLX_METAL_JIT=ON and then throw
[metal::Device] Unable to load kernel ... under the actual default
build the first time M >= 4096 was hit.

Config sweep

Measured on M2 Ultra (non-NAX), production shape M=37723, K=5376, N=14336, group_size=32, bits=8:

config	median ms	vs dense	threadgroups	notes
dense-bf16 (reference)	~359	1.00x	132,160	steel bm=64,bn=64,bk=16,wm=1,wn=2
old (32x32x32, wm=2,wn=2)	624.1	1.74x	528,192	baseline
64x64x32, wm=2,wn=2	559.0	1.56x	132,160	first attempt — helped, short of target
64x64x32, wm=1,wn=2	557.1	1.55x	132,160	mirrors dense's simdgroup config; no measurable gain
128x64x32, wm=2,wn=2 (shipped)	474.4	1.31x	66,080	closest to dense of the configs tried
128x128x32, wm=2,wn=2	1993.1	5.53x	33,040	severe regression — too few threadgroups to hide memory latency; stopped here

Tried 64x64 first, then 128x64 (the winner) and 128x128 (a regression) as
the two additional configs per a fixed "at most 2 more configs" search
budget, then stopped.

Correctness

Reference for all checks: mx.dequantize(w_q, ...) followed by an fp32
dense matmul — not the original bf16 weight — so quantization noise is
isolated from tile-dispatch bugs.

Targeted edge cases (10/10 pass, identical on both MLX_METAL_JIT=ON
and the default OFF build): M just below/at/above the 4096 threshold;
N a multiple of 32 but not 64 (exercises the corrected alignment check);
N not a multiple of 32 at all; group_size 32 and 64; bits 4 and 8;
decode-shaped M (old path, sanity that it's untouched); the production
shape.

Broad stress sweep: 550 combinations (M in {1, 7, 33, 4095, 4096, 4097, 4159, 5000, 8192, 20000, 40000} x N in {1, 17, 32, 63, 64, 65, 100, 4096, 4160, 8000} x (group_size, bits) in {(32,4), (32,8), (64,4), (64,8), (128,4)}, K = group_size * 7) against the same
dequantize+matmul reference, rel_l2 < 1e-2 threshold: 547/550 passed.
The 3 failures are all at M=1 (far below the 4096 threshold — the
large-M branch is never taken) with group_size=64, bits=4 and tiny N;
these reproduce identically on unmodified upstream/main and are
pre-existing 4-bit quantization noise at a small-sample output size,
unrelated to this change.

Guard verification: out-of-scope combinations at large M
(bits=6, group_size=128, dtype=float32, M=5000) correctly fall
through to the original 32x32x32 path with correct output — spot-checked
on the MLX_METAL_JIT=OFF build, where a missing-AOT-kernel bug would
actually throw rather than silently JIT-compile.

Test suite

python -m unittest test_quantized -v (from python/tests/): 34/34
pass on both MLX_METAL_JIT=ON (35.6s) and the default
MLX_METAL_JIT=OFF build (9.1s). Includes test_qmm_large_dims,
test_qmm_shapes, test_qmm, test_qmv*, test_gather_qmm*,
test_non_multiples, test_small_matrix.

Scoping / what's out of scope

• qmm_n (transpose=False) and qmv/qvm/gather variants are
  untouched — this stays scoped to the qmm_t/transpose=true path.
• AOT coverage is {fp16, bf16} x {gs 32, 64} x {bits 4, 8}, not the
  full {float, float16_t, bfloat16_t} x {32, 64, 128} x {2, 3, 4, 5, 6, 8} matrix. Widening is mechanical — one more
  instantiate_quantized_large_m_tile_types(...) call per combination —
  but not done here; happy to extend if maintainers want broader coverage
  before merge.
• wm/wn are template parameters now but not varied from the existing
  default (2, 2) in the shipped config — left as knobs for a future shape
  class if one needs them, not exercised here.

Open items

• 474ms doesn't quite clear a 470ms internal viability target (~0.8-0.9%
  over) or dense parity (359ms). The remaining gap is most plausibly the
  per-K-tile dequantization cost inside QuantizedBlockLoader (bit-unpack
  • affine transform every K-iteration, real additive-vs-dense work that
    no tile-size change removes) rather than remaining tile-geometry
    overhead, since the K-loop trip count is identical across every config
    tried. Not separated from tile-geometry cost by a profiler in this
    round — GPU capture / per-shader-stage counters would settle it.
• End-to-end pipeline result: on our H3 workload (the M=37723 GEMMs
  measured above occur inside a 31B-parameter DiT's denoise step, run on a
  64GB machine already at its wired-memory limit for this configuration),
  the per-GEMM 1.31x improvement does not change end-to-end step time —
  ~466 s/step post-fix vs. ~442 s/step pre-fix, within machine-to-machine
  variance but not an improvement — because the step is bound by
  unified-memory pressure (denoise-phase peak ~58 GB, unchanged by this
  fix) rather than by GEMM dispatch cost. The kernel win here is real and
  should benefit workloads that aren't already sitting at the memory
  ceiling; it just doesn't compose with this particular resident-INT8
  choreography on this machine.

This surfaced a separate, latent bug in QuantizedBlockLoader::load_safe
(comparing the wrong short2 field, unreachable until BROWS != BK,
which this PR is the first change to trigger) — split out into its own
PR (#4203) since it stands alone and is correct independent of
this change.