| """Matched compiled-block validation of the operator finalists.""" |
|
|
| import argparse |
| import copy |
| import json |
| import types |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from experiments.r3.common import paired |
| from experiments.r3.fusions import TunedFC2, qk_native_reduce |
| 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.ops import attention, modulate, rms_norm, rotate |
| from h3_mlx.runtime import configure_memory, environment |
|
|
|
|
| def main(a): |
| configure_memory(wired_gb=14) |
| report = {"environment": environment(), "trace": a.trace, "trials": []} |
| with FootprintMonitor() as 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") |
| native_fc2 = block.fc2 |
| control = mx.compile( |
| lambda x, params, indices, rope: block( |
| x, params=params, indices=indices, rope=rope, mlp_chunk=0 |
| ) |
| ) |
| ref = control(x, params, indices, rope) |
| mx.eval(ref) |
| xx = modulate(x, block.n1, params[0], params[1], indices, block.config.norm_eps) |
| qkv = block.qkv(xx).reshape(-1, 3, 56, 128) |
| qref = tuple( |
| rotate(rms_norm(qkv[None, :, j], w, block.config.qk_norm_eps), *rope) |
| for j, w in ((0, block.qn), (1, block.kn)) |
| ) |
| mx.eval(qkv, qref) |
| for head_major in (False, True): |
|
|
| def fn(head_major=head_major): |
| return qk_native_reduce( |
| qkv, |
| block.qn, |
| block.kn, |
| *rope, |
| block.config.qk_norm_eps, |
| head_major=head_major, |
| ) |
|
|
| result = fn() |
| if head_major: |
| result = tuple(t.transpose(0, 2, 1, 3) for t in result) |
| row = { |
| "group": "qk_native_reduce", |
| "head_major": head_major, |
| **error(mx.stack(qref), mx.stack(result)), |
| } |
|
|
| def native(): |
| return tuple( |
| rotate( |
| rms_norm(qkv[None, :, j], w, block.config.qk_norm_eps), *rope |
| ) |
| for j, w in ((0, block.qn), (1, block.kn)) |
| ) |
|
|
| row.update(paired(native, fn, a.repeats)) |
| report["trials"].append(row) |
| print(json.dumps(row), flush=True) |
| block.fc2 = TunedFC2(native_fc2) |
| candidate = mx.compile( |
| lambda x, params, indices, rope: block( |
| x, params=params, indices=indices, rope=rope, mlp_chunk=0 |
| ) |
| ) |
| row = { |
| "group": "fc2_compiled_block", |
| **error(ref, candidate(x, params, indices, rope)), |
| } |
| row.update( |
| paired( |
| lambda: control(x, params, indices, rope), |
| lambda: candidate(x, params, indices, rope), |
| a.repeats, |
| ) |
| ) |
| report["trials"].append(row) |
| print(json.dumps(row), flush=True) |
| |
| for tuned_fc2 in (False, True): |
| fused = copy.copy(block) |
| fused.fc2 = TunedFC2(native_fc2) if tuned_fc2 else native_fc2 |
|
|
| def fused_attn( |
| self, xx, rope, query_chunk, kernel="native", projection_weights=None |
| ): |
| qq = self.qkv(xx).reshape(-1, 3, 56, 128) |
| q, k = qk_native_reduce( |
| qq, self.qn, self.kn, *rope, self.config.qk_norm_eps |
| ) |
| v = qq[None, :, 2] |
| return self.out( |
| attention(q, k, v, query_chunk=query_chunk).reshape(xx.shape[0], -1) |
| ) |
|
|
| fused.attn = types.MethodType(fused_attn, fused) |
| runner = mx.compile( |
| lambda x, params, indices, rope, fused=fused: fused( |
| x, params=params, indices=indices, rope=rope, mlp_chunk=0 |
| ) |
| ) |
| row = { |
| "group": "qk_and_fc2_compiled_block" |
| if tuned_fc2 |
| else "qk_compiled_block", |
| **error(ref, runner(x, params, indices, rope)), |
| } |
| row.update( |
| paired( |
| lambda: control(x, params, indices, rope), |
| lambda runner=runner: runner(x, params, indices, rope), |
| a.repeats, |
| ) |
| ) |
| report["trials"].append(row) |
| print(json.dumps(row), flush=True) |
| model.clear() |
| report["memory"] = memory.result() |
| Path(a.output).write_text(json.dumps(report, indent=2) + "\n") |
|
|
|
|
| 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=10) |
| p.add_argument("--output", default="reports/optimization-r3/blocks.json") |
| main(p.parse_args()) |
|
|
