| """Benchmark full first DiT block with real saved conditioning/noise and weights.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import error, measure |
| from h3_mlx.config import SIGMAS |
| from h3_mlx.dit import DiT |
| from h3_mlx.layout import image_positions, rope_cache |
| from h3_mlx.runtime import configure_memory, environment |
|
|
|
|
| def main(a): |
| configure_memory() |
| data = mx.load(a.trace) |
| model = DiT(a.model) |
| 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") |
| ) |
| indices = 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") |
| dest = Path(a.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "trace": a.trace, |
| "warmup": 2, |
| "repeats": a.repeats, |
| "trials": [], |
| } |
| ref = block(x, params=params, indices=indices, rope=rope, mlp_chunk=0) |
| mx.eval(ref) |
| for kernel in ("native", "metal", "metal-qk"): |
| for chunk in (0, 512, 1024, 2048): |
| for compiled in (False, True): |
| if compiled and chunk: |
| continue |
|
|
| def forward(x, params, indices, rope, chunk=chunk, kernel=kernel): |
| return block( |
| x, |
| params=params, |
| indices=indices, |
| rope=rope, |
| mlp_chunk=chunk, |
| kernel=kernel, |
| ) |
|
|
| fn = mx.compile(forward) if compiled else forward |
| row = {"kernel": kernel, "mlp_chunk": chunk, "compiled": compiled} |
| try: |
| row.update( |
| measure( |
| lambda fn=fn: fn(x, params, indices, rope), |
| repeats=a.repeats, |
| ), |
| **error(ref, fn(x, params, indices, rope)), |
| ) |
| except Exception as exc: |
| row["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) |
| model.clear() |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument( |
| "--trace", default="outputs/suite/product/native-3.trace.safetensors" |
| ) |
| p.add_argument("--repeats", type=int, default=5) |
| p.add_argument("--output", default="reports/optimization-v2/blocks.json") |
| main(p.parse_args()) |
|
|
