diff --git a/CHANGELOG.md b/CHANGELOG.md index f981d0f..3265c94 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/benchmarks/bench_verify_band.py b/benchmarks/bench_verify_band.py new file mode 100644 index 0000000..b63cb64 --- /dev/null +++ b/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)) diff --git a/benchmarks/bench_verify_band_ab.py b/benchmarks/bench_verify_band_ab.py new file mode 100644 index 0000000..f989d85 --- /dev/null +++ b/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() diff --git a/docs/kernels.md b/docs/kernels.md index 2192f11..d7b537b 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -66,15 +66,16 @@ Tuning levers (defaults are right for normal use): - `KQ_MV_EXT_NR` - `2` selects the two-rows-per-thread `mv_ext` variant (q6_k, M 5-12), which halves activation cache traffic but measured no faster than the shipped kernels. Kept as a probe for future silicon. Default `1` (shipped behavior). -- `KQ_QMM_SPLITK_NAX` - split-K on the NAX BM=32 tile; the value is the target slice count (`1` = - auto 32, `0` off). Lifts the collapsed M 9-16 band 65-76% on M5 Max. q6_k/q8_0, M <= 32; read - live per call. Default off. -- `KQ_QMM_SPLITK` - split-K for the plain small-M qmm (target slice count, `0` off). Measured flat - to negative on M5 Max; kept as a probe. K-quants plus q8_0, M <= 32. Default off. -- `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` / `KQ_MV_EXT_TS` - `mv_ext` activation-traffic - experiments: shuffle-broadcast (`1`), wide nxpsg (`16`/`32`), half-precision chunk dots (`1`), - threadgroup-staged activations (`1`). q6_k M 4-12 only. `HD` measured +4-5% at M 8; the rest flat - to negative on M5 Max. Kept as probes. Default off. +- `KQ_QMM_SPLITK_NAX` - split-K on the NAX BM=32 tile; `0` disables the route, a value at or above + `1` forces it and sets the target slice count. Unset takes the per-codec entry M in + `kq_splitk_nax_min_m`, measured on M5 Max. Every codec with NAX kernels, M <= 32; read live per + call, so both arms can share one process. +- `KQ_QMM_SPLITK` - the same lever for the plain small-M qmm, used when NAX is absent or disabled. + Entry points come from a per-device table. K-quants, legacy quants and the IQ codecs, M <= 32. +- `KQ_MV_EXT_SB` / `KQ_MV_EXT_NX` / `KQ_MV_EXT_HD` - `mv_ext` activation-traffic experiments: + shuffle-broadcast (`1`), wide nxpsg (`16`/`32`), half-precision chunk dots (`1`). q6_k M 4-12 + only. `HD` measured +4-5% at M 8; the rest flat to negative on M5 Max. Kept as probes. Default + off. ## MoE GLU diff --git a/metal/kq_quantized.metal b/metal/kq_quantized.metal index da9cff9..7f830c3 100644 --- a/metal/kq_quantized.metal +++ b/metal/kq_quantized.metal @@ -37,6 +37,18 @@ bits, \ aligned_N) +// BM=16 split-K tile for M <= 16 (codecs with the bm16 template arm). +#define instantiate_kquant_qmm_t_splitk_bm16(type, gs, bits, aligned_N, codec) \ + instantiate_kernel( \ + "kquant_" #codec "_qmm_t_splitk_bm16_" #type "_gs_" #gs "_b_" #bits \ + "_alN_" #aligned_N, \ + kq_ ## codec ## _qmm_t_splitk, \ + type, \ + gs, \ + bits, \ + aligned_N, \ + true) + #define instantiate_kquant_qmm_n(type, gs, bits, batched, codec) \ instantiate_kernel( \ "kquant_" #codec "_qmm_n_" #type "_gs_" #gs "_b_" #bits \ @@ -110,6 +122,8 @@ instantiate_kquant_qmm_t(type, 32, 8, false, 1, q8_0) \ instantiate_kquant_qmm_t_splitk(type, 32, 8, true, q8_0) \ instantiate_kquant_qmm_t_splitk(type, 32, 8, false, q8_0) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 32, 8, true, q8_0) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 32, 8, false, q8_0) \ instantiate_kquant_qmm_n(type, 32, 8, 0, q8_0) \ instantiate_kquant_qmm_n(type, 32, 8, 1, q8_0) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 32, 8, q8_0) \ @@ -255,6 +269,8 @@ instantiate_kquant_q5_0_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 4, false, 1, q4_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 4, true, q4_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 4, false, q4_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 4, true, q4_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 4, false, q4_k) \ instantiate_kquant_qmm_n(type, 256, 4, 0, q4_k) \ instantiate_kquant_qmm_n(type, 256, 4, 1, q4_k) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 4, q4_k) \ @@ -284,6 +300,8 @@ instantiate_kquant_q4_k_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 5, false, 1, q5_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 5, true, q5_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 5, false, q5_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 5, true, q5_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 5, false, q5_k) \ instantiate_kquant_qmm_n(type, 256, 5, 0, q5_k) \ instantiate_kquant_qmm_n(type, 256, 5, 1, q5_k) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 5, q5_k) \ @@ -314,6 +332,8 @@ instantiate_kquant_q5_k_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 6, false, 1, q6_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 6, true, q6_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 6, false, q6_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 6, true, q6_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 6, false, q6_k) \ instantiate_kquant_qmm_n(type, 256, 6, 0, q6_k) \ instantiate_kquant_qmm_n(type, 256, 6, 1, q6_k) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 6, q6_k) \ @@ -475,30 +495,6 @@ instantiate_mv_ext_nx_all(q6_k, 256, 6, 32) instantiate_mv_ext_hd_for_type(codec, gs, bits, float16_t) instantiate_mv_ext_hd_all(q6_k, 256, 6) -// Staged-activation experiment (KQ_MV_EXT_TS=1): the M x 128 activation -// window stages into threadgroup memory once per K-step and 8 simdgroups -// (32 output rows) share it -- the activation-path lever the sb/nr2/nx -// falsifications never isolated. Dot math identical to base. q6_k M 4-12. -#define instantiate_mv_ext_ts(codec, type, gs, bits, m) \ - instantiate_kernel( \ - "kquant_" #codec "_mv_ext_" #type "_gs_" #gs "_b_" #bits "_m" #m \ - "_ts", \ - kq_ ## codec ## _mv_ext_ts, type, m, 8, 8) -#define instantiate_mv_ext_ts_for_type(codec, gs, bits, type) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 4) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 5) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 6) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 7) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 8) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 9) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 10) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 11) \ - instantiate_mv_ext_ts(codec, type, gs, bits, 12) -#define instantiate_mv_ext_ts_all(codec, gs, bits) \ - instantiate_mv_ext_ts_for_type(codec, gs, bits, float) \ - instantiate_mv_ext_ts_for_type(codec, gs, bits, bfloat16_t) \ - instantiate_mv_ext_ts_for_type(codec, gs, bits, float16_t) -instantiate_mv_ext_ts_all(q6_k, 256, 6) #define instantiate_kquant_q3_k_for_type(type) \ instantiate_kquant_batched(verify_qmv, type, 256, 3, 0, q3_k) \ @@ -514,6 +510,8 @@ instantiate_mv_ext_ts_all(q6_k, 256, 6) instantiate_kquant_qmm_t(type, 256, 3, false, 1, q3_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, true, q3_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, false, q3_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, true, q3_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, false, q3_k) \ instantiate_kquant_qmm_n(type, 256, 3, 0, q3_k) \ instantiate_kquant_qmm_n(type, 256, 3, 1, q3_k) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 3, q3_k) \ @@ -543,6 +541,8 @@ instantiate_kquant_q3_k_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 2, false, 1, q2_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, true, q2_k) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, false, q2_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, true, q2_k) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, false, q2_k) \ instantiate_kquant_qmm_n(type, 256, 2, 0, q2_k) \ instantiate_kquant_qmm_n(type, 256, 2, 1, q2_k) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 2, q2_k) \ @@ -573,6 +573,8 @@ instantiate_kquant_q2_k_for_type(float16_t) instantiate_kquant_qmm_t(type, 32, 4, false, 1, iq4_nl) \ instantiate_kquant_qmm_t_splitk(type, 32, 4, true, iq4_nl) \ instantiate_kquant_qmm_t_splitk(type, 32, 4, false, iq4_nl) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 32, 4, true, iq4_nl) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 32, 4, false, iq4_nl) \ instantiate_kquant_qmm_n(type, 32, 4, 0, iq4_nl) \ instantiate_kquant_qmm_n(type, 32, 4, 1, iq4_nl) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 32, 4, iq4_nl) \ @@ -657,6 +659,8 @@ instantiate_kquant_nvfp4_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 4, false, 1, iq4_xs) \ instantiate_kquant_qmm_t_splitk(type, 256, 4, true, iq4_xs) \ instantiate_kquant_qmm_t_splitk(type, 256, 4, false, iq4_xs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 4, true, iq4_xs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 4, false, iq4_xs) \ instantiate_kquant_qmm_n(type, 256, 4, 0, iq4_xs) \ instantiate_kquant_qmm_n(type, 256, 4, 1, iq4_xs) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 4, iq4_xs) \ @@ -685,6 +689,8 @@ instantiate_kquant_iq4_xs_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 3, false, 1, iq3_xxs) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, true, iq3_xxs) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, false, iq3_xxs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, true, iq3_xxs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, false, iq3_xxs) \ instantiate_kquant_qmm_n(type, 256, 3, 0, iq3_xxs) \ instantiate_kquant_qmm_n(type, 256, 3, 1, iq3_xxs) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 3, iq3_xxs) \ @@ -713,6 +719,8 @@ instantiate_kquant_iq3_xxs_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 3, false, 1, iq3_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, true, iq3_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 3, false, iq3_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, true, iq3_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 3, false, iq3_s) \ instantiate_kquant_qmm_n(type, 256, 3, 0, iq3_s) \ instantiate_kquant_qmm_n(type, 256, 3, 1, iq3_s) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 3, iq3_s) \ @@ -741,6 +749,8 @@ instantiate_kquant_iq3_s_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 2, false, 1, iq2_xxs) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, true, iq2_xxs) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, false, iq2_xxs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, true, iq2_xxs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, false, iq2_xxs) \ instantiate_kquant_qmm_n(type, 256, 2, 0, iq2_xxs) \ instantiate_kquant_qmm_n(type, 256, 2, 1, iq2_xxs) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 2, iq2_xxs) \ @@ -769,6 +779,8 @@ instantiate_kquant_iq2_xxs_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 2, false, 1, iq2_xs) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, true, iq2_xs) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, false, iq2_xs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, true, iq2_xs) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, false, iq2_xs) \ instantiate_kquant_qmm_n(type, 256, 2, 0, iq2_xs) \ instantiate_kquant_qmm_n(type, 256, 2, 1, iq2_xs) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 2, iq2_xs) \ @@ -797,6 +809,8 @@ instantiate_kquant_iq2_xs_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 2, false, 1, iq2_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, true, iq2_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 2, false, iq2_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, true, iq2_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 2, false, iq2_s) \ instantiate_kquant_qmm_n(type, 256, 2, 0, iq2_s) \ instantiate_kquant_qmm_n(type, 256, 2, 1, iq2_s) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 2, iq2_s) \ @@ -825,6 +839,8 @@ instantiate_kquant_iq2_s_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 1, false, 1, iq1_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 1, true, iq1_s) \ instantiate_kquant_qmm_t_splitk(type, 256, 1, false, iq1_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 1, true, iq1_s) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 1, false, iq1_s) \ instantiate_kquant_qmm_n(type, 256, 1, 0, iq1_s) \ instantiate_kquant_qmm_n(type, 256, 1, 1, iq1_s) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 1, iq1_s) \ @@ -852,6 +868,8 @@ instantiate_kquant_iq1_s_for_type(float16_t) instantiate_kquant_qmm_t(type, 256, 1, false, 1, iq1_m) \ instantiate_kquant_qmm_t_splitk(type, 256, 1, true, iq1_m) \ instantiate_kquant_qmm_t_splitk(type, 256, 1, false, iq1_m) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 1, true, iq1_m) \ + instantiate_kquant_qmm_t_splitk_bm16(type, 256, 1, false, iq1_m) \ instantiate_kquant_qmm_n(type, 256, 1, 0, iq1_m) \ instantiate_kquant_qmm_n(type, 256, 1, 1, iq1_m) \ instantiate_kquant_gather_qmv(gather_qmv_fast, type, 256, 1, iq1_m) \ diff --git a/metal/kq_quantized_nax.metal b/metal/kq_quantized_nax.metal index e5be078..662e097 100644 --- a/metal/kq_quantized_nax.metal +++ b/metal/kq_quantized_nax.metal @@ -120,8 +120,8 @@ // grid.z K-slices into T partials + shared accum fold. The plain small-M // grid is TG-count starved (ceil(N/64) x 1 threadgroups at decode shapes); // splitting K multiplies occupancy without touching the fragment shape. -// q6_k + q8_0 only; no batched or float x variants (route gates match -// qmm_nax and non_batched). +// Full K/IQ codec coverage, matching the bm32 tile's instantiation set; no +// batched or float x variants (route gates match qmm_nax and non_batched). #define instantiate_kquant_nax_qmm_t_splitk(type, gs, bits, aligned_N, codec) \ instantiate_kernel( \ "kquant_" #codec "_qmm_t_nax_splitk_" #type "_gs_" #gs "_b_" #bits \ @@ -135,6 +135,23 @@ instantiate_kquant_nax_qmm_t_splitk(bfloat16_t, gs, bits, false, codec) instantiate_kquant_nax_splitk(q6_k, 256, 6) instantiate_kquant_nax_splitk(q8_0, 32, 8) +instantiate_kquant_nax_splitk(q4_k, 256, 4) +instantiate_kquant_nax_splitk(q5_k, 256, 5) +instantiate_kquant_nax_splitk(q3_k, 256, 3) +instantiate_kquant_nax_splitk(q2_k, 256, 2) +instantiate_kquant_nax_splitk(q4_0, 32, 4) +instantiate_kquant_nax_splitk(q4_1, 32, 4) +instantiate_kquant_nax_splitk(q5_0, 32, 5) +instantiate_kquant_nax_splitk(q5_1, 32, 5) +instantiate_kquant_nax_splitk(iq4_nl, 32, 4) +instantiate_kquant_nax_splitk(iq4_xs, 256, 4) +instantiate_kquant_nax_splitk(iq3_xxs, 256, 3) +instantiate_kquant_nax_splitk(iq3_s, 256, 3) +instantiate_kquant_nax_splitk(iq2_xxs, 256, 2) +instantiate_kquant_nax_splitk(iq2_xs, 256, 2) +instantiate_kquant_nax_splitk(iq2_s, 256, 2) +instantiate_kquant_nax_splitk(iq1_s, 256, 1) +instantiate_kquant_nax_splitk(iq1_m, 256, 1) // Double-buffered BM=64 qmm_t, name-suffixed _db: dispatched by the host // solely for the M33-64 decode band (kq_smallbm_policy db64 + KQ_NAX_DB64). diff --git a/metal/mlx/backend/metal/kernels/kq_quantized.h b/metal/mlx/backend/metal/kernels/kq_quantized.h index 0d0449d..ebe1144 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized.h @@ -786,111 +786,6 @@ METAL_FUNC void kq_mv_ext_hd_impl( } } -// Staged-activation variant of kq_mv_ext_impl (suffix _ts). Weights stay in -// registers exactly as in the base kernel; the change is the activation -// path, the one lever the falsified variants never isolated: sb swapped -// loads for shuffles (worse throughput), nr2 amortized via registers -// (spilled), x16/x32 only changed which thread issues the loads. Here the -// M x (nxpsg*16) activation window stages into threadgroup memory once per -// K-step via a cooperative load, every row-thread dots from on-core SRAM, -// and the threadgroup carries nsg_ts simdgroups (32 rows at nsg_ts=8) so -// cross-TG device activation traffic drops rows_per_tg/8-fold vs the base -// kernel. Dot arithmetic is bit-identical to base (staging is a copy). -// K must be a multiple of nxpsg*16 (q6_k superblock 256 guarantees it). -template -METAL_FUNC void kq_mv_ext_ts_impl( - const device uint8_t* w, - const device T* x, - device T* y, - threadgroup T* staged, // r1ptg * nxpsg * 16 elements - const constant int& in_vec_size, // K - const constant int& out_vec_size, // N - uint3 tgpig, - ushort tiisg, - ushort sgitg) { - constexpr short nypsg = 32 / nxpsg; // output rows per simdgroup - constexpr short chpb = Codec::superblock / 16; // 16-weight chunks per block - constexpr short stage_w = nxpsg * 16; // staged K-window elements per row - const short tx = tiisg % nxpsg; // K position within the row group - const short ty = tiisg / nxpsg; // which of nypsg rows this thread owns - - const int i01 = tgpig.x * (nypsg * nsg) + nypsg * sgitg + ty; // output row - const int i11 = tgpig.y * r1ptg; // first activation column (grid.y==1 -> 0) - - const int nb = in_vec_size / Codec::superblock; - const int row_bytes = nb * Codec::block_bytes; - // Clamp OOB rows to row 0 for a valid read; the store is masked below. - const device uint8_t* w_row = - (i01 < out_vec_size) ? w + static_cast(i01) * row_bytes : w; - - const short lin = sgitg * 32 + tiisg; // linear thread id in the TG - constexpr short tg_threads = nsg * 32; - - float sumf[r1ptg]; -#pragma unroll - for (short ir1 = 0; ir1 < r1ptg; ++ir1) { - sumf[ir1] = 0.0f; - } - - for (int base = 0; 16 * base < in_vec_size; base += nxpsg) { - // Cooperative stage of the M x stage_w activation window. - threadgroup_barrier(mem_flags::mem_threadgroup); - const int kw = 16 * base; - for (short f = lin; f < r1ptg * stage_w; f += tg_threads) { - const short ir1 = f / stage_w; - const short j = f % stage_w; - staged[f] = x[static_cast(i11 + ir1) * in_vec_size + kw + j]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - const int ich = base + tx; - const int ib = ich / chpb; // super-block index - const short cch = ich % chpb; // chunk within super-block - const device uint8_t* block = - w_row + static_cast(ib) * Codec::block_bytes; - float4x4 lx; - Codec::deq_chunk16(block, cch, lx); - const threadgroup T* sp = staged + tx * 16; -#pragma unroll - for (short ir1 = 0; ir1 < r1ptg; ++ir1) { - const threadgroup T* yp = sp + ir1 * stage_w; - const float4 a0 = float4(*(const threadgroup vec*)(yp + 0)); - const float4 a1 = float4(*(const threadgroup vec*)(yp + 4)); - const float4 a2 = float4(*(const threadgroup vec*)(yp + 8)); - const float4 a3 = float4(*(const threadgroup vec*)(yp + 12)); - sumf[ir1] += - dot(lx[0], a0) + dot(lx[1], a1) + dot(lx[2], a2) + dot(lx[3], a3); - } - } - -#pragma unroll - for (short ir1 = 0; ir1 < r1ptg; ++ir1) { - if (nxpsg >= 32) { - sumf[ir1] += simd_shuffle_down(sumf[ir1], 16); - } - if (nxpsg >= 16) { - sumf[ir1] += simd_shuffle_down(sumf[ir1], 8); - } - if (nxpsg >= 8) { - sumf[ir1] += simd_shuffle_down(sumf[ir1], 4); - } - if (nxpsg >= 4) { - sumf[ir1] += simd_shuffle_down(sumf[ir1], 2); - } - if (nxpsg >= 2) { - sumf[ir1] += simd_shuffle_down(sumf[ir1], 1); - } - } - - if (tx == 0 && i01 < out_vec_size) { -#pragma unroll - for (short ir1 = 0; ir1 < r1ptg; ++ir1) { - y[static_cast(i11 + ir1) * out_vec_size + i01] = - static_cast(sumf[ir1]); - } - } -} - // Wide-M variant of kq_mv_ext_impl: each thread owns nr0 CONSECUTIVE output // rows instead of one. The nr0=1 kernel re-loads all r1ptg activation columns // per 16-weight chunk per row, so activation cache traffic scales as @@ -1528,7 +1423,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q8_0_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1546,7 +1446,8 @@ template static_assert( group_size == KQ_Q8_0_GROUP, "Q8_0 kernel requires group_size=32"); static_assert(bits == 8, "Q8_0 kernel requires bits=8"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 7c301d4..185d13d 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h @@ -1266,7 +1266,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq4_xs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1283,7 +1288,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ4_XS_SUPERBLOCK, "IQ4_XS requires gs=256"); static_assert(bits == 4, "IQ4_XS requires bits=4"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -1717,7 +1723,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq3_xxs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1734,7 +1745,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ3_XXS_SUPERBLOCK, "IQ3_XXS requires gs=256"); static_assert(bits == 3, "IQ3_XXS requires bits=3"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2173,7 +2185,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq3_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2190,7 +2207,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ3_S_SUPERBLOCK, "IQ3_S requires gs=256"); static_assert(bits == 3, "IQ3_S requires bits=3"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2625,7 +2643,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq2_xxs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2642,7 +2665,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ2_XXS_SUPERBLOCK, "IQ2_XXS requires gs=256"); static_assert(bits == 2, "IQ2_XXS requires bits=2"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -3080,7 +3104,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq2_xs_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3097,7 +3126,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ2_XS_SUPERBLOCK, "IQ2_XS requires gs=256"); static_assert(bits == 2, "IQ2_XS requires bits=2"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -3537,7 +3567,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq2_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3554,7 +3589,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ2_S_SUPERBLOCK, "IQ2_S requires gs=256"); static_assert(bits == 2, "IQ2_S requires bits=2"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -3988,7 +4024,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq1_s_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4005,7 +4046,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ1_S_SUPERBLOCK, "IQ1_S requires gs=256"); static_assert(bits == 1, "IQ1_S requires bits=1"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -4461,7 +4503,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq1_m_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4478,7 +4525,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ1_M_SUPERBLOCK, "IQ1_M requires gs=256"); static_assert(bits == 1, "IQ1_M requires bits=1"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h index c5e499c..61f8fb4 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_kquants.h @@ -613,7 +613,14 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +// bm16=true is the M<=16 tile: the BM=32 tile spends half of its MMA +// issues on row padding there; BM=16 halves that waste. +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q4_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -631,7 +638,8 @@ template static_assert( group_size == KQ_Q4_K_SUPERBLOCK, "Q4_K kernel requires group_size=256"); static_assert(bits == 4, "Q4_K kernel requires bits=4"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -1675,7 +1683,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q5_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -1693,7 +1706,8 @@ template static_assert( group_size == KQ_Q5_K_SUPERBLOCK, "Q5_K kernel requires group_size=256"); static_assert(bits == 5, "Q5_K kernel requires bits=5"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2588,7 +2602,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q6_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2606,7 +2625,8 @@ template static_assert( group_size == KQ_Q6_K_SUPERBLOCK, "Q6_K kernel requires group_size=256"); static_assert(bits == 6, "Q6_K kernel requires bits=6"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -2898,23 +2918,6 @@ template w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); } -template -[[kernel]] void kq_q6_k_mv_ext_ts( - const device uint8_t* w, - const device uint8_t* /* scales */, - const device T* x, - device T* y, - const constant int& in_vec_size, // K - const constant int& out_vec_size, // N - const constant int& /* vm */, // == r1ptg - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - threadgroup T staged[r1ptg * nxpsg * 16]; - kq_mv_ext_ts_impl( - w, x, y, staged, in_vec_size, out_vec_size, tgpig, tiisg, sgitg); -} - template [[kernel]] void kq_q6_k_mv_ext_hd( const device uint8_t* w, @@ -3704,7 +3707,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q3_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -3722,7 +3730,8 @@ template static_assert( group_size == KQ_Q3_K_SUPERBLOCK, "Q3_K kernel requires group_size=256"); static_assert(bits == 3, "Q3_K kernel requires bits=3"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; @@ -4595,7 +4604,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_q2_k_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -4613,7 +4627,8 @@ template static_assert( group_size == KQ_Q2_K_SUPERBLOCK, "Q2_K kernel requires group_size=256"); static_assert(bits == 2, "Q2_K kernel requires bits=2"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h index 3980977..bf41f1b 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h @@ -2600,7 +2600,12 @@ template w, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool aligned_N, + bool bm16 = false> [[kernel]] void kq_iq4_nl_qmm_t_splitk( const device uint8_t* w, const device uint8_t* /* scales */, @@ -2617,7 +2622,8 @@ template uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(group_size == KQ_IQ4_NL_GROUP, "IQ4_NL requires gs=32"); static_assert(bits == 4, "IQ4_NL requires bits=4"); - constexpr int BM = 32, BK = 32, BN = 32; + constexpr int BM = bm16 ? 16 : 32; + constexpr int BK = 32, BN = bm16 ? 64 : 32; constexpr int BK_padded = (BK + 16 / sizeof(T)); threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index 6b2555b..5f37b85 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h @@ -3417,6 +3417,23 @@ KQ_NAX_DEFINE_KERNELS(q2_k, 256, 2, KqNaxQ2_KBlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(q6_k, 256, 6, KqNaxQ6_KBlockLoader) KQ_NAX_DEFINE_SPLITK_KERNEL(q8_0, 32, 8, KqNaxQ8_0BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q4_k, 256, 4, KqNaxQ4_KBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q5_k, 256, 5, KqNaxQ5_KBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q3_k, 256, 3, KqNaxQ3_KBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q2_k, 256, 2, KqNaxQ2_KBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q4_0, 32, 4, KqNaxQ4_0BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q4_1, 32, 4, KqNaxQ4_1BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q5_0, 32, 5, KqNaxQ5_0BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(q5_1, 32, 5, KqNaxQ5_1BlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq4_nl, 32, 4, KqNaxIq4_nlBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq4_xs, 256, 4, KqNaxIq4_xsBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq3_xxs, 256, 3, KqNaxIq3_xxsBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq3_s, 256, 3, KqNaxIq3_sBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq2_xxs, 256, 2, KqNaxIq2_xxsBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq2_xs, 256, 2, KqNaxIq2_xsBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq2_s, 256, 2, KqNaxIq2_sBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq1_s, 256, 1, KqNaxIq1_sBlockLoader) +KQ_NAX_DEFINE_SPLITK_KERNEL(iq1_m, 256, 1, KqNaxIq1_mBlockLoader) template < typename T, diff --git a/mlx_kquant/_version.py b/mlx_kquant/_version.py index df0ed33..8a3be2e 100644 --- a/mlx_kquant/_version.py +++ b/mlx_kquant/_version.py @@ -1 +1 @@ -__version__ = "0.3.12" +__version__ = "0.3.13" diff --git a/src/kquant_gather.cpp b/src/kquant_gather.cpp index 8dfac06..e6e11ea 100644 --- a/src/kquant_gather.cpp +++ b/src/kquant_gather.cpp @@ -14,6 +14,12 @@ // right_sorted_ == do_sort: MoE PREFILL takes this sorted per-expert GEMM // (~=6-8x faster than B separate gather_qmv vector-matmuls), while decode // (top_k<64 -> no sort -> B<16) falls through to gather_qmv. +// +// No split-K here, unlike the dense path: this grid already spans the +// active experts, and verify cost grows with M because more experts get +// touched, not from starvation. At [E=256, N=2048, K=7168, top_k=8] it +// holds 78-83% of roofline across M8-M32, so a split has nothing to +// recover. #include #include #include diff --git a/src/kquant_matmul.cpp b/src/kquant_matmul.cpp index 6df497f..b22f845 100644 --- a/src/kquant_matmul.cpp +++ b/src/kquant_matmul.cpp @@ -190,6 +190,120 @@ static int kq_smallm_route_min(const std::string& t, int N, int K) { return m; } +// Codecs with a split-K kernel and its bm16 tile arm. Gates both the +// route and the tile pick, so the two cannot drift apart. +static bool kq_splitk_codec(const std::string& t) { + return t == "q4_k" || t == "q5_k" || t == "q6_k" || t == "q3_k" || + t == "q2_k" || t == "q8_0" || t == "iq4_xs" || t == "iq4_nl" || + t == "iq3_xxs" || t == "iq3_s" || t == "iq2_xxs" || t == "iq2_xs" || + t == "iq2_s" || t == "iq1_s" || t == "iq1_m"; +} + +// Non-NAX default split-K entry M per codec (0 = env lever only). +// Below the entry, the mv paths win. From the entry through M32, +// qmm_splitk is flat at the measured shapes (the bm16 tile carries +// M <= 16; BM=32 carries M 17-32). The default route decays past M~4 +// and falls into the plain qmm_t hole at M13. +// +// Two device tables, picked at dispatch. On NAX silicon this route only +// runs with NAX off (KQ_DISABLE_NAX, or a codec with no NAX kernels); +// the band is otherwise served by kq_splitk_nax_min_m. That ALU path +// crosses later than pre-NAX silicon does. Measured 2026-08-13 on M5 +// Max under KQ_DISABLE_NAX=1, three shapes; wins from the entry are +// 2.0-2.3x. +static int kq_splitk_min_m_nax_alu(const std::string& t) { + if (t == "q4_k" || t == "q3_k" || t == "iq3_xxs" || t == "iq4_nl") { + return 8; + } + if (t == "q5_k") { + return 10; + } + if (t == "q6_k" || t == "iq3_s") { + return 12; + } + if (t == "q2_k" || t == "q8_0" || t == "iq4_xs" || t == "iq2_s") { + return 16; + } + return 0; +} + +// Pre-NAX silicon. Measured 2026-08 on M3 Max, bm16 tile: +// - q4_k, muse-glimmer MLP: bm16 flat 1.83-1.91 ms M2-16; mv_ext wins +// through M4 (1.05-1.42), loses from M6 (~2.3). Entry 6. +// - q5_k, muse-glimmer vocab head [6656x202k]: bm16 flat 5.9-6.1; +// default wins through M4 (4.4), loses from M6 (6.5). Entry 6. +// - q3_k, Qwen3-4B MLP: bm16 flat ~0.7 from M6. Entry 6. +// - q6_k, Qwen3-4B vocab head: bm16 flat 1.83-1.94; default 2.25 at +// M6, 3.93 at M8. Entry 6. +// - iq3_s, iq3_xxs and iq4_nl, same shape as q4_k: bm16 flat 1.5-1.9; +// the default wins through M4 (0.69-0.75x), ties at M6 (1.04-1.20x) +// and loses 2.2-2.7x at M16. Entry 6. iq3_s holds at a second shape +// [12288x4096] from real weights (M6 0.98x, M8 1.32x). +// - iq4_xs, same shape: its bm16 tile is slower (2.6 flat), so the +// default holds through M8 (0.90x) and the tile takes over at M10 +// (1.02x), 1.82x at M16. Entry 10. +// - iq2_s at [4096x12288] from real weights: bm16 flat ~1.5, the +// default crosses at M8 (1.07x) and loses 1.61x at M16. Entry 8. +// q2_k and q8_0 are not measured. They enter at the M13 cliff (q8_0's +// mv paths are the strongest in the fleet; measure before lowering). +// iq2_xxs, iq2_xs, iq1_s and iq1_m stay on the env lever: ggml refuses +// to encode them without an importance matrix, so no synthetic weights +// exist, and the local imatrix models hold iq2_xxs only as 3D expert +// tensors, which this route never takes. +static int kq_splitk_min_m(const std::string& t) { + if (kq_is_nax_available()) { + return kq_splitk_min_m_nax_alu(t); + } + if (t == "q4_k" || t == "q3_k" || t == "q5_k" || t == "q6_k" || + t == "iq3_s" || t == "iq3_xxs" || t == "iq4_nl") { + return 6; + } + if (t == "iq2_s") { + return 8; + } + if (t == "iq4_xs") { + return 10; + } + if (t == "q2_k" || t == "q8_0") { + return 13; + } + return 0; +} + +// Default NAX split-K entry M per codec (0 = env lever only). The +// un-split BM=32 grid is ceil(N/64) x 1 threadgroups with a serial +// in-tile K walk, so the verify band pays per row; splitting K +// multiplies threadgroup count. The win scales with how starved the +// grid is: 1.6-2.5x at [256x6656], 1.0-1.2x at vocab-head N. +// +// Entry is the lowest M with no regression on any measured shape. +// Measured 2026-08-13 on M5 Max at target 16 over seven shapes +// (benchmarks/bench_verify_band_ab.py). No N gate needed: vocab-head +// shapes stay >= 1.0 at every entry. +static int kq_splitk_nax_min_m(const std::string& t) { + if (t == "iq2_xs" || t == "iq2_s") { + return 16; + } + if (t == "q6_k" || t == "iq1_m") { + return 12; + } + if (t == "iq3_s" || t == "iq2_xxs" || t == "iq1_s") { + return 10; + } + if (t == "q2_k" || t == "q3_k" || t == "q4_k" || t == "q5_k" || t == "q8_0" || + t == "q4_0" || t == "q4_1" || t == "q5_0" || t == "q5_1" || + t == "iq4_nl" || t == "iq4_xs" || t == "iq3_xxs") { + return 8; + } + return 0; +} + +// Default split target; resolves down to a divisor of K / max(gs, BK), +// so the realised count is coarser (K=6656 and K=19968 both give 13). +// One target for all codecs: over routed cells only, 16 is best for 15 +// of 19 and within 1% for the rest. +static constexpr int kq_splitk_nax_target = 16; + // NAX (tensor-core) GEMM dispatch for the quantized matmul (no biases). void qmm_nax( const array& x, @@ -430,8 +544,12 @@ void qmm_splitk( Device& d, const Stream& s, const std::string& kquant_type) { - constexpr int bm = 32, bn = 32; constexpr int wm = 2, wn = 2; + // The bm16 tile halves the MMA row-padding waste at small M and + // widens BN to 64 so each K-step does more MMA work per barrier. + const bool bm16 = M <= 16 && kq_splitk_codec(kquant_type); + const int bm = bm16 ? 16 : 32; + const int bn = bm16 ? 64 : 32; const int k_partition = (K / group_size / splits) * group_size; const int part_stride = M * N; @@ -447,7 +565,8 @@ void qmm_splitk( kname.reserve(64); mx::concatenate( kname, - kq_kname_prefix(kquant_type) + "qmm_t_splitk_", + kq_kname_prefix(kquant_type) + + (bm16 ? "qmm_t_splitk_bm16_" : "qmm_t_splitk_"), type_string, "_gs_", group_size, @@ -856,7 +975,8 @@ void verify_mv_ext( const char* e = std::getenv("KQ_MV_EXT_NR"); return e != nullptr ? std::atoi(e) : 1; }(); - const bool use_nr2 = mv_ext_nr == 2 && M >= 5 && kquant_type == "q6_k"; + const bool use_nr2 = + mv_ext_nr == 2 && M >= 5 && M <= 12 && kquant_type == "q6_k"; // Shuffle-broadcast experiment (KQ_MV_EXT_SB=1): ty-lanes exchange // activation quarters over simd_shuffle instead of each loading the full // window -- activation cache traffic / 4, same grid. q6_k M 4-12 only. @@ -864,7 +984,8 @@ void verify_mv_ext( const char* e = std::getenv("KQ_MV_EXT_SB"); return e != nullptr && std::atoi(e) == 1; }(); - const bool use_sb = !use_nr2 && mv_ext_sb && M >= 4 && kquant_type == "q6_k"; + const bool use_sb = + !use_nr2 && mv_ext_sb && M >= 4 && M <= 12 && kquant_type == "q6_k"; // Wide-nxpsg experiment (KQ_MV_EXT_NX=16|32): fewer redundant activation // readers per element (nypsg*nsg drops 8 -> 4 -> 2) + more threadgroups. // q6_k M 4-12 only; wins here would generalize per codec. @@ -873,8 +994,8 @@ void verify_mv_ext( const int v = e != nullptr ? std::atoi(e) : 0; return (v == 16 || v == 32) ? v : 0; }(); - const bool use_nx = - !use_nr2 && !use_sb && mv_ext_nx != 0 && M >= 4 && kquant_type == "q6_k"; + const bool use_nx = !use_nr2 && !use_sb && mv_ext_nx != 0 && M >= 4 && + M <= 12 && kquant_type == "q6_k"; // T-precision-dot experiment (KQ_MV_EXT_HD=1): chunk dots in half/bfloat // at 2x issue rate + no per-row activation converts, f32 fold per chunk. // q6_k M 4-12, half/bfloat x only. @@ -883,17 +1004,8 @@ void verify_mv_ext( return e != nullptr && std::atoi(e) == 1; }(); const bool use_hd = !use_nr2 && !use_sb && !use_nx && mv_ext_hd && M >= 4 && - kquant_type == "q6_k" && x.dtype() != mx::float32; - // Staged-activation experiment (KQ_MV_EXT_TS=1): cooperative TG-memory - // stage of the activation window, 8 simdgroups / 32 rows per TG. q6_k - // M 4-12; K must cover a full 128-element window (q6_k geometry does). - static const bool mv_ext_ts = []() { - const char* e = std::getenv("KQ_MV_EXT_TS"); - return e != nullptr && std::atoi(e) == 1; - }(); - const bool use_ts = !use_nr2 && !use_sb && !use_nx && !use_hd && mv_ext_ts && - M >= 4 && kquant_type == "q6_k" && K % 128 == 0; - const int nsg_eff = use_ts ? 8 : nsg; + M <= 12 && kquant_type == "q6_k" && x.dtype() != mx::float32; + const int nsg_eff = nsg; const int nxpsg_eff = use_nx ? mv_ext_nx : nxpsg; const int rows_per_tg = (32 / nxpsg_eff) * nsg_eff * (use_nr2 ? 2 : 1); MTL::Size group_dims(32, nsg_eff, 1); @@ -915,7 +1027,7 @@ void verify_mv_ext( use_nr2 ? "_nr2" : (use_sb ? "_sb" : (use_nx ? (mv_ext_nx == 16 ? "_x16" : "_x32") - : (use_hd ? "_hd" : (use_ts ? "_ts" : ""))))); + : (use_hd ? "_hd" : "")))); auto kernel = kq_get_kernel(d, kname); auto& ce = mx::metal::get_command_encoder(s); @@ -1136,19 +1248,22 @@ void KQuantMatmul::eval_gpu( return; } - // NAX split-K qmm experiment (KQ_QMM_SPLITK_NAX=, 0 = off, - // 1 = auto target 32; read LIVE per call so --ab-env can flip it on one - // generator): K-slices on the tensor-core BM=32 tile; see qmm_nax_splitk. - // Slice quantum is max(superblock, BK) so every slice starts a loader at - // kt_base 0. q6_k + q8_0 instantiations only. + // NAX split-K, entering at kq_splitk_nax_min_m through M32; see + // qmm_nax_splitk. Slice quantum max(superblock, BK) keeps every slice + // starting a loader at kt_base 0 for both gs families. + // KQ_QMM_SPLITK_NAX= forces the route for all M <= 32, 1 uses + // the default target, 0 disables. Read live so an A/B can flip arms + // on one generator. const char* sk_nax_e = std::getenv("KQ_QMM_SPLITK_NAX"); - int qmm_splitk_nax_env = sk_nax_e != nullptr ? std::atoi(sk_nax_e) : 0; - if (qmm_splitk_nax_env == 1) { - qmm_splitk_nax_env = 32; - } - if (qmm_splitk_nax_env > 1 && transpose_ && non_batched && M <= 32 && - (kquant_type_ == "q6_k" || kquant_type_ == "q8_0") && - kq_is_nax_available() && (K % 64 == 0) && x.dtype() != mx::float32) { + const int sk_nax_env = sk_nax_e != nullptr ? std::atoi(sk_nax_e) : -1; + const int sk_nax_min_m = kq_splitk_nax_min_m(kquant_type_); + const bool sk_nax_route = sk_nax_env >= 1 || + (sk_nax_env == -1 && sk_nax_min_m > 0 && M >= sk_nax_min_m); + if (sk_nax_route && transpose_ && non_batched && M <= 32 && + codec_has_nax(kquant_type_) && kq_is_nax_available() && (K % 64 == 0) && + x.dtype() != mx::float32) { + const int qmm_splitk_nax_env = + sk_nax_env > 1 ? sk_nax_env : kq_splitk_nax_target; const int sliceq = std::max(group_size_, 64); const int nblk = K / sliceq; int sp = std::min(qmm_splitk_nax_env, nblk); @@ -1174,20 +1289,29 @@ void KQuantMatmul::eval_gpu( } } - // Split-K qmm experiment (KQ_QMM_SPLITK=, 0 = off, read - // once): the small-M band's occupancy lever; see qmm_splitk. Routes when - // a >1 divisor of the wire-block count exists at or under the target. - // K-quants + q8_0 only for now (instantiation coverage). + // Split-K qmm: the occupancy lever for the small-M band; see + // qmm_splitk. The route needs a >1 divisor of the wire-block count + // at or under the target. K-quants + q8_0 only for now + // (instantiation coverage). Default routing is non-NAX only and + // enters at kq_splitk_min_m. KQ_QMM_SPLITK= forces + // the route for all M <= 32 (A/B lever). KQ_QMM_SPLITK=0 disables + // both. static const int qmm_splitk_env = []() { const char* e = std::getenv("KQ_QMM_SPLITK"); - return e != nullptr ? std::atoi(e) : 0; + return e != nullptr ? std::atoi(e) : -1; // -1 = per-codec default }(); - if (qmm_splitk_env > 1 && transpose_ && non_batched && M <= 32 && - (kquant_type_ == "q6_k" || kquant_type_ == "q5_k" || - kquant_type_ == "q4_k" || kquant_type_ == "q3_k" || - kquant_type_ == "q2_k" || kquant_type_ == "q8_0")) { + const int splitk_min_m = kq_splitk_min_m(kquant_type_); + // Hardware AND codec: KQ_DISABLE_NAX lands in codec_has_nax, so + // keying off availability alone left this unreachable on NAX silicon. + const bool nax_path = kq_is_nax_available() && codec_has_nax(kquant_type_); + const bool splitk_route = qmm_splitk_env > 1 || + (qmm_splitk_env == -1 && splitk_min_m > 0 && M >= splitk_min_m && + !nax_path); + if (splitk_route && transpose_ && non_batched && M <= 32 && + kq_splitk_codec(kquant_type_)) { + const int splitk_target = qmm_splitk_env > 1 ? qmm_splitk_env : 16; const int nblk = K / group_size_; - int sp = std::min(qmm_splitk_env, nblk); + int sp = std::min(splitk_target, nblk); while (sp > 1 && nblk % sp != 0) { --sp; } diff --git a/tests/test_nax_smallm.py b/tests/test_nax_smallm.py index bce22cc..64c1a82 100644 --- a/tests/test_nax_smallm.py +++ b/tests/test_nax_smallm.py @@ -142,3 +142,27 @@ def test_db64_band_dispatch(codec): def test_smallm_routing_iq(codec, n_out): w, s, ref_w = _iq_setup(codec, n_out) _sweep(codec, w, s, ref_w, n_out) + + +# Non-NAX split-K band, which the rest of this file cannot reach on NAX +# silicon: M 2 sits below every codec entry, 8-32 spans the routed band up +# to its M <= 32 ceiling, and 33 is the handoff back to plain qmm. +# KQ_DISABLE_NAX is read live, so toggling it re-routes in-process. +ALU_SPLITK_MS = [2, 8, 10, 12, 16, 17, 24, 32, 33] + + +@pytest.fixture +def nax_off(monkeypatch): + monkeypatch.setenv("KQ_DISABLE_NAX", "1") + + +@pytest.mark.parametrize("codec", ENCODABLE) +def test_alu_splitk_band(codec, nax_off): + w, s, ref_w = _encodable_setup(codec, 1000) + _sweep(codec, w, s, ref_w, 1000, ms=ALU_SPLITK_MS) + + +@pytest.mark.parametrize("codec", IQ) +def test_alu_splitk_band_iq(codec, nax_off): + w, s, ref_w = _iq_setup(codec, 1000) + _sweep(codec, w, s, ref_w, 1000, ms=ALU_SPLITK_MS)