| """Real first-block inputs: attention layout, Metal 4 FC2 tiling, lossless decode.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import itertools |
| import json |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from experiments.r3.common import paired |
| from experiments.v2_mpp import gemm |
| 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 rms_norm, rotate |
| from h3_mlx.packed import PackedBF16 |
| from h3_mlx.runtime import configure_memory, environment |
|
|
|
|
| @lru_cache(None) |
| def decoder(lanes): |
| return mx.fast.metal_kernel( |
| name=f"h3_unpack_bf16_palette12_vec{lanes}", |
| input_names=["low", "nibbles", "escape_indices", "escape_values", "base"], |
| output_names=["out"], |
| source=r""" |
| uint first = thread_position_in_grid.x * LANES; |
| #pragma unroll |
| for (uint lane = 0; lane < LANES; ++lane) { |
| uint i = first + lane; |
| if (i >= low_shape[0]) break; |
| uint code = (nibbles[i / 2] >> ((i % 2) * 4)) & 15; |
| uint count = escape_indices_shape[0]; |
| if (code == 0 && count != 0) { |
| uint left = 0, right = count; |
| while (left < right) { |
| uint mid = (left + right) / 2; |
| if (escape_indices[mid] < i) left = mid + 1; |
| else right = mid; |
| } |
| out[i] = escape_values[left]; |
| } else { |
| out[i] = ushort(low[i]) | ushort(((code & 7) + base[0]) << 8) |
| | ushort((code & 8) << 12); |
| } |
| } |
| """, |
| ) |
|
|
|
|
| def unpack(packed, lanes, threads): |
| raw = decoder(lanes)( |
| inputs=packed.state(), |
| template=[("LANES", lanes)], |
| grid=((packed.low.size + lanes - 1) // lanes, 1, 1), |
| threadgroup=(threads, 1, 1), |
| output_shapes=[(packed.low.size,)], |
| output_dtypes=[mx.uint16], |
| )[0] |
| return raw.view(mx.bfloat16).reshape(packed.shape) |
|
|
|
|
| class Capture: |
| def __init__(self, linear): |
| self.linear = linear |
|
|
| def __call__(self, x): |
| self.input = x |
| mx.eval(x) |
| return self.linear(x) |
|
|
|
|
| def main(a): |
| configure_memory(wired_gb=14) |
| dest = Path(a.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "scope": "Actual saved first-block activations, one GPU job, counterbalanced warmed calls including layout/decode costs. Operator evidence only.", |
| "trials": [], |
| } |
|
|
| def record(row): |
| report["trials"].append(row) |
| dest.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
|
|
| 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") |
| captures = {} |
| for name in ("qkv", "out", "fc1", "fc2"): |
| capture = Capture(getattr(block, name)) |
| captures[name] = capture |
| setattr(block, name, capture) |
| y = block(x, params=params, indices=indices, rope=rope, mlp_chunk=0) |
| mx.eval(y) |
| report["activation_ranges"] = { |
| name: { |
| "shape": list(v.input.shape), |
| "max_abs": mx.max(mx.abs(v.input.astype(mx.float32))).item(), |
| } |
| for name, v in captures.items() |
| } |
| if "attention" in a.groups: |
| qkv = ( |
| captures["qkv"] |
| .linear(captures["qkv"].input) |
| .reshape(1, x.shape[0], 3, 56, 128) |
| ) |
| q, k = ( |
| rotate( |
| rms_norm(qkv[:, :, j], weight, block.config.qk_norm_eps), *rope |
| ).transpose(0, 2, 1, 3) |
| for j, weight in ((0, block.qn), (1, block.kn)) |
| ) |
| v = qkv[:, :, 2].transpose(0, 2, 1, 3) |
| mx.eval(q, k, v) |
|
|
| def attention(layout="strided"): |
| args = [ |
| mx.contiguous(t) if axis in layout else t |
| for axis, t in zip("qkv", (q, k, v), strict=True) |
| ] |
| return mx.fast.scaled_dot_product_attention(*args, scale=128**-0.5) |
|
|
| ref = attention() |
| mx.eval(ref) |
| for layout in ("q", "k", "v", "kv", "qk", "qkv"): |
|
|
| def fn(layout=layout): |
| return attention(layout) |
|
|
| record( |
| { |
| "group": "attention", |
| "variant": layout, |
| **paired(attention, fn, a.repeats), |
| **error(ref, fn()), |
| } |
| ) |
| del qkv, q, k, v, ref |
| if "mpp" in a.groups: |
| count = 0 |
| for name in ("fc2", "fc1"): |
| cap = captures[name] |
| layer, xx = cap.linear, cap.input |
| w = layer.weight |
| b = ( |
| layer.bias |
| if layer.bias is not None |
| else mx.zeros(w.shape[0], w.dtype) |
| ) |
| ref = layer(xx) |
| mx.eval(ref) |
| tiles = itertools.product( |
| ( |
| (64, 128, 4), |
| (64, 128, 8), |
| (128, 128, 4), |
| (128, 128, 8), |
| (128, 256, 8), |
| (64, 256, 8), |
| ), |
| (128, 256, 512), |
| (0, 8), |
| ) |
| for (bm, bn, sg), bk, walk in tiles: |
| if w.shape[1] % bk or w.shape[0] % bn: |
| continue |
| count += 1 |
| row = { |
| "group": "mpp", |
| "operator": name, |
| "tile": [bm, bn, bk], |
| "simdgroups": sg, |
| "walk": walk, |
| } |
| try: |
|
|
| def fn(bm=bm, bn=bn, bk=bk, sg=sg, walk=walk, xx=xx, w=w, b=b): |
| return gemm(xx, w, b, bm, bn, bk, sg, walk) |
|
|
| row.update( |
| paired(lambda layer=layer, xx=xx: layer(xx), fn, a.repeats), |
| **error(ref, fn()), |
| ) |
| except Exception as exc: |
| row["error"] = f"{type(exc).__name__}: {exc}" |
| record(row) |
| if count >= a.max_tiles: |
| break |
| if count >= a.max_tiles: |
| break |
| if "unpack" in a.groups: |
| for name in ("qkv", "fc2"): |
| cap = captures[name] |
| layer, xx = cap.linear, cap.input |
| packed = PackedBF16.pack(layer.weight) |
| ref = packed.unpack() |
| mx.eval(ref) |
| for lanes, threads in itertools.product((2, 4, 8), (128, 256)): |
|
|
| def fn(lanes=lanes, threads=threads, packed=packed): |
| return unpack(packed, lanes, threads) |
|
|
| candidate = fn() |
| mx.eval(candidate) |
| bit_equal = mx.array_equal( |
| ref.view(mx.uint16), candidate.view(mx.uint16) |
| ).item() |
| row = { |
| "group": "unpack", |
| "operator": name, |
| "lanes": lanes, |
| "threads": threads, |
| "bit_equal": bit_equal, |
| "decode": paired(packed.unpack, fn, a.repeats), |
| } |
|
|
| def project(w, xx=xx, layer=layer): |
| return ( |
| xx @ w.T |
| if layer.bias is None |
| else mx.addmm(layer.bias, xx, w.T) |
| ) |
|
|
| row["with_gemm"] = paired( |
| lambda project=project, packed=packed: project(packed.unpack()), |
| lambda project=project, fn=fn: project(fn()), |
| a.repeats, |
| ) |
| record(row) |
| del packed, ref |
| model.clear() |
| gc.collect() |
| report["memory"] = memory.result() |
| dest.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("--groups", nargs="+", default=["attention", "mpp", "unpack"]) |
| p.add_argument("--max-tiles", type=int, default=72) |
| p.add_argument("--repeats", type=int, default=6) |
| p.add_argument("--output", default="reports/optimization-r3/screen.json") |
| main(p.parse_args()) |
|
|
