[ahuzmeza]

Affine quantized_matmul at 3–16 rows on GPUs without NAX.
On main use_qmv_mma returns false for affine (quantized.cpp L645-647), and #4781 opens it only behind is_nax_available(), so those chips stay on qmv / qmv_wide / qmm. This is the non-NAX side of the picture.

Port of #4712's fp_qmv_mma to affine. Differences:

1. q·scale + bias decode, biases as an extra input
2. dispatch on 3 to 16 rows on gen 13/14 and 4 to 16 on gen 15+ (fp: 6),
   N·K >= 2^21 (fp: 2^22). Not on Ultra or NAX
3. wk 4 on every generation (fp: 8 on gen 15+), bn 32 below gen 15 at ≤ 8 rows

4 and 8 bits at group sizes 32, 64 and 128.

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.

Overlaps #4781 (@XXXXRT666, the NAX side) in a few lines of quantized.cpp. Decided to bring them here to not chain PRs.

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

At 6–8 rows main (gen 13/14) moves the MLP projections to qmm_splitk; against main with that limit raised the gain is 2.09–2.84×.

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

End to end on the M2 Pro, mlx-community/gemma-4-e4b-it-4bit with 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 (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 under my direction. the investigation and the numbers are mine.