[ahuzmeza]

On GPUs without NAX, affine quantized_matmul with 3–16 rows never uses qmv_mma, because use_qmv_mma returns false for affine (quantized.cpp L645-647).
Those shapes (batched decode, speculative decoding) fall back to qmv/qmv_wide or to qmm, whose 32-row tiles are mostly padding at this size. #4712 added qmv_mma for the fp formats only. This is the non-NAX side of the picture.

This PR ports #4712's fp_qmv_mma to affine for non-NAX GPUs:

• dequantizes as q·scale + bias, with biases as an extra input
• dispatches at 3–16 rows on gen 13/14 and 4–16 on gen 15+, for N·K ≥ 2^21 (not on Ultra or NAX)
• uses wk 4 on every generation and bn 32 below gen 15 at ≤ 8 rows
• covers 4 and 8 bits at group sizes 32, 64 and 128; all other dispatch paths are unchanged

Results

End to end on an M2 Pro, mlx-community/gemma-4-e4b-it-4bit, mlx_lm.benchmark -p 512 -g 128 -b 8 -n 2:
batch-8 decode goes from 70.0 to 202.1 tok/s (2.89×). Batch 1 remains unchanged.

Speedup over main, chained Qwen3-4B decoder shapes, bf16, group size 64 (95% intervals in the details below):

Rows	3	4	5	6	7	8	10	12	16
M2 Pro, affine8	1.15	1.46	1.77	3.00	3.08	3.15	1.97	1.95	1.92
M2 Pro, affine4	1.16	1.50	1.84	3.27	3.34	3.41	2.07	2.05	2.02
M4, affine8	off	1.02	1.03	1.10	1.26	1.46	1.56	1.88	1.83
M4, affine4	off	1.12	1.34	1.58	1.84	2.21	1.52	1.82	1.93

mxfp8 is the control and measures 1.00 throughout.

Part of the 6–8-row gain on the M2 Pro comes from main's own dispatch. On gen 13/14, main's batch limit sends the MLP projections to qmm, which is slower there than main's qmv. Measured against main with that limit raised, the gain at 6–8 rows is 2.09–2.84×, and batch-8 decode goes from 81.5 to 201.0 tok/s (2.47×).

Overlap with #4781

Both PRs use the same plumbing change. I'm happy to rebase onto #4781 if it lands first.

Measured on the M2 Pro (gen 14) and the M4 (gen 16).
Unmeasured: M1 (gen 13), M3 (gen 15) and A-series phones, enabled by generation.
Non-affine and NAX dispatches are unchanged.

95% intervals

M2 Pro (16-core), main VS. this branch, speedup [95% interval]

Rows	3	4	5	6	7	8	10	12	16
affine8	1.15 [1.14, 1.15]	1.46 [1.43, 1.46]	1.77 [1.74, 1.79]	3.00 [2.88, 3.00]	3.08 [3.05, 3.09]	3.15 [3.14, 3.15]	1.97 [1.96, 1.97]	1.95 [1.91, 1.96]	1.92 [1.88, 1.95]
affine4	1.16 [1.15, 1.18]	1.50 [1.50, 1.51]	1.84 [1.84, 1.86]	3.27 [3.27, 3.28]	3.34 [3.34, 3.34]	3.41 [3.41, 3.43]	2.07 [2.07, 2.08]	2.05 [2.05, 2.06]	2.02 [2.02, 2.02]
mxfp8 (control)	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00

M4 (10-core), same

Rows	3	4	5	6	7	8	10	12	16
affine8	1.00 (off)	1.02 [1.02, 1.03]	1.03 [1.03, 1.03]	1.10 [1.07, 1.11]	1.26 [1.23, 1.27]	1.46 [1.42, 1.47]	1.56 [1.56, 1.60]	1.88 [1.87, 1.89]	1.83 [1.82, 1.83]
affine4	1.00 (off)	1.12 [1.10, 1.13]	1.34 [1.33, 1.36]	1.58 [1.55, 1.60]	1.84 [1.81, 1.86]	2.21 [2.16, 2.22]	1.52 [1.52, 1.55]	1.82 [1.82, 1.85]	1.93 [1.92, 1.93]
mxfp8 (control)	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00	1.00

Method

Chained decoder forward pass on Qwen3-4B shapes, bf16, group size 64. 10 rounds per build in separate processes, alternating. (ratio of medians with a 95% bootstrap interval). mxfp8 is the control. Details are in the script below.

Benchmark script (to reproduce)

Install main and this branch in two venvs, then:

python affine_small_m_bench.py --base <main venv>/bin/python --other <branch venv>/bin/python

"""Chained-forward benchmark for quantized_matmul at few rows, comparing two MLX builds.

One forward = LAYERS decoder-shaped blocks (q/k/v/o, gate/up/down, rms_norm between), each layer's output
feeding the next, so the GPU cannot overlap one layer with the next (as in decode; see mlx #4265). Weights
cycle over 8 distinct layer sets so reads come from DRAM. The non-matmul ops are the same in every mode,
so differences between builds are matmul differences; mxfp8 serves as the control.

The two builds run as separate processes, alternating, for --rounds rounds. The figure per row count is
base/other as the ratio of the medians over rounds, with a 95% bootstrap interval (whole rounds resampled).

usage: python affine_small_m_bench.py --base <main venv>/bin/python --other <port venv>/bin/python
       [--rounds 10] [--rows 1,2,3,4,5,6,7,8,10,12,16] [--modes affine8,affine4,mxfp8] [--shape qwen3-4b]
"""
import argparse
import json
import os
import random
import statistics
import subprocess
import sys
import time

