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