perf(matmul): speed up the speculative verify band for NAX and ALU paths

CHANGELOG.md

@@ -6,7 +6,42 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
-## [0.3.11]
+## [0.3.13]
+
+### Changed
+- Quantized matmul is faster from 2 to 32 rows on machines with no NAX tile,
+  which is the band a speculative-decoding verify step runs in. A 16-row
+  split-K tile carries M <= 16, and each codec enters the route at its own
+  measured row count. Measured on M3 Max at [17920x6656], 16 rows: q4_k
+  5.6 -> 1.9 ms, iq4_nl 4.1 -> 1.5, iq3_xxs 4.2 -> 1.7, iq3_s 4.3 -> 1.8,
+  iq4_xs 4.7 -> 2.6. Single-row decode keeps its own route and is unchanged.
+- `KQ_QMM_SPLITK` forces or disables that route for every codec it supports.
+  iq2_xxs, iq2_xs, iq1_s and iq1_m have no measured entry point, because ggml
+  refuses to encode them without an importance matrix, so they stay on the
+  environment lever.
+- Speculative verify is faster on NAX hardware: a target forward of 8-32
+  rows now costs about 1.4x a single-row forward instead of about 2x, so
+  drafted tokens ride the weight read instead of paying per row. Measured
+  1.40x per full forward at verify widths on a 30B q4_k model.
+- The NAX split-K tile covers every codec that has NAX kernels, not just
+  q6_k and q8_0. Per-call wins from the routing entry are 1.05-1.25x
+  worst-shape and up to 2.5x on small-N projections, biggest for the
+  grid-dequant IQ codecs.
+- `KQ_QMM_SPLITK_NAX` unset now takes a measured per-codec entry M rather
+  than disabling the route. Set it to 0 to disable, or to a split count to
+  force the route at every width up to 32.
+- The non-NAX split-K entry points are picked per device instead of from one
+  table, so NAX machines running with the tile forced off get their own
+  measured entries.
+- `KQ_QMM_SPLITK` now takes effect when NAX is disabled by environment on
+  NAX hardware. It keyed off the hardware rather than the active route, so
+  that combination silently fell back to the plain tile.
+
+### Removed
+- `KQ_MV_EXT_TS` and its staged-activation kernels. Against the current
+  BM=32 tile the route is 0.27-0.72x, so it loses at every width it covered.
+
+## [0.3.12]
 
 ### Added
 - `dsa_kv_qat` takes `f16_round=False`, which stops at the fp8 result and

benchmarks/bench_verify_band.py

