File size: 3,155 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 | """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 # Chunking intentionally synchronizes intermediates.
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())
|