| """Measure exact palette12 packing versus quantized and decoded-weight GEMM.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import error, measure |
| from h3_mlx.packed import PackedLinear |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.weights import Linear, WeightStore |
|
|
|
|
| def main(a): |
| configure_memory(wired_gb=8) |
| dest = Path(a.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "warmup": 2, |
| "repeats": a.repeats, |
| "scope": "Original block-0 weights; synthetic BF16 activations. Includes per-call GPU weight reconstruction, excludes initial packing.", |
| "trials": [], |
| } |
| store = WeightStore(Path(a.model) / "transformer") |
| mx.random.seed(731) |
| 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 = layer.weight |
| x = mx.random.normal((a.tokens, w.shape[1])).astype(mx.bfloat16) |
| ref = layer(x) |
| mx.eval(ref, x) |
| tick = time.perf_counter() |
| packed = PackedLinear.from_linear(layer) |
| pack_s = time.perf_counter() - tick |
| exact = mx.all( |
| packed.weight.unpack().view(mx.uint16) == w.view(mx.uint16) |
| ).item() |
| if not exact: |
| raise AssertionError("Lossless weight reconstruction failed") |
| functions = [ |
| ("bf16", lambda layer=layer, x=x: layer(x), {}), |
| ( |
| "packed12", |
| lambda packed=packed, x=x: packed(x), |
| { |
| "pack_s": pack_s, |
| "weight_bytes": packed.nbytes, |
| "ratio": packed.nbytes / layer.nbytes, |
| "weight_bits_identical": exact, |
| "escape_fraction": packed.weight.escape_indices.size / w.size, |
| }, |
| ), |
| ] |
| |
| def packed_forward(x, state, packed=packed): |
| return packed(x, state=state) |
| fn = mx.compile(packed_forward) |
| state = packed.weight.state() |
| functions.append(("packed12_compiled", lambda fn=fn, x=x, state=state: fn(x, state), {})) |
| for variant, call, extra in functions: |
| row = { |
| "operator": name, |
| "variant": variant, |
| "shape": [a.tokens, w.shape[1], w.shape[0]], |
| **extra, |
| **measure(call, repeats=a.repeats), |
| **error(ref, call()), |
| } |
| report["trials"].append(row) |
| dest.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
| del functions, fn, packed |
| for bits in (8, 6, 4): |
| q, s, z = mx.quantize(w, bits=bits, group_size=32) |
| mx.eval(q, s, z) |
|
|
| def decoded(q=q, s=s, z=z, bits=bits, x=x, layer=layer): |
| d = mx.dequantize(q, s, z, bits=bits, group_size=32) |
| return x @ d.T if layer.bias is None else mx.addmm(layer.bias, x, d.T) |
|
|
| def fused(q=q, s=s, z=z, bits=bits, x=x, layer=layer): |
| y = mx.quantized_matmul(x, q, s, z, bits=bits, group_size=32) |
| return y if layer.bias is None else y + layer.bias |
|
|
| for variant, call in ( |
| ( |
| f"q{bits}_fused", |
| fused, |
| ), |
| (f"q{bits}_dequant_gemm", decoded), |
| ): |
| row = { |
| "operator": name, |
| "variant": variant, |
| "shape": [a.tokens, w.shape[1], w.shape[0]], |
| **measure(call, repeats=a.repeats), |
| **error(ref, call()), |
| } |
| report["trials"].append(row) |
| dest.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
| del q, s, z, decoded |
| del layer, w, x, ref |
| gc.collect() |
| mx.clear_cache() |
|
|
|
|
| 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=7) |
| p.add_argument("--output", default="reports/optimization-v2/storage.json") |
| main(p.parse_args()) |
|
|