@@ -0,0 +1,99 @@
+"""Verify-band row-scaling bench: matmul groups and full forward vs M.
+
+Times one decoder layer's MLP, the lm_head, and (with --full) the full
+forward against a warm cache, at row counts M = 1..32. A near-flat
+curve means verify rows ride the weight read. A linear curve shows the
+speculative verify defect.
+
+Use one process per kernel-env config. Most KQ_* levers latch at first
+dispatch. Label each run with CFG:
+
+    CFG=default python benchmarks/bench_verify_band.py --model m.gguf
+    CFG=splitk16 KQ_QMM_SPLITK=16 python benchmarks/bench_verify_band.py ...
+    CFG=nax_splitk KQ_QMM_SPLITK_NAX=1 python benchmarks/bench_verify_band.py ...
+
+Requires gmlx in the environment (loads the model through the gmlx
+loader so weights come in as real kquant wire tensors).
+"""
+
+import argparse
+import os
+import time
+
+import mlx.core as mx
+from gmlx.loader import load_model
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", required=True, help="GGUF path or gmlx model name")
+parser.add_argument("--rows", default="1,2,3,4,6,8,12,16,17,24,32")
+parser.add_argument("--reps", type=int, default=10)
+parser.add_argument(
+    "--depth", type=int, default=320, help="cache depth for the full-forward sweep"
+)
+parser.add_argument(
+    "--full", action="store_true", help="also run the full-forward sweep (slower)"
+)
+args = parser.parse_args()
+
+ROWS = [int(r) for r in args.rows.split(",")]
+label = os.environ.get("CFG", "default")
+
+model, config, tok = load_model(args.model, verbose=False)
+lm = getattr(model, "language_model", model)
+mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"])
+inner = lm.model if hasattr(lm, "model") else lm
+layer = inner.layers[0]
+hidden = layer.input_layernorm.weight.shape[0]
+head = getattr(lm, "lm_head", None) or inner.embed_tokens.as_linear
+
+
+def sweep(name, fn):
+    out = []
+    for n in ROWS:
+        x = mx.random.normal((1, n, hidden)).astype(mx.float16)
+        for _ in range(3):
+            mx.eval(fn(x))
+        ts = []
+        for _ in range(args.reps):
+            t0 = time.perf_counter()
+            mx.eval(fn(x))
+            ts.append((time.perf_counter() - t0) * 1e3)
+        ts.sort()
+        out.append(f"M{n}={ts[len(ts) // 2]:.2f}")
+    print(f"[{label}] {name:>12}: " + "  ".join(out))
+
+
+sweep("mlp", lambda x: layer.mlp(x))
+sweep("lm_head", lambda x: head(x))
+
+if args.full:
+    ids = tok.encode("The quick brown fox jumps over the lazy dog. " * 60)
+    ids = ids[: args.depth]
+
+    def build_cache():
+        cache = lm.make_cache() if hasattr(lm, "make_cache") else model.make_cache()
+        for i in range(0, len(ids), 128):
+            out = lm(mx.array([ids[i : i + 128]]), cache=cache)
+            mx.eval(out if isinstance(out, mx.array) else out[0])
+        return cache
+
+    cache = build_cache()
+    out = []
+    for n in ROWS:
+        x = mx.array([[ids[-1]] * n])
+        for _ in range(2):
+            o = lm(x, cache=cache)
+            mx.eval(o if isinstance(o, mx.array) else o[0])
+            for c in cache:
+                c.trim(n)
+        ts = []
+        for _ in range(args.reps):
+            t0 = time.perf_counter()
+            o = lm(x, cache=cache)
+            mx.eval(o if isinstance(o, mx.array) else o[0])
+            ts.append((time.perf_counter() - t0) * 1e3)
+            for c in cache:
+                c.trim(n)
+        ts.sort()
+        out.append(f"M{n}={ts[len(ts) // 2]:.1f}")
+    print(f"[{label}] full@{args.depth:>5}: " + "  ".join(out))

benchmarks/bench_verify_band_ab.py

