File size: 3,350 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 | """Screen a 512 MiB free-allocator cache on the actual fully tuned first block."""
import argparse
import json
import statistics
import time
from pathlib import Path
import mlx.core as mx
from h3_mlx.benchmark import error
from h3_mlx.config import SIGMAS
from h3_mlx.dit import DiT
from h3_mlx.layout import image_positions, rope_cache
from h3_mlx.memory import FootprintMonitor
from h3_mlx.runtime import configure_memory, environment
def main(a):
configure_memory(wired_gb=14)
data = mx.load("outputs/suite/product/native-3.trace.safetensors")
model = DiT(a.model, kernel="metal-exact", tune_matmul=True)
model.prepare_modulation(SIGMAS)
text = model.refine(data["conditioning"])
x = mx.concatenate(
(
text,
model._linear("video_patch_proj")(data["initial_rows"]).astype(mx.bfloat16),
),
axis=0,
)
rope = rope_cache(
image_positions(len(text), 128, 128), model._constant("rope.inv_freq")
)
ids = mx.concatenate((data["text_tags"], mx.zeros(4096, mx.int32)))
params = tuple(model.modulation["block.0"][0, :, j] for j in range(6))
block = model._block("blocks.0")
fn = mx.compile(
lambda x, p, i, r: block(
x, params=p, indices=i, rope=r, kernel="metal-exact", mlp_chunk=0
)
)
ref = fn(x, params, ids, rope)
mx.eval(ref)
report = {
"environment": environment(),
"scope": "Real complete first block, not full-image performance. Cache is free allocations only; same compiled graph and original weights.",
"runs": [],
}
for repeat in range(6):
for cap in (1.0, 0.5) if repeat % 2 == 0 else (0.5, 1.0):
mx.set_cache_limit(int(cap * 1024**3))
for _ in range(2):
mx.eval(fn(x, params, ids, rope))
mx.reset_peak_memory()
with FootprintMonitor(interval=0.001) as memory:
mx.synchronize()
tick = time.perf_counter()
result = fn(x, params, ids, rope)
mx.eval(result)
mx.synchronize()
elapsed = time.perf_counter() - tick
row = {
"cache_gb": cap,
"repeat": repeat,
"milliseconds": elapsed * 1000,
"physical": memory.result(),
"mlx_peak_bytes": mx.get_peak_memory(),
"mlx_cache_bytes": mx.get_cache_memory(),
**error(ref, result),
}
report["runs"].append(row)
print(json.dumps(row), flush=True)
report["summary"] = {
str(cap): {
"median_ms": statistics.median(
r["milliseconds"] for r in report["runs"] if r["cache_gb"] == cap
),
"max_physical_bytes": max(
r["physical"]["sampled_peak_physical_bytes"]
for r in report["runs"]
if r["cache_gb"] == cap
),
}
for cap in (1.0, 0.5)
}
Path(a.output).write_text(json.dumps(report, indent=2) + "\n")
model.clear()
if __name__ == "__main__":
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--model", required=True)
p.add_argument("--output", default="reports/optimization-r3/cache-screen.json")
main(p.parse_args())
|