[bri-prism]

FP32 quantized_matmul inputs currently require FP16/BF16 affine scales and biases to be widened before each call. This keeps those constants narrow on Metal and reads them directly in the unbatched, one-row qmv_fast path, widening each load in registers. It avoids two conversion kernels per projection without changing the arithmetic.

Other affine dispatches widen the constants inside eval_gpu before entering their existing kernels. CPU streams and floating-point quantization modes retain their existing input handling. Scale and bias gradients are cast back to their original dtypes. Both JIT and precompiled Metal entrypoints are included.

Measured result

Apple M5 Pro, a local quantized text decoder, greedy 256-token decode, excluding prefill:

Prompt	Existing widen-first path	Narrow reads	Speedup
Explanation	110.53 tokens/s	126.05 tokens/s	1.140x
Coding	110.64 tokens/s	126.09 tokens/s	1.140x

These measurements used one JIT binary with the new path temporarily disabled/enabled in the order off/on/off/on/off, with 32 warmup tokens per path. The table reports the median of three control samples and two candidate samples. Maximum adjacent control drift was 1.05% and 0.21%, respectively. Tokens and logits matched bitwise across all samples. The temporary control is absent from this PR, and the clean rebuilt candidate reproduced the same tokens and logits. These are workload-specific measurements; the model checkpoint and full-model harness are not included in this PR.

Validation

On M5 Pro, both JIT and precompiled Metal builds pass 85 tests, with one skipped and 6,261 subtests passed, using upstream's pytest configuration:

python python/tests/run.py python/tests/test_quantized.py python/tests/test_export_import.py -q

The tests cover FP16/BF16 constants, all seven affine bit widths, all three group sizes, aligned and partial output rows, strided and batched layouts, vector inputs, CPU streams, fallback dispatches, gradient dtypes, compiled forward/backward calls with fresh inputs, and export/import round trips. Temporary pipeline tracing confirmed selection of all 84 mixed-kernel variants; the tracing is removed.

Formatting, Python lint and stream-argument checks pass. Other Apple GPU generations and CUDA execution have not been tested. The throughput numbers above are from JIT execution; precompiled Metal was checked for correctness.

Operator reproduction

Run the following on main and this branch with the same build configuration. It measures synchronized operator wall time, including Python and dispatch overhead, on synthetic weights:

import statistics
import time
import mlx.core as mx

mx.random.seed(7)
x = mx.random.normal((1, 4096))
w = mx.random.normal((14336, 4096)).astype(mx.float16)
wq, scales, biases = mx.quantize(w, group_size=128, bits=2)
mx.eval(x, wq, scales, biases)

def call():
    return mx.quantized_matmul(x, wq, scales, biases, True, 128, 2)

warmup_end = time.perf_counter() + 0.25
while time.perf_counter() < warmup_end:
    mx.eval(call())

samples = []
for _ in range(5):
    start = time.perf_counter_ns()
    for _ in range(50):
        mx.eval(call())
    samples.append((time.perf_counter_ns() - start) / 50 / 1000)
print(f"median: {statistics.median(samples):.2f} us/call")

AI usage disclosure: Implementation, tests, benchmark analysis, and this description were prepared with AI assistance.

[zcbenz]

It seems that all affine qmv/qmm kernels can actually benefit from this trick? If they do I think it would be a cleaner change updating them all than adding a new kernel. And btw please create pull requests from personal accounts as maintainers can't rebase PRs from org accounts.