| """Actual H3 projection dimensions, M5 layouts, and new quantization formats. |
| |
| Run as a module from the repository root. Each result is persisted immediately. |
| The inputs are synthetic; weights come from the original checkpoint. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import error, measure |
| from h3_mlx.ops import swiglu |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.weights import Linear, WeightStore |
|
|
|
|
| def run(args): |
| configure_memory() |
| mx.random.seed(731) |
| dest = Path(args.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "warmup": 2, |
| "repeats": args.repeats, |
| "scope": "Original block-0 weights; synthetic BF16 inputs. Operator latency only, excludes conversion and loading.", |
| "trials": [], |
| } |
|
|
| def record(op, variant, fn, ref, **extra): |
| try: |
| timing = measure(fn, repeats=args.repeats) |
| row = { |
| "operator": op, |
| "variant": variant, |
| **extra, |
| **timing, |
| **error(ref, fn()), |
| } |
| except Exception as exc: |
| row = { |
| "operator": op, |
| "variant": variant, |
| **extra, |
| "error": f"{type(exc).__name__}: {exc}", |
| } |
| report["trials"].append(row) |
| dest.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
|
|
| store = WeightStore(Path(args.model) / "transformer") |
| for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"): |
| layer = Linear.load( |
| store, |
| "blocks.0." + name, |
| grouped_heads=(56, 128) if name == "attn.qkv_proj" else None, |
| ) |
| w, bias = layer.weight, layer.bias |
| x = mx.random.normal((args.tokens, w.shape[1])).astype(mx.bfloat16) |
| mx.eval(x, w) |
| ref = layer(x) |
| mx.eval(ref) |
| shape = [args.tokens, w.shape[1], w.shape[0]] |
| record(name, "bf16_nt", lambda layer=layer, x=x: layer(x), ref, shape=shape) |
| wt = mx.contiguous(w.T) |
| mx.eval(wt) |
| fn = ( |
| (lambda x=x, wt=wt: x @ wt) |
| if bias is None |
| else (lambda bias=bias, x=x, wt=wt: mx.addmm(bias, x, wt)) |
| ) |
| record(name, "bf16_nn", fn, ref, shape=shape) |
| del wt, fn |
| for pad in (64, 128, 256): |
| n = ((args.tokens + pad - 1) // pad) * pad |
| padded = mx.pad(x, ((0, n - args.tokens), (0, 0))) |
| mx.eval(padded) |
| record( |
| name, |
| f"bf16_pad{pad}", |
| lambda layer=layer, padded=padded: layer(padded)[: args.tokens], |
| ref, |
| shape=shape, |
| padded_tokens=n, |
| ) |
| del padded |
| half, halfw = x.astype(mx.float16), w.astype(mx.float16) |
| halfb = None if bias is None else bias.astype(mx.float16) |
| mx.eval(half, halfw, halfb) |
| fn = ( |
| (lambda half=half, halfw=halfw: (half @ halfw.T).astype(mx.bfloat16)) |
| if bias is None |
| else ( |
| lambda halfb=halfb, half=half, halfw=halfw: mx.addmm( |
| halfb, half, halfw.T |
| ).astype(mx.bfloat16) |
| ) |
| ) |
| record(name, "fp16_round_bf16", fn, ref, shape=shape) |
| del half, halfw, halfb, fn |
| for mode, bits, group in ( |
| ("mxfp8", 8, 32), |
| ("mxfp4", 4, 32), |
| ("nvfp4", 4, 16), |
| ("affine", 8, 32), |
| ("affine", 6, 32), |
| ("affine", 4, 32), |
| ): |
| try: |
| packed = mx.quantize(w, mode=mode, bits=bits, group_size=group) |
| mx.eval(packed) |
|
|
| def quant( |
| x=x, packed=packed, mode=mode, bits=bits, group=group, bias=bias |
| ): |
| y = mx.quantized_matmul( |
| x, *packed, mode=mode, bits=bits, group_size=group |
| ) |
| return y if bias is None else y + bias |
|
|
| record( |
| name, |
| f"{mode}{bits}_g{group}", |
| quant, |
| ref, |
| shape=shape, |
| weight_gib=sum(v.nbytes for v in packed) / 1024**3, |
| ) |
| del packed, quant |
| except Exception as exc: |
| record( |
| name, |
| f"{mode}{bits}_g{group}", |
| lambda exc=exc: (_ for _ in ()).throw(exc), |
| ref, |
| ) |
| del x, w, bias, layer, ref |
| gc.collect() |
| mx.clear_cache() |
| a = Linear.load(store, "blocks.0.mlp.fc1") |
| b = Linear.load(store, "blocks.0.mlp.fc2") |
| x = mx.random.normal((args.tokens, 5376)).astype(mx.bfloat16) |
|
|
| def mlp(x): |
| return b(swiglu(a(x))) |
|
|
| ref = mlp(x) |
| mx.eval(ref) |
| record("mlp", "eager", lambda: mlp(x), ref) |
| compiled = mx.compile(mlp) |
| record("mlp", "compiled", lambda: compiled(x), ref) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--tokens", type=int, default=4181) |
| p.add_argument("--repeats", type=int, default=5) |
| p.add_argument("--output", default="reports/optimization-v2/operators.json") |
| run(p.parse_args()) |
|
|
