From 05411c701bbaac7fda90370ea42283c6ea72ecbf Mon Sep 17 00:00:00 2001 From: AndreSlavescu Date: Wed, 29 Jul 2026 21:07:34 -0700 Subject: [PATCH] Measure attention projections instead of serving them, and record the amortisation trap The q, k, v and o projections are about a quarter of a decode token and have no meTile path, so they look like the obvious next thing to serve. At matched bf16 they are not worth serving, and the reason is the same one the standing goal already records: MLX's decode GEMVs run at the hardware streaming limit. At one row MLX reads these weights at 121 to 254 GB/s against a 120.6 GB/s streaming ceiling, so for three of four shapes the weights are cache-resident and DRAM is not even the constraint. meTile measures 0.993x, 0.721x, 1.007x and 0.996x there, and 0.93x to 1.03x at eight rows. A win appears only at thirty-two rows, 1.33x to 1.56x, which is batch or prefill rather than decode, and two of the four shapes stop being bit-exact there. Routing that would trade the logit-equality contract for a win outside the case the task was about. So this closes as measured rather than built. Adding the path would deliver parity at decode along with a code path, a tuning surface and a bit-exactness risk. The measurement is the deliverable, and getting it wrong was instructive enough to encode. A first version sized the inner dispatch count for 64 MB of weight traffic per eval. For a 4.7 MB projection that is thirteen dispatches, about 390 us of work against a roughly 200 us mx.eval round trip, so nearly a third of every sample was overhead -- and because the two sides pay it differently it reported MLX at 52 GB/s with a 2.3x gap and meTile at 1.828x. Both were fiction. Targeting a gigabyte per eval puts the round trip near 2%, and the same shape then reads 146 GB/s at 0.993x. That is the third time this session that insufficient amortisation manufactured a win: module reloads inside a timing loop reported the scheduling pass at 0.83x to 0.92x when it was measuring the Metal compiler, and a per-call API reported the int4 multi-row QMV at 0.05x when the published path uses a prepared executor. The rule that would have caught all three: if a measurement of a small kernel shows a large difference, suspect the harness before believing it. benchmarks/projection_headroom.py takes the traffic target as a flag, so lowering it reproduces the error rather than leaving the claim to be trusted. 680 pass, 7 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- benchmarks/projection_headroom.py | 136 ++++++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 benchmarks/projection_headroom.py diff --git a/benchmarks/projection_headroom.py b/benchmarks/projection_headroom.py new file mode 100644 index 0000000..bb3184f --- /dev/null +++ b/benchmarks/projection_headroom.py @@ -0,0 +1,136 @@ +"""Whether attention projections have anything to give, at matched bf16. + +The q, k, v and o projections are about a quarter of a decode token and have no meTile path, which +makes them look like the obvious next thing to serve. They are not, and this measures why rather than +arguing it: at one row MLX already runs them at or above the streaming ceiling, so there is nothing +for a kernel to win. + + shape rows MLX GB/s gap meTile bit-exact + Qwen2.5-1.5B q/o 1 148 0.8x 0.967x yes + Qwen2.5-1.5B k/v 1 242 0.5x 0.780x yes + Qwen3-8B q/o 1 120 1.0x 0.990x yes + Qwen3-8B k/v 1 148 0.8x 0.991x yes + +A gap below 1.0x means the weights are cache-resident and the DRAM ceiling is not even the limit. At +eight rows it is still parity. Only at thirty-two rows does a win appear, 1.30x to 1.57x, and that is +batch or prefill rather than decode; two of the four shapes stop being bit-exact there, so routing it +would trade the logit-equality contract for a win outside the case the task was about. + +Amortisation is the whole methodology here and it is easy to get wrong in the direction of good news. +An earlier version of this measurement used enough inner dispatches for 64 MB of weight traffic per +eval. For a 4.7 MB projection that is thirteen dispatches, about 390 us of work against a roughly +200 us `mx.eval` round trip, so nearly a third of every sample was overhead -- and because the two +sides pay it differently it reported MLX at 52 GB/s and meTile at 1.828x. Both were fiction. Targeting +a gigabyte per eval puts the round trip near 2% and the same shape reads 148 GB/s at 0.967x. + +If a measurement of a small kernel shows a large win, suspect the harness first. +""" + +import argparse +import statistics +import sys +import time +from pathlib import Path + +_root = str(Path(__file__).resolve().parent.parent) +sys.path.insert(0, _root) + +from metile.target import agx + +# Real projection shapes: (label, hidden, output features). Grouped-query attention makes k and v much +# narrower than q and o, and the two behave differently, so both are measured. +SHAPES = ( + ("Qwen2.5-1.5B q/o", 1536, 1536), + ("Qwen2.5-1.5B k/v", 1536, 256), + ("Qwen3-8B q/o", 4096, 4096), + ("Qwen3-8B k/v", 4096, 1024), +) +ROWS = (1, 8, 32) + +# Weight traffic per eval. The round trip is fixed, so this decides the error floor: a gigabyte is +# about 8ms at the streaming ceiling, putting a 200us round trip at 2.5%. +TRAFFIC_TARGET = 1e9 + + +def _arguments(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rounds", type=int, default=9) + parser.add_argument( + "--traffic", + type=float, + default=TRAFFIC_TARGET, + help="weight bytes per eval; lowering this reintroduces round-trip error", + ) + return parser.parse_args() + + +def _inner(hidden, output_features, traffic): + return max(16, min(1024, int(traffic // (output_features * hidden * 2)))) + + +def _median(mx, build, inner, rounds): + for _ in range(3): + mx.eval([build() for _ in range(inner)]) + mx.synchronize() + samples = [] + for _ in range(rounds): + started = time.perf_counter_ns() + mx.eval([build() for _ in range(inner)]) + samples.append((time.perf_counter_ns() - started) / inner / 1e9) + return statistics.median(samples) + + +def main(): + arguments = _arguments() + try: + import mlx.core as mx + except ImportError: + print("mlx is required") + return 1 + + from metile.backends.mlx_dense import MLXDenseWeight, mlx_dense_matmul + + ceiling = agx.STREAMING_READ_GBPS + print(f"attention projections, matched bf16, streaming ceiling {ceiling} GB/s") + print(f"{int(arguments.traffic / 1e6)} MB of weight traffic per eval\n") + header = ( + f"{'shape':<20}{'rows':>5}{'inner':>7}{'MLX GB/s':>10}{'gap':>7}{'meTile':>9}{'exact':>7}" + ) + print(header) + print("-" * len(header)) + + for label, hidden, output_features in SHAPES: + inner = _inner(hidden, output_features, arguments.traffic) + for rows in ROWS: + mx.random.seed(0) + dense = mx.random.normal((output_features, hidden)).astype(mx.bfloat16) + values = mx.random.normal((rows, hidden)).astype(mx.bfloat16) + mx.eval(dense, values) + weight = MLXDenseWeight.from_mlx(dense) + mx.eval(weight.k_major) + + def native(values=values, dense=dense): + return values @ dense.T + + def generated(values=values, weight=weight): + return mlx_dense_matmul(values, weight) + + produced, reference = generated(), native() + mx.eval(produced, reference) + exact = bool(mx.array_equal(produced, reference).item()) + + base = _median(mx, native, inner, arguments.rounds) + ours = _median(mx, generated, inner, arguments.rounds) + gbps = (output_features * hidden * 2) / base / 1e9 + print( + f"{label:<20}{rows:>5}{inner:>7}{gbps:>10.0f}{ceiling / gbps:>6.1f}x" + f"{base / ours:>8.3f}x{'yes' if exact else 'NO':>7}" + ) + + print("\ngap is the ceiling divided by what MLX achieves: what any kernel could win.") + print("Below 1.0x the weights are cache-resident and DRAM is not the limit at all.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())