[0xamlab]

Fixes #4621.

Summary

• Add 64 x 64 instantiations for transposed affine QMM with 4-bit and 8-bit weights.
• Use them for non-NAX, non-batched products with enough output tiles: M >= 512 for 4-bit and M >= 2048 for 8-bit.
• Keep NAX, small-M, batched, non-transposed, and split-K dispatch unchanged.

Tests

• .venv/bin/python python/tests/run.py — 982 passed, 58 skipped.
• .venv/bin/python -m pytest python/tests/test_quantized.py -k qmm_affine_large_tiles -q — 90 subtests passed, covering exact and partial tiles across group sizes 32/64/128, 4/8 bits, and float32/float16/bfloat16.
• .venv/bin/pre-commit run --all — passed.

Benchmark

Apple M4 Max, bfloat16 activations, N = 12288, K = 4096, group size 64. Each value is the median of three process-level medians, with 20 measured calls after 3 warmups and alternating main/PR order. Main is c69303b00f5f61860436e832f23fa8c34b96f93a.

M	Bits	Main (ms)	PR (ms)	Reduction
512	4	5.894	4.465	24.2%
2048	4	28.483	17.618	38.1%
2048	8	19.422	17.519	9.8%
8192	4	164.584	155.337	5.6%
8192	8	160.275	133.895	16.5%

Reproduction script (run three times on main and this branch, alternating order):

import statistics
import time
import mlx.core as mx

K, N, GROUP = 4096, 12288, 64
mx.random.seed(0)
w = (mx.random.normal((N, K)) * 0.02).astype(mx.bfloat16)
for bits in (4, 8):
    wq, scales, biases = mx.quantize(w, group_size=GROUP, bits=bits)
    for M in (512, 2048, 8192):
        x = mx.random.normal((1, M, K)).astype(mx.bfloat16)
        fn = lambda: mx.quantized_matmul(
            x, wq, scales, biases, transpose=True,
            group_size=GROUP, bits=bits,
        )
        for _ in range(3):
            mx.eval(fn())
        samples = []
        for _ in range(20):
            start = time.perf_counter()
            mx.eval(fn())
            samples.append((time.perf_counter() - start) * 1000)
        print(M, bits, f"{statistics.median(samples):.3f} ms")