| """Shape-specific Metal 4 TensorOps GEMM tile, K barrier and walk-order search.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import error, measure |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.weights import Linear, WeightStore |
|
|
|
|
| @lru_cache(maxsize=None) |
| def kernel(bm, bn, bk, groups, walk): |
| header = """ |
| #include <metal_tensor> |
| #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> |
| using namespace mpp::tensor_ops; |
| """ |
| source = r""" |
| const int M = x_shape[0], K = x_shape[1], N = w_shape[0]; |
| int tiles_m = (M + BM - 1) / BM, tiles_n = N / BN; |
| int tile = threadgroup_position_in_grid.x; |
| int ti, tj; |
| if (WALK == 0) { |
| ti = tile / tiles_n; tj = tile % tiles_n; |
| } else { |
| // Group adjacent M rows, interleaving N tiles for L2 reuse. |
| int group = tile / (WALK * tiles_n); |
| int first_m = group * WALK; |
| int size_m = min(tiles_m - first_m, WALK); |
| int local = tile - group * WALK * tiles_n; |
| ti = first_m + local % size_m; |
| tj = local / size_m; |
| } |
| auto A = tensor(const_cast<device T*>(x), dextents<int, 2>(K, M), array<int, 2>{1, K}); |
| auto B = tensor(const_cast<device T*>(w), dextents<int, 2>(K, N), array<int, 2>{1, K}); |
| auto D = tensor(out, dextents<int, 2>(N, M), array<int, 2>{1, N}); |
| constexpr auto mode = matmul2d_descriptor::mode::multiply_accumulate; |
| constexpr auto desc = matmul2d_descriptor(BM, BN, BK, false, true, false, mode); |
| matmul2d<desc, execution_simdgroups<SG>> op; |
| auto a0 = A.template slice<BK, dynamic_extent>(0, ti * BM); |
| auto b0 = B.template slice<BK, BN>(0, tj * BN); |
| auto accum = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), float>(); |
| for (int e = 0; e < accum.get_capacity(); ++e) accum[e] = 0.0f; |
| for (int k = 0; k < K; k += BK) { |
| threadgroup_barrier(mem_flags::mem_none); |
| auto a = A.template slice<BK, dynamic_extent>(k, ti * BM); |
| auto b = B.template slice<BK, BN>(k, tj * BN); |
| op.run(a, b, accum); |
| } |
| auto bias_tensor = tensor(const_cast<device T*>(bias), dextents<int, 2>(N, M), array<int, 2>{1, 0}); |
| auto bias_tile = bias_tensor.template slice<BN, dynamic_extent>(tj * BN, ti * BM); |
| auto bias_regs = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), T>(); |
| bias_regs.load(bias_tile); |
| auto result = op.template get_destination_cooperative_tensor<decltype(a0), decltype(b0), T>(); |
| for (int e = 0; e < accum.get_capacity(); ++e) result[e] = T(accum[e] + float(bias_regs[e])); |
| result.store(D.template slice<BN, dynamic_extent>(tj * BN, ti * BM)); |
| """ |
| return mx.fast.metal_kernel( |
| name=f"h3_mpp_{bm}_{bn}_{bk}_{groups}_{walk}", |
| input_names=["x", "w", "bias"], |
| output_names=["out"], |
| source=source, |
| header=header, |
| ) |
|
|
|
|
| def gemm(x, w, bias, bm, bn, bk, groups, walk): |
| m, k = x.shape |
| n = w.shape[0] |
| if k % bk or n % bn: |
| raise ValueError("This experimental kernel requires aligned K and N") |
| return kernel(bm, bn, bk, groups, walk)( |
| inputs=[x, w, bias], |
| template=[ |
| ("T", x.dtype), |
| ("BM", bm), |
| ("BN", bn), |
| ("BK", bk), |
| ("SG", groups), |
| ("WALK", walk), |
| ], |
| grid=(((m + bm - 1) // bm) * (n // bn) * groups * 32, 1, 1), |
| threadgroup=(groups * 32, 1, 1), |
| output_shapes=[(m, n)], |
| output_dtypes=[x.dtype], |
| )[0] |
|
|
|
|
| def main(a): |
| configure_memory() |
| mx.random.seed(731) |
| dest = Path(a.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "warmup": 2, |
| "repeats": a.repeats, |
| "scope": "Original projection weights, synthetic BF16 inputs. Custom Metal 4 TensorOps with FP32 accumulation and BF16 bias epilogue.", |
| "trials": [], |
| } |
| store = WeightStore(Path(a.model) / "transformer") |
| for name in a.operators: |
| layer = Linear.load( |
| store, |
| "blocks.0." + name, |
| grouped_heads=(56, 128) if name == "attn.qkv_proj" else None, |
| ) |
| w = layer.weight |
| b = layer.bias if layer.bias is not None else mx.zeros(w.shape[0], w.dtype) |
| x = mx.random.normal((a.tokens, w.shape[1])).astype(w.dtype) |
| ref = layer(x) |
| mx.eval(x, w, b, ref) |
| variants = [(0, 0, 0, 0, 0)] |
| variants += [ |
| (bm, bn, bk, sg, walk) |
| for bm, bn, sg in ( |
| (32, 64, 2), |
| (64, 64, 4), |
| (64, 128, 4), |
| (128, 64, 4), |
| (128, 128, 4), |
| (128, 128, 8), |
| (64, 128, 8), |
| ) |
| for bk in (128, 256, 512) |
| if w.shape[1] % bk == 0 |
| for walk in (0, 8) |
| ] |
| for bm, bn, bk, sg, walk in variants: |
| row = { |
| "operator": name, |
| "tokens": a.tokens, |
| "tile": [bm, bn, bk], |
| "simdgroups": sg, |
| "walk_m": walk, |
| } |
| try: |
| fn = ( |
| (lambda layer=layer, x=x: layer(x)) |
| if bm == 0 |
| else ( |
| lambda x=x, w=w, b=b, bm=bm, bn=bn, bk=bk, sg=sg, walk=walk: ( |
| gemm(x, w, b, bm, bn, bk, sg, walk) |
| ) |
| ) |
| ) |
| row.update(measure(fn, repeats=a.repeats), **error(ref, fn())) |
| 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) |
| if "error" in row and bm == 32: |
| |
| return |
| del layer, w, b, x, ref |
| gc.collect() |
| mx.clear_cache() |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--tokens", type=int, default=4181) |
| p.add_argument( |
| "--operators", |
| nargs="+", |
| default=["attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"], |
| ) |
| p.add_argument("--repeats", type=int, default=5) |
| p.add_argument("--output", default="reports/optimization-v2/mpp.json") |
| main(p.parse_args()) |
|
|
