File size: 4,656 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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,
},
),
]
# Explicit state inputs prevent compile from retaining a dense constant.
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())
|