"""Measure bandwidth and matrix-multiply throughput with MLX on Apple silicon, beside PyTorch. Purpose: run the same two measurements through Apple's own array framework, so that the Track M reader has two numbers per quantity for one machine and can see how much of the gap to the specification is the framework rather than the hardware; and time MLX's quantised matrix multiply, the kernel mlx-lm runs for 4-bit and 8-bit models, at four lower precisions Platform: mac (Apple silicon; the lab runs it only on Track M) Minimum memory: 8 GB Assumes: mlx is installed in the active environment (Part 1's lab installs it); arrays live in unified memory, so the buffer size is bounded by the machine's memory and by the wired-memory limit rather than by a separate GPU memory pool Usage: python3 mlx-bandwidth-test.py [--size-mb 512] [--iters 20] [--repeats 3] [--matmul-size 4096] [--matmul-iters 30] [--max-seconds 30] [--labbook labbook.md] Method. MLX is lazy: nothing is computed until mx.eval asks for it, and mx.eval blocks until it is done, so every timed operation ends in mx.eval. 1 GB = 1e9 bytes. pass mx.eval(src * 0.5) over an N-byte float32 array reads N bytes and writes N: (2 x N x iters) / seconds, the convention of bandwidth-test.py's copy test read mx.eval(mx.matmul(W, x)) with W a square float32 matrix of about N bytes: (bytes of W x iters) / seconds, the convention of bandwidth-test.py's read test matmul A @ B for two N x N matrices, counted as 2 x N^3 operations, in float32, bfloat16 and float16 through mx.matmul; then with B quantised by mx.quantize and multiplied by mx.quantized_matmul (A stays float16) in four modes: affine 8-bit and 4-bit with group size 64, mxfp8 and mxfp4 with group size 32 Each matmul format gets one untimed warm-up and one timed multiplication to size the loop to --max-seconds, never fewer than three multiplications. """ import argparse import json import statistics import sys import time from datetime import datetime, timezone from pathlib import Path import mlx.core as mx GB = 1_000_000_000 MIB = 1024 * 1024 FLOAT_FORMATS = {"float32": mx.float32, "bfloat16": mx.bfloat16, "float16": mx.float16} QUANT_FORMATS = { # name: (mode, bits, group_size), all documented for mx.quantize "affine8": ("affine", 8, 64), "affine4": ("affine", 4, 64), "mxfp8": ("mxfp8", 8, 32), "mxfp4": ("mxfp4", 4, 32), } INFO_KEYS = ("architecture", "device_name", "memory_size", "max_recommended_working_set_size", "max_buffer_length") def timed(fn, count): start = time.perf_counter() for _ in range(count): fn() return time.perf_counter() - start def bandwidth(size_mb, iters, repeats): nbytes = size_mb * MIB src = mx.random.uniform(shape=(nbytes // 4,), dtype=mx.float32) mx.eval(src) for _ in range(3): mx.eval(src * 0.5) passes = [2 * nbytes * iters / timed(lambda: mx.eval(src * 0.5), iters) / GB for _ in range(repeats)] del src side = int((nbytes / 4) ** 0.5) w = mx.random.uniform(shape=(side, side), dtype=mx.float32) x = mx.random.uniform(shape=(side,), dtype=mx.float32) mx.eval(w, x) for _ in range(3): mx.eval(mx.matmul(w, x)) reads = [side * side * 4 * iters / timed(lambda: mx.eval(mx.matmul(w, x)), iters) / GB for _ in range(repeats)] del w, x mx.clear_cache() return passes, reads, side def multiply_fn(name, size): a = mx.random.normal(shape=(size, size), dtype=mx.float32) b = mx.random.normal(shape=(size, size), dtype=mx.float32) if name in FLOAT_FORMATS: a, b = a.astype(FLOAT_FORMATS[name]), b.astype(FLOAT_FORMATS[name]) mx.eval(a, b) return lambda: mx.eval(mx.matmul(a, b)) mode, bits, group = QUANT_FORMATS[name] a = a.astype(mx.float16) packed = mx.quantize(b.astype(mx.float16), group_size=group, bits=bits, mode=mode) wq, scales = packed[0], packed[1] biases = packed[2] if len(packed) > 2 else None # affine returns biases, mxfp does not mx.eval(a, *packed) return lambda: mx.eval(mx.quantized_matmul(a, wq, scales, biases, transpose=True, group_size=group, bits=bits, mode=mode)) def matmul(name, size, iters, max_seconds): fn = multiply_fn(name, size) fn() # warm-up: Metal kernel compilation and allocation one = timed(fn, 1) runs = max(min(iters, 3), min(iters, int(max_seconds / max(one, 1e-9)))) seconds = timed(fn, runs) del fn mx.clear_cache() return 2.0 * size ** 3 * runs / seconds / 1e12, runs, seconds def main(): parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--size-mb", type=int, default=512) parser.add_argument("--iters", type=int, default=20) parser.add_argument("--repeats", type=int, default=3) parser.add_argument("--matmul-size", type=int, default=4096, help="square matrix side; must be divisible by 64") parser.add_argument("--matmul-iters", type=int, default=30) parser.add_argument("--max-seconds", type=float, default=30.0) parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file") args = parser.parse_args() if args.matmul_size % 64 or min(args.size_mb, args.iters, args.repeats, args.matmul_iters) < 1: sys.exit("mlx-bandwidth-test: --matmul-size must be a positive multiple of 64 and " "the counts at least 1") print(f"mlx {mx.__version__}") print(f"default device: {mx.default_device()}") info = mx.device_info() shown = {k: info[k] for k in INFO_KEYS if k in info} for key, value in shown.items(): if isinstance(value, int) and value > 1024 ** 3: print(f"{key}: {value} bytes ({value / 1024 ** 3:.1f} GiB)") else: print(f"{key}: {value}") passes, reads, side = bandwidth(args.size_mb, args.iters, args.repeats) print(f"\nbuffer {args.size_mb} MiB; {args.iters} operations x {args.repeats} repeats; " f"1 GB = 1e9 bytes") for label, runs in (("pass", passes), (f"read ({side}^2 x 4)", reads)): median = statistics.median(runs) spread = (max(runs) - min(runs)) / median * 100.0 print(f" mlx {label:<22} {median:8.1f} GB/s spread {spread:.1f}%") print(f"\nmatrices: {args.matmul_size} x {args.matmul_size}; up to {args.matmul_iters} " f"multiplications or {args.max_seconds:.0f} s per format") tflops = {} for name in list(FLOAT_FORMATS) + list(QUANT_FORMATS): try: value, runs, seconds = matmul(name, args.matmul_size, args.matmul_iters, args.max_seconds) except (RuntimeError, ValueError, TypeError) as exc: reason = str(exc).splitlines()[0][:100] print(f" {name:<9} unsupported: {reason}") tflops[name] = f"unsupported: {reason}" continue print(f" {name:<9} {value:8.3f} TFLOPS ({runs} in {seconds:.2f} s)") tflops[name] = round(value, 3) if args.labbook: record = { "lab": "part-05/mlx-bandwidth-test", "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), "tool": "mlx", "mlx": mx.__version__, "device": str(mx.default_device()), "method": "pass: bytes read + written per second; read: W @ x, bytes of W per " "second; matmul: 2 x N^3 flops, quantised formats via quantized_matmul", "size_mb": args.size_mb, "iters": args.iters, "repeats": args.repeats, "gbps": {"pass": round(statistics.median(passes), 1), "read": round(statistics.median(reads), 1), "pass_runs": [round(r, 1) for r in passes], "read_runs": [round(r, 1) for r in reads]}, "matmul_size": args.matmul_size, "tflops": tflops, "device_info": shown, } with Path(args.labbook).open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True, default=str) + "\n") print(f"\nrecorded in {args.labbook}") if __name__ == "__main__": main()