File size: 6,681 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """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:
# A syntax or SDK incompatibility must be fixed before a sweep.
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())
|