# (hidden, q_out, kv_out, intermediate, layers)
SHAPES = {"qwen3-4b": (2560, 4096, 1024, 9728, 36), "2b-class": (2048, 2048, 1024, 6144, 28)}
MODES = {"affine8": dict(group_size=64, bits=8, mode="affine"), "affine4": dict(group_size=64, bits=4, mode="affine"),
         "mxfp8": dict(mode="mxfp8"), "mxfp4": dict(mode="mxfp4")}
SETS, WARM, REPS = 8, 3, 15


def child(shape_name, modes, rows):
    import mlx.core as mx

    d, qo, kvo, inter, layers = SHAPES[shape_name]
    dims = {"q": (qo, d), "k": (kvo, d), "v": (kvo, d), "o": (d, qo), "g": (inter, d), "u": (inter, d),
            "dn": (d, inter)}

    def qweight(n, k, kw):
        w = (mx.random.normal((n, k)) / (k**0.5)).astype(mx.bfloat16)
        q = mx.quantize(w, **kw)
        b = q[2] if kw["mode"] == "affine" else None
        mx.eval(*[a for a in (q[0], q[1], b) if a is not None])
        return q[0], q[1], b

    def qmm(x, w, kw):
        return mx.quantized_matmul(x, w[0], w[1], w[2], transpose=True, **kw)

    def forward(h, sets, kw, norm_w):
        for li in range(layers):
            L = sets[li % SETS]
            x = mx.fast.rms_norm(h, norm_w, 1e-6)
            q, k, v = qmm(x, L["q"], kw), qmm(x, L["k"], kw), qmm(x, L["v"], kw)
            h = h + qmm(q + (k[:, :1] + v[:, :1]), L["o"], kw)  # k/v feed the chain, as attention would
            x = mx.fast.rms_norm(h, norm_w, 1e-6)
            g = qmm(x, L["g"], kw)
            h = h + qmm(g * mx.sigmoid(g) * qmm(x, L["u"], kw), L["dn"], kw)
        return h

    res = {"mlx": mx.__version__, "arch": mx.device_info().get("architecture"), "ms": {}}
    norm_w = mx.ones((d,), dtype=mx.bfloat16)
    for mode in modes:
        kw = MODES[mode]
        sets = [{name: qweight(n, k, kw) for name, (n, k) in dims.items()} for _ in range(SETS)]
        for M in rows:
            h0 = mx.random.normal((M, d)).astype(mx.bfloat16)
            mx.eval(h0, norm_w)
            for _ in range(WARM):
                mx.eval(forward(h0, sets, kw, norm_w))
            ts = []
            for _ in range(REPS):
                t = time.perf_counter()
                mx.eval(forward(h0, sets, kw, norm_w))
                ts.append((time.perf_counter() - t) * 1000)
            res["ms"][f"{mode}|{M}"] = statistics.median(ts)
        del sets
        mx.clear_cache()
    print(json.dumps(res))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--base", required=True, help="python of the reference build (main)")
    p.add_argument("--other", required=True, help="python of the build under test")
    p.add_argument("--rounds", type=int, default=10)
    p.add_argument("--rows", default="1,2,3,4,5,6,7,8,10,12,16")
    p.add_argument("--modes", default="affine8,affine4,mxfp8")
    p.add_argument("--shape", default="qwen3-4b", choices=SHAPES)
    a = p.parse_args()
    cmd = [os.path.abspath(__file__), "--child", a.shape, a.modes, a.rows]
    runs = {"base": [], "other": []}
    for r in range(a.rounds):
        for arm in (("base", "other") if r % 2 == 0 else ("other", "base")):
            py = a.base if arm == "base" else a.other
            out = subprocess.run([py, *cmd], check=True, capture_output=True, text=True).stdout
            res = json.loads(out.strip().splitlines()[-1])
            runs[arm].append(res["ms"])
            print(f"round {r + 1}, {arm}: mlx {res['mlx']} {res['arch']}", file=sys.stderr, flush=True)
    rng = random.Random(0)
    n = a.rounds
    for mode in a.modes.split(","):
        print(f"[{a.shape}] {mode}   base ms  other ms   base/other  95% CI")
        for M in a.rows.split(","):
            k = f"{mode}|{M}"
            b = [x[k] for x in runs["base"]]
            o = [x[k] for x in runs["other"]]
            boot = []
            for _ in range(10000):
                pick = [rng.randrange(n) for _ in range(n)]
                boot.append(statistics.median(b[i] for i in pick) / statistics.median(o[i] for i in pick))
            boot.sort()
            ratio = statistics.median(b) / statistics.median(o)
            print(f"  M={M:>2}  {statistics.median(b):8.1f}  {statistics.median(o):8.1f}   {ratio:6.2f}x"
                  f"   [{boot[249]:.2f}, {boot[9749]:.2f}]")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--child":
        child(sys.argv[2], sys.argv[3].split(","), [int(m) for m in sys.argv[4].split(",")])
    else:
        main()

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude Code wrote the kernel, its dispatch, the tests and the benchmark script, investigated the dispatch paths and prior work, and ran the measurements, all under my direction

[zcbenz]

This basically looks good to me, thanks for taking care of adding the kernel. Can you please rebase since now #4781 has been merged?