| """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()) |
|
|
