| """Synchronized operator experiments on real H3 dimensions, with numerical error.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import statistics |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from .kernels import norm_modulate, qk_norm_rope |
| from .ops import attention, modulate, rms_norm, rotate |
| from .runtime import configure_memory, environment |
|
|
|
|
| def measure(fn, *, warmup=2, repeats=7): |
| for _ in range(warmup): |
| mx.eval(fn()) |
| samples = [] |
| for _ in range(repeats): |
| mx.synchronize() |
| tick = time.perf_counter() |
| mx.eval(fn()) |
| mx.synchronize() |
| samples.append((time.perf_counter() - tick) * 1000) |
| return { |
| "median_ms": statistics.median(samples), |
| "min_ms": min(samples), |
| "max_ms": max(samples), |
| "samples_ms": samples, |
| } |
|
|
|
|
| def error(reference, candidate): |
| a, b = reference.astype(mx.float32), candidate.astype(mx.float32) |
| diff = a - b |
| return { |
| "max_abs": mx.max(mx.abs(diff)).item(), |
| "relative_rmse": mx.sqrt( |
| mx.mean(diff * diff) / (mx.mean(a * a) + 1e-30) |
| ).item(), |
| "equal_fraction": mx.mean((a == b).astype(mx.float32)).item(), |
| } |
|
|
|
|
| def run(output, *, tokens=(256, 4096), repeats=7): |
| configure_memory() |
| mx.random.seed(731) |
| result = { |
| "environment": environment(), |
| "warmup": 2, |
| "repeats": repeats, |
| "scope": "Synthetic operator microbenchmarks; excludes loading/quantization. Not image latency.", |
| "trials": [], |
| } |
|
|
| def record(group, variant, shape, fn, reference): |
| data = { |
| "operator": group, |
| "variant": variant, |
| "shape": shape, |
| **measure(fn, repeats=repeats), |
| **error(reference, fn()), |
| } |
| result["trials"].append(data) |
| print(json.dumps(data), flush=True) |
| |
| Path(output).parent.mkdir(parents=True, exist_ok=True) |
| Path(output).write_text(json.dumps(result, indent=2) + "\n") |
|
|
| for n in tokens: |
| x = mx.random.normal((n, 5376)).astype(mx.bfloat16) |
| w = mx.random.normal((5376,)).astype(mx.bfloat16) |
| shift = (mx.random.normal((3, 5376)) * 0.1).astype(mx.bfloat16) |
| scale = (mx.random.normal((3, 5376)) * 0.1).astype(mx.bfloat16) |
| ids = mx.arange(n, dtype=mx.int32) % 3 |
| mx.eval(x, w, shift, scale, ids) |
| ref = modulate(x, w, shift, scale, ids) |
| mx.eval(ref) |
| record( |
| "rms_adaln", |
| "native_rounding", |
| [n, 5376], |
| lambda x=x, w=w, shift=shift, scale=scale, ids=ids: modulate( |
| x, w, shift, scale, ids |
| ), |
| ref, |
| ) |
| for threads in (128, 256, 512): |
| record( |
| "rms_adaln", |
| f"metal_tg{threads}", |
| [n, 5376], |
| lambda x=x, w=w, shift=shift, scale=scale, ids=ids, threads=threads: ( |
| norm_modulate(x, w, shift, scale, ids, threads=threads) |
| ), |
| ref, |
| ) |
| qkv = mx.random.normal((n, 3, 56, 128)).astype(mx.bfloat16) |
| qw = mx.random.normal((128,)).astype(mx.bfloat16) |
| angles = mx.random.normal((1, n, 1, 48)) |
| angles = mx.concatenate((angles, angles), axis=-1) |
| cos, sin = ( |
| mx.cos(angles).astype(mx.bfloat16), |
| mx.sin(angles).astype(mx.bfloat16), |
| ) |
|
|
| def native_qk(qkv=qkv, qw=qw, cos=cos, sin=sin): |
| return tuple( |
| rotate(rms_norm(qkv[None, :, i], qw), cos, sin) for i in (0, 1) |
| ) |
|
|
| def fused_qk(qkv=qkv, qw=qw, cos=cos, sin=sin): |
| return qk_norm_rope(qkv, qw, qw, cos, sin) |
|
|
| qr, kr = native_qk() |
| both = mx.concatenate((qr, kr), axis=0) |
| mx.eval(both) |
| record( |
| "qk_norm_rope", |
| "native", |
| [n, 56, 128], |
| lambda fn=native_qk: mx.concatenate(fn(), axis=0), |
| both, |
| ) |
| record( |
| "qk_norm_rope", |
| "metal", |
| [n, 56, 128], |
| lambda fn=fused_qk: mx.concatenate(fn(), axis=0), |
| both, |
| ) |
| del qkv, qw, angles, cos, sin, qr, kr, both, native_qk, fused_qk |
| mx.clear_cache() |
| for k, out in ((5376, 21504), (5376, 28672), (14336, 5376)): |
| a = mx.random.normal((n, k)).astype(mx.bfloat16) |
| weight = (mx.random.normal((out, k)) * 0.02).astype(mx.bfloat16) |
| mx.eval(a, weight) |
| baseline = a @ weight.T |
| mx.eval(baseline) |
| record( |
| "gemm", |
| "bf16", |
| [n, k, out], |
| lambda a=a, weight=weight: a @ weight.T, |
| baseline, |
| ) |
| for bits in (8, 4): |
| for size in (32, 64, 128): |
| q, s, z = mx.quantize(weight, bits=bits, group_size=size) |
| mx.eval(q, s, z) |
| record( |
| "gemm", |
| f"q{bits}_g{size}", |
| [n, k, out], |
| lambda a=a, q=q, s=s, z=z, size=size, bits=bits: ( |
| mx.quantized_matmul( |
| a, q, s, z, transpose=True, group_size=size, bits=bits |
| ) |
| ), |
| baseline, |
| ) |
| del a, weight, baseline, q, s, z |
| gc.collect() |
| mx.clear_cache() |
| q = mx.random.normal((1, n, 56, 128)).astype(mx.bfloat16) |
| k = mx.random.normal(q.shape).astype(mx.bfloat16) |
| v = mx.random.normal(q.shape).astype(mx.bfloat16) |
| mx.eval(q, k, v) |
| baseline = attention(q, k, v) |
| mx.eval(baseline) |
| for chunk in (0, 128, 512, 1024): |
| record( |
| "sdpa", |
| f"query_chunk_{chunk}", |
| list(q.shape), |
| lambda q=q, k=k, v=v, chunk=chunk: attention( |
| q, k, v, query_chunk=chunk |
| ), |
| baseline, |
| ) |
| del q, k, v, baseline, x, w, shift, scale, ids, ref |
| gc.collect() |
| mx.clear_cache() |
| return result |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--output", default="reports/operators.json") |
| p.add_argument("--tokens", type=int, nargs="+", default=[256, 4096]) |
| p.add_argument("--repeats", type=int, default=7) |
| a = p.parse_args() |
| if a.repeats < 1 or any(n < 1 for n in a.tokens): |
| p.error("tokens and repeats must be positive") |
| run(a.output, tokens=a.tokens, repeats=a.repeats) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