@@ -0,0 +1,248 @@
+"""Thermally-paired A/B of the NAX split-K tile (KQ_QMM_SPLITK_NAX) per codec.
+
+Source of the kq_splitk_nax_min_m entries in src/kquant_matmul.cpp.
+Measured on an M5 Max; re-run on new silicon before trusting them.
+
+KQ_QMM_SPLITK_NAX is read live per dispatch, so all arms share one
+process and one resident copy of the weights. Arms alternate in
+Thue-Morse slot order, not ABBA: an ABBA contrast is the quadratic
+contrast and aliases thermal curvature into the arm difference.
+
+Weights are synthetic random codes in the quantized layout; every bit
+pattern is valid, and it skips the minutes-long IQ encodes.
+
+Cells: codec x (N,K) shape x M x split target. Markdown + JSON out.
+"""
+
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+
+DEFAULT_CODECS = [
+    "q2_k",
+    "q3_k",
+    "q4_k",
+    "q5_k",
+    "q6_k",
+    "q8_0",
+    "q4_0",
+    "q4_1",
+    "q5_0",
+    "q5_1",
+    "iq4_nl",
+    "iq4_xs",
+    "iq3_s",
+    "iq3_xxs",
+    "iq2_xxs",
+    "iq2_xs",
+    "iq2_s",
+    "iq1_s",
+    "iq1_m",
+]
+
+# Muse-Glimmer-30B MLP shapes: gate/up [19968x6656], down [6656x19968].
+DEFAULT_SHAPES = "19968x6656,6656x19968"
+DEFAULT_MS = [1, 2, 4, 6, 8, 10, 12, 16, 20, 24, 32]
+DEFAULT_TARGETS = [32, 16, 8]
+
+# Thue-Morse: t[i] = parity of popcount(i). Balances linear drift without
+# the quadratic aliasing an ABBA block introduces.
+THUE_MORSE = [bin(i).count("1") & 1 for i in range(8)]
+
+
+def group_size_of(codec):
+    return 32 if codec in ("q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl") else 256
+
+
+def effective_sp(target, K, gs):
+    """Mirror the host-side split resolution: the target resolves down to
+    a divisor of K / max(gs, BK). Different targets can be one kernel."""
+    sliceq = max(gs, 64)
+    nblk = K // sliceq
+    sp = min(target, nblk)
+    while sp > 1 and nblk % sp != 0:
+        sp -= 1
+    return sp
+
+
+def probe_layout(codec, N, K):
+    import mlx.core as mx
+    import numpy as np
+
+    import mlx_kquant as kq
+
+    if codec.startswith("iq"):
+        tests_dir = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tests"
+        )
+        sys.path.insert(0, tests_dir)
+        from test_codecs import CODECS, _synth_iq_wire
+
+        _, wpb, bpb, _, _ = CODECS[codec]
+        rng = np.random.default_rng(N + K)
+        wire = _synth_iq_wire(rng, bpb, N * (K // wpb))
+        return (
+            mx.array(wire.reshape(N, (K // wpb) * bpb)),
+            mx.array(np.zeros((1,), dtype=np.uint8)),
+        )
+
+    wf = mx.random.normal((8, K)).astype(mx.float32)
+    w8, s8 = kq.quantize(wf, codec)
+    mx.eval(w8, s8)
+    rng = np.random.default_rng(N + K)
+
+    def full(sample):
+        a = np.asarray(sample)
+        shape = (N,) + a.shape[1:]
+        if np.issubdtype(a.dtype, np.floating):
+            return (rng.standard_normal(shape) * 0.01).astype(a.dtype)
+        info = np.iinfo(a.dtype)
+        return rng.integers(
+            info.min, info.max, size=shape, endpoint=True, dtype=a.dtype
+        )
+
+    return mx.array(full(w8)), mx.array(full(s8))
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--codecs", nargs="+", default=DEFAULT_CODECS)
+    ap.add_argument("--shapes", default=DEFAULT_SHAPES)
+    ap.add_argument("--ms", type=int, nargs="+", default=DEFAULT_MS)
+    ap.add_argument("--targets", type=int, nargs="+", default=DEFAULT_TARGETS)
+    ap.add_argument("--iters", type=int, default=12)
+    ap.add_argument("--warmup", type=int, default=4)
+    ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16"])
+    ap.add_argument("--json-out")
+    ap.add_argument("--md-out")
+    args = ap.parse_args()
+
+    import mlx.core as mx
+
+    import mlx_kquant as kq
+
+    if not kq.nax_available():
+        sys.exit("NAX not available on this device; nothing to measure.")
+
+    mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"])
+    dt = mx.bfloat16 if args.dtype == "bfloat16" else mx.float16
+    shapes = [tuple(int(v) for v in s.split("x")) for s in args.shapes.split(",")]
+
+    def time_arm(target, x, w, s, codec):
+        os.environ["KQ_QMM_SPLITK_NAX"] = str(target)
+
+        def call(xx):
+            return kq.quantized_matmul(xx, w, s, codec, transpose=True)
+
+        o = call(x)
+        mx.eval(o)
+        t0 = time.perf_counter()
+        o = call(x)
+        mx.eval(o)
+        return (time.perf_counter() - t0) * 1e3
+
+    results = []
+    for codec in args.codecs:
+        gs = group_size_of(codec)
+        for N, K in shapes:
+            try:
+                w, s = probe_layout(codec, N, K)
+                mx.eval(w, s)
+            except Exception as e:  # codec cannot synthesize at this shape
+                print(f"skip {codec} [{N}x{K}]: {e}", file=sys.stderr)
+                continue
+            wbytes = w.nbytes + s.nbytes
+            for M in args.ms:
+                x = mx.random.normal((M, K), key=mx.random.key(M)).astype(dt)
+                mx.eval(x)
+                for target in args.targets:
+                    sp = effective_sp(target, K, gs)
+                    if sp <= 1:
+                        continue
+                    for _ in range(args.warmup):
+                        time_arm(0, x, w, s, codec)
+                        time_arm(target, x, w, s, codec)
+                    off, on = [], []
+                    for _rep in range(args.iters):
+                        for slot in THUE_MORSE:
+                            # slot 0 -> off first, slot 1 -> on first
+                            if slot == 0:
+                                off.append(time_arm(0, x, w, s, codec))
+                                on.append(time_arm(target, x, w, s, codec))
+                            else:
+                                on.append(time_arm(target, x, w, s, codec))
+                                off.append(time_arm(0, x, w, s, codec))
+                    off_ms = statistics.median(off)
+                    on_ms = statistics.median(on)
+                    results.append(
+                        {
+                            "codec": codec,
+                            "N": N,
+                            "K": K,
+                            "M": M,
+                            "target": target,
+                            "sp": sp,
+                            "off_ms": off_ms,
+                            "on_ms": on_ms,
+                            "speedup": off_ms / on_ms,
+                            "off_gbs": wbytes / (off_ms * 1e-3) / 1e9,
+                            "on_gbs": wbytes / (on_ms * 1e-3) / 1e9,
+                        }
+                    )
+                    print(
+                        f"{codec:8s} [{N}x{K}] M{M:<3d} t{target:<3d} sp={sp:<3d} "
+                        f"off={off_ms:.3f} on={on_ms:.3f} "
+                        f"speedup={off_ms / on_ms:.3f}",
+                        flush=True,
+                    )
+            del w, s
+
+    if args.json_out:
+        with open(args.json_out, "w") as f:
+            json.dump({"device": mx.device_info()["device_name"], "rows": results}, f)
+
+    lines = ["# NAX split-K verify band A/B", ""]
+    lines.append(f"Device: {mx.device_info()['device_name']}  dtype: {args.dtype}")
+    lines.append("")
+    lines.append("speedup = off / on; >1 means split-K is faster.")
+    lines.append("")
+    for codec in args.codecs:
+        rows = [r for r in results if r["codec"] == codec]
+        if not rows:
+            continue
+        lines.append(f"## {codec}")
+        lines.append("")
+        for N, K in shapes:
+            sub = [r for r in rows if r["N"] == N and r["K"] == K]
+            if not sub:
+                continue
+            targets = sorted({r["target"] for r in sub}, reverse=True)
+            lines.append(f"### [{N}x{K}]")
+            lines.append("")
+            lines.append(
+                "| M | off ms | " + " | ".join(f"t{t}" for t in targets) + " |"
+            )
+            lines.append("|---|---|" + "---|" * len(targets))
+            for M in args.ms:
+                cells = [r for r in sub if r["M"] == M]
+                if not cells:
+                    continue
+                row = f"| {M} | {cells[0]['off_ms']:.3f} |"
+                for t in targets:
+                    c = [r for r in cells if r["target"] == t]
+                    row += f" {c[0]['speedup']:.3f} |" if c else " - |"
+                lines.append(row)
+            lines.append("")
+    md = "\n".join(lines)
+    if args.md_out:
+        with open(args.md_out, "w") as f:
+            f.write(md + "\n")
+    else:
+        print(md)
+
+
+if __name__ == "__main__":
+    main()