From d33c891768184c6aeadeb3f6ebcc342e4355e629 Mon Sep 17 00:00:00 2001 From: AndreSlavescu Date: Wed, 29 Jul 2026 21:03:37 -0700 Subject: [PATCH] Measure the int8 multi-row headroom, and keep the four-bit guard that was protecting us The int4 multi-row win reproduces and is stronger than published: 1.45x, 1.73x and 1.54x at rows 8, 16 and 32 through the prepared executor. The int8 half of this task does not exist, and finding that out was more instructive than the number. Why int4 wins is bandwidth, not batching. MLX re-reads weights per row tile, and at four bits that drops its effective weight-read bandwidth to 36 GB/s at sixteen rows against a measured 120 GB/s streaming ceiling -- a 3.4x gap, of which the kernel captures about half. At eight bits MLX runs 62 GB/s, a 1.9x gap. benchmarks/quantized_row_headroom.py measures this per width so the question "should we build a multi-row kernel for N bits" is answerable before anyone builds one. Effective bandwidth rather than time, because eight-bit weights are twice the bytes and a ratio of times cannot tell a slower kernel from a worse use of the memory system. Half of the eight-bit gap would be around 1.4x, which is worth wanting, so I tried it. The repacking generalises to eight bits in a few lines, the weights load, results came back and the kernel measured 1.6x to 2.0x faster than MLX. It was also wrong, at a relative error of 2.5 to 2.9, and the speed and the wrongness are the same fact: `lower_affine_matmul` emits NAX affine fragments with block_size=4 and takes no bit width, so it decoded eight nibbles per word where the data held four bytes, read half the values, and finished early. A large speedup arriving together with a wrong answer is one bug, not one win and one bug. Two things worth keeping from that. The four-bit check in MLXAffineWeight.from_mlx is load-bearing rather than a formatting preference, and I loosened it without checking that anything downstream honoured the parameter -- the bit width was threaded into cache keys but never into the lowering. It is restored, its message now says why, and a test pins it, because relaxing it looks harmless and the failure is silent. Only the tuner's agreement gate stopped the fast wrong kernel being selected, which is the gate doing exactly its job and not a reason to rely on it. The generalised repacking is reverted too. It was correct, but a correct-looking eight-bit path sitting beside a four-bit-only lowering is precisely the trap that just caught me. Supporting eight bits means teaching the lowering a bit width, which is real compiler work against a 1.4x prize, not a check to relax. 680 pass, 7 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- benchmarks/quantized_row_headroom.py | 114 +++++++++++++++++++++++++++ metile/backends/mlx_affine.py | 16 +++- metile/backends/mlx_quantized.py | 8 +- tests/test_config_admissibility.py | 29 +++++++ 4 files changed, 165 insertions(+), 2 deletions(-) create mode 100644 benchmarks/quantized_row_headroom.py diff --git a/benchmarks/quantized_row_headroom.py b/benchmarks/quantized_row_headroom.py new file mode 100644 index 0000000..b6d8a49 --- /dev/null +++ b/benchmarks/quantized_row_headroom.py @@ -0,0 +1,114 @@ +"""How much room MLX leaves at each weight width, which is what decides where a kernel can win. + +The multi-row win at four bits is not a general property of batching quantized decode. It exists +because MLX re-reads weights per row tile and, at four bits, that drops its effective weight-read +bandwidth to about a third of what the part can stream. This measures that directly for each width +so the question "should we build a multi-row kernel for N bits" has an answer before anyone builds +one. + +What it shows on M5, against a measured streaming ceiling near 120 GB/s: + + rows int4 int8 + 1 60 83 + 8 42 65 + 16 36 62 + 32 37 62 + +So four bits leaves roughly 3.3x on the table at sixteen rows and eight bits roughly 1.9x. The +four-bit kernel captures about half of its share, measuring 1.45x to 1.73x at rows 8 to 32. Half of +the eight-bit share would be around 1.4x, which is worth wanting, and it is not reachable through +the same path: the matrix-unit affine fragment format is four bits wide and `lower_affine_matmul` +has no bit width to thread through it. + +Reading effective bandwidth rather than time is what makes the widths comparable. Eight-bit weights +are twice the bytes, so a slower wall time can still be the better use of the memory system, and a +ratio of times cannot tell the two apart. +""" + +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 + +HIDDEN, INTERMEDIATE = 1536, 8960 +GROUP = 64 +ROWS = (1, 2, 4, 8, 16, 32, 64) + + +def _arguments(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rounds", type=int, default=25) + parser.add_argument("--hidden", type=int, default=HIDDEN) + parser.add_argument("--intermediate", type=int, default=INTERMEDIATE) + return parser.parse_args() + + +def _bandwidth(mx, bits, rows, hidden, intermediate, rounds, inner=3): + """MLX's effective weight-read bandwidth in GB/s for one width and row count.""" + mx.random.seed(0) + dense = mx.random.normal((intermediate, hidden)).astype(mx.float16) + packed, scales, biases = mx.quantize(dense, group_size=GROUP, bits=bits, mode="affine") + activations = mx.random.normal((rows, hidden)).astype(mx.float16) + mx.eval(packed, scales, biases, activations) + + def run(): + return mx.quantized_matmul( + activations, + packed, + scales=scales, + biases=biases, + transpose=True, + group_size=GROUP, + bits=bits, + mode="affine", + ) + + for _ in range(3): + mx.eval([run() for _ in range(inner)]) + mx.synchronize() + samples = [] + for _ in range(rounds): + started = time.perf_counter_ns() + mx.eval([run() for _ in range(inner)]) + samples.append((time.perf_counter_ns() - started) / inner / 1e9) + + # Weights plus one scale and one bias per group, which is what a single pass has to read. + parameters = intermediate * (hidden // GROUP) * 2 * 2 + weight_bytes = intermediate * hidden * bits / 8 + parameters + return weight_bytes / statistics.median(samples) / 1e9 + + +def main(): + arguments = _arguments() + try: + import mlx.core as mx + except ImportError: + print("mlx is required") + return 1 + + ceiling = agx.STREAMING_READ_GBPS + print("MLX's effective weight-read bandwidth, GB/s") + print(f"shape {arguments.hidden}x{arguments.intermediate}, group {GROUP}, ") + print(f"measured streaming ceiling {ceiling} GB/s\n") + print(f"{'rows':>6}{'int4':>9}{'int8':>9}{'int4 gap':>11}{'int8 gap':>11}") + + for rows in ROWS: + four = _bandwidth(mx, 4, rows, arguments.hidden, arguments.intermediate, arguments.rounds) + eight = _bandwidth(mx, 8, rows, arguments.hidden, arguments.intermediate, arguments.rounds) + print( + f"{rows:>6}{four:>9.1f}{eight:>9.1f}{ceiling / four:>10.2f}x{ceiling / eight:>10.2f}x" + ) + + print("\nThe gap columns bound what any kernel could win at that width and row count.") + print("A width whose gap is near 1.00x has nothing to give however the kernel is written.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metile/backends/mlx_affine.py b/metile/backends/mlx_affine.py index 19f2644..a2db544 100644 --- a/metile/backends/mlx_affine.py +++ b/metile/backends/mlx_affine.py @@ -79,8 +79,22 @@ class MLXAffineWeight: @classmethod def from_mlx(cls, weight, scales, biases, *, group_size=64, bits=4): + # Four bits is not a formatting preference, it is the operand format the matrix unit + # takes. `lower_affine_matmul` builds NAX affine fragments with block_size=4 and has no + # bit-width parameter to thread anything else through. + # + # Relaxing this to accept eight bits was tried and is a trap worth naming. The repacking + # generalises cleanly, the weights load, and the kernel runs 1.6x to 2.0x faster than MLX -- + # because it decodes eight nibbles per word where the data holds four bytes, so it reads + # half the values and returns garbage, at a relative error of 2.5 to 2.9. A large speedup + # arriving together with a wrong answer is one bug, not one win and one bug. Only the + # tuner's agreement gate kept it out of a selection. if group_size != 64 or bits != 4: - raise ValueError("MLX affine NAX weights require group size 64 and 4 bits") + raise ValueError( + f"MLX affine NAX weights require group size 64 and 4 bits, got {group_size} and " + f"{bits}. The matrix-unit affine fragment format is 4-bit; supporting another " + f"width means teaching lower_affine_matmul the bit width, not relaxing this check." + ) packed, repacked_scales, repacked_biases = repack_mlx_affine_weight( weight, scales, diff --git a/metile/backends/mlx_quantized.py b/metile/backends/mlx_quantized.py index bc8d853..81a91b9 100644 --- a/metile/backends/mlx_quantized.py +++ b/metile/backends/mlx_quantized.py @@ -752,7 +752,13 @@ def mlx_affine_swiglu_qmv( def repack_mlx_affine_weight(weight, scales, biases): - """Repack MLX output-major affine uint4 weights into K-major NAX layout.""" + """Repack MLX output-major affine uint4 weights into K-major NAX layout. + + Four bits throughout, matching the matrix unit's affine fragment format. Generalising this to + eight is a few lines and does not help: the consumer, `lower_affine_matmul`, emits NAX affine + fragments with block_size=4 and takes no bit width, so a wider repack only produces weights the + kernel will decode as nibbles. See the guard in MLXAffineWeight.from_mlx. + """ import mlx.core as mx if biases is None: diff --git a/tests/test_config_admissibility.py b/tests/test_config_admissibility.py index 72879e0..58c9c25 100644 --- a/tests/test_config_admissibility.py +++ b/tests/test_config_admissibility.py @@ -73,3 +73,32 @@ def test_attention_decode_matches_mlx_at_every_head_dimension(dimension): reference = mx.fast.scaled_dot_product_attention(query, keys, values, scale=scale) mx.eval(got, reference) assert mx.allclose(got, reference, rtol=2e-3, atol=2e-3).item() + + +def test_affine_nax_weights_reject_widths_the_matrix_unit_cannot_decode(): + """The 4-bit guard is load-bearing, not a formatting preference. + + `lower_affine_matmul` emits NAX affine fragments with block_size=4 and has no bit-width + parameter, so an 8-bit weight loaded through here is decoded as nibbles: the kernel reads eight + values per word where the data holds four, returns a relative error of 2.5 to 2.9, and runs + 1.6x to 2.0x faster than MLX precisely because it does half the work. A large speedup arriving + with a wrong answer is one bug, not one win and one bug. + + Pinned as a test because relaxing the check looks harmless and the failure is silent: only the + tuner's agreement gate stopped the wrong kernel being selected. + """ + mx = pytest.importorskip("mlx.core") + + from metile.backends.mlx_affine import MLXAffineWeight + + dense = mx.random.normal((256, 128)).astype(mx.float16) + for bits in (2, 8): + packed, scales, biases = mx.quantize(dense, group_size=64, bits=bits, mode="affine") + mx.eval(packed, scales, biases) + with pytest.raises(ValueError, match="4 bits"): + MLXAffineWeight.from_mlx(packed, scales, biases, group_size=64, bits=bits) + + packed, scales, biases = mx.quantize(dense, group_size=64, bits=4, mode="affine") + mx.eval(packed, scales, biases) + weight = MLXAffineWeight.from_mlx(packed, scales, biases, group_size=64, bits=4) + assert weight.bits == 4