Measure the int8 multi-row headroom, and keep the four-bit guard that was protecting us

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())

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,

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:

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