"""Measure matrix-multiply throughput in TFLOPS at FP32, BF16, FP16 and two lower precisions. Purpose: measure the compute half of the pair that predicts everything. Bandwidth sets the decode ceiling; matrix-multiply throughput sets what prefill and training can do. This is the compute-bound counterpart to bandwidth-test.py. Platform: all (cuda on Tracks S and N and on Track X with the ROCm build, mps on Track M, cpu everywhere as a fallback) Minimum memory: 8 GB Assumes: torch is installed in the active environment; three matrices of --size fit on the device (at the default 4096, 64 MiB each in FP32, 256 MiB each at 8192) Usage: python3 matmul-test.py [--size 4096] [--iters 30] [--device auto] [--dtypes fp32,bf16,fp16,fp8,int8w] [--max-seconds 30] [--bandwidth-gbps 250] [--labbook labbook.md] Method. C = A @ B for two N x N matrices performs N^3 multiplies and about N^3 adds, so each multiplication counts as 2 x N^3 floating-point operations, and the rate is 2 x N^3 x multiplications / seconds / 1e12 TFLOPS. The formats: fp32, bf16, fp16 A and B both in that dtype, one torch.matmul fp8 A and B as torch.float8_e4m3fn through the same plain torch.matmul; a build without that kernel raises, recorded as unsupported int8w weight-only INT8: B is stored as int8 with one bf16 scale per row and converted back to bf16 before every multiplication, the way a weight-only format without a native kernel runs; same operation count Before each format a 512 x 512 probe estimates one full-size multiplication. A format whose estimate exceeds --max-seconds is recorded as "slow path" with the probe's rate rather than left running for many minutes, which is what a CPU without a fast kernel for that dtype would otherwise do; otherwise the loop is shortened to fit --max-seconds, but never below three multiplications. With --bandwidth-gbps, the ridge point (TFLOPS x 1e12) / (GB/s x 1e9) is printed in floating-point operations per byte. """ import argparse import json import sys import time from datetime import datetime, timezone from pathlib import Path import torch FORMATS = ("fp32", "bf16", "fp16", "fp8", "int8w") DTYPES = {"fp32": "float32", "bf16": "bfloat16", "fp16": "float16", "fp8": "float8_e4m3fn"} PROBE = 512 FAILURES = (RuntimeError, TypeError, NotImplementedError) def pick_device(requested): if requested != "auto": return torch.device(requested) if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def synchronize(device): if device.type == "cuda": torch.cuda.synchronize() elif device.type == "mps": torch.mps.synchronize() def release(device): if device.type == "cuda": torch.cuda.empty_cache() elif device.type == "mps": torch.mps.empty_cache() def describe(device): if device.type == "cuda": if torch.version.hip: runtime = f"ROCm (HIP {torch.version.hip})" else: runtime = f"CUDA {torch.version.cuda}" return f"{torch.cuda.get_device_name(0)}, {runtime}" if device.type == "mps": return "Apple silicon GPU through MPS" return f"CPU, {torch.get_num_threads()} threads" def build(fmt, side, device): """Return a zero-argument function that performs one multiplication in this format.""" a32 = torch.randn(side, side, device=device, dtype=torch.float32) b32 = torch.randn(side, side, device=device, dtype=torch.float32) if fmt == "int8w": a = a32.to(torch.bfloat16) scale = b32.abs().amax(dim=1, keepdim=True).clamp(min=1e-8) / 127.0 q = torch.round(b32 / scale).clamp(-127, 127).to(torch.int8) s = scale.to(torch.bfloat16) del a32, b32 return lambda: torch.matmul(a, q.to(torch.bfloat16) * s) dtype = getattr(torch, DTYPES[fmt]) a, b = a32.to(dtype), b32.to(dtype) del a32, b32 return lambda: torch.matmul(a, b) def rate(device, fn, side, count=None, min_seconds=None): """Time count multiplications, or as many as fill min_seconds; return (flops/s, seconds).""" synchronize(device) start = time.perf_counter() done = 0 while True: fn() done += 1 if count is not None and done >= count: break if min_seconds is not None: synchronize(device) if time.perf_counter() - start >= min_seconds: break synchronize(device) seconds = time.perf_counter() - start return 2.0 * side ** 3 * done / seconds, seconds def measure(device, fmt, size, iters, max_seconds): """Return (result, runs, seconds, probe GFLOP/s); result is TFLOPS or a string.""" probe_side = min(PROBE, size) fn = build(fmt, probe_side, device) fn() # warm-up: kernel selection, allocation probe_flops, _ = rate(device, fn, probe_side, min_seconds=0.25) del fn one = 2.0 * size ** 3 / probe_flops # estimated seconds for one full-size multiply if one > max_seconds: release(device) note = (f"slow path: {probe_side} probe {probe_flops / 1e9:.2f} GFLOP/s, " f"one {size} multiply ~{one:.0f} s") return note, 0, 0.0, probe_flops / 1e9 fn = build(fmt, size, device) fn() # warm-up at full size _, one = rate(device, fn, size, count=1) # the loop is sized from a real multiplication runs = max(min(iters, 3), min(iters, int(max_seconds / max(one, 1e-9)))) flops, seconds = rate(device, fn, size, count=runs) del fn release(device) return flops / 1e12, runs, seconds, probe_flops / 1e9 def main(): parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--size", type=int, default=4096, help="square matrix side") parser.add_argument("--iters", type=int, default=30, help="multiplications per format, fewer if --max-seconds binds") parser.add_argument("--device", default="auto", help="auto, cuda, mps or cpu") parser.add_argument("--dtypes", default=",".join(FORMATS), help=f"comma-separated, from {','.join(FORMATS)}") parser.add_argument("--max-seconds", type=float, default=30.0, help="time budget per format") parser.add_argument("--bandwidth-gbps", type=float, default=None, help="read figure from bandwidth-test.py, to print the ridge point") parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file") args = parser.parse_args() wanted = [d.strip() for d in args.dtypes.split(",") if d.strip()] unknown = [d for d in wanted if d not in FORMATS] if unknown or args.size < 2 or args.iters < 1 or args.max_seconds <= 0: sys.exit(f"matmul-test: formats must come from {','.join(FORMATS)} " f"(got {','.join(unknown) or 'none unknown'}); --size, --iters and " f"--max-seconds must be positive") device = pick_device(args.device) print(f"torch {torch.__version__}") print(f"device: {device.type} ({describe(device)})") print(f"matrices: {args.size} x {args.size}; up to {args.iters} multiplications or " f"{args.max_seconds:.0f} s per format; 2 x N^3 = {2 * args.size ** 3 / 1e9:,.1f} " f"GFLOP per multiplication\n") ridge = args.bandwidth_gbps is not None header = f"{'format':<7}{'TFLOPS':>9}{'runs':>6}{'seconds':>9}" header += f"{'FLOP/byte':>11}" if ridge else "" print(header + " note") print("-" * (len(header) + 6)) tflops, runs_by_format, probes = {}, {}, {} for fmt in wanted: try: result, runs, seconds, probe = measure(device, fmt, args.size, args.iters, args.max_seconds) except FAILURES as exc: release(device) reason = str(exc).splitlines()[0][:100] print(f"{fmt:<7}{'-':>9}{'-':>6}{'-':>9}" + (f"{'-':>11}" if ridge else "") + f" unsupported: {reason}") tflops[fmt] = f"unsupported: {reason}" continue probes[fmt] = round(probe, 2) if isinstance(result, str): print(f"{fmt:<7}{'-':>9}{'-':>6}{'-':>9}" + (f"{'-':>11}" if ridge else "") + f" {result}") tflops[fmt] = result continue line = f"{fmt:<7}{result:>9.3f}{runs:>6}{seconds:>9.2f}" if ridge: line += f"{result * 1e12 / (args.bandwidth_gbps * 1e9):>11.0f}" print(line) tflops[fmt] = round(result, 3) runs_by_format[fmt] = runs if args.labbook: record = { "lab": "part-05/matmul-test", "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), "tool": "torch", "torch": torch.__version__, "device": device.type, "device_name": describe(device), "method": "torch.matmul, 2 x N^3 flops per multiplication; int8w dequantises B first", "size": args.size, "iters": args.iters, "max_seconds": args.max_seconds, "tflops": tflops, "runs": runs_by_format, "probe_gflops": probes, } if ridge: record["bandwidth_gbps"] = args.bandwidth_gbps with Path(args.labbook).open("a", encoding="utf-8") as handle: handle.write(json.dumps(record, sort_keys=True) + "\n") print(f"\nrecorded in {args.labbook}") if __name__ == "__main__": main()