Merlin78
120 lines
3.1 KiB
Python
120 lines
3.1 KiB
Python
# Copyright © 2026 Apple Inc.
|
|
|
|
import math
|
|
import time
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import torch
|
|
|
|
N_WARMUP = 5
|
|
N_BENCH = 20
|
|
|
|
|
|
def bench_mlx(a, b):
|
|
for _ in range(N_WARMUP):
|
|
mx.eval(a @ b)
|
|
|
|
times = []
|
|
for _ in range(N_BENCH):
|
|
start = time.perf_counter_ns()
|
|
mx.eval(a @ b)
|
|
end = time.perf_counter_ns()
|
|
times.append((end - start) * 1e-9)
|
|
|
|
return np.mean(times), np.std(times)
|
|
|
|
|
|
@torch.no_grad()
|
|
def bench_torch(a, b):
|
|
for _ in range(N_WARMUP):
|
|
_ = a @ b
|
|
torch.mps.synchronize()
|
|
|
|
times = []
|
|
for _ in range(N_BENCH):
|
|
start = time.perf_counter_ns()
|
|
_ = a @ b
|
|
torch.mps.synchronize()
|
|
end = time.perf_counter_ns()
|
|
times.append((end - start) * 1e-9)
|
|
|
|
return np.mean(times), np.std(times)
|
|
|
|
|
|
def check_correctness(out_mx, out_pt, rtol, M, N, K):
|
|
if not np.allclose(out_pt, out_mx, rtol=rtol, atol=0):
|
|
abs_diff = np.abs(out_pt - out_mx)
|
|
rel_diff = abs_diff / np.maximum(np.abs(out_pt), 1e-10)
|
|
|
|
print(
|
|
f" WARNING: Correctness failed at {M}x{N}x{K}: "
|
|
f"max_abs={np.max(abs_diff):.6e}, max_rel={np.max(rel_diff):.6e}"
|
|
)
|
|
|
|
|
|
def bench_gemm(M, N, K, dtype, rtol):
|
|
scale = 0.5 / math.sqrt(K)
|
|
a_np = np.random.uniform(0, scale, (M, K)).astype(np.float32)
|
|
b_np = np.random.uniform(0, scale, (K, N)).astype(np.float32)
|
|
|
|
a_mx = mx.array(a_np).astype(getattr(mx, dtype))
|
|
b_mx = mx.array(b_np).astype(getattr(mx, dtype))
|
|
|
|
a_pt = torch.from_numpy(a_np).to(dtype=getattr(torch, dtype), device="mps")
|
|
b_pt = torch.from_numpy(b_np).to(dtype=getattr(torch, dtype), device="mps")
|
|
torch.mps.synchronize()
|
|
|
|
torch_mean, torch_std = bench_torch(a_pt, b_pt)
|
|
mlx_mean, mlx_std = bench_mlx(a_mx, b_mx)
|
|
|
|
out_mx = (a_mx @ b_mx).astype(mx.float32)
|
|
out_pt = (a_pt @ b_pt).to(torch.float32).to("cpu").numpy(force=True)
|
|
check_correctness(out_mx, out_pt, rtol, M, N, K)
|
|
|
|
return mlx_mean, mlx_std, torch_mean, torch_std
|
|
|
|
|
|
if __name__ == "__main__":
|
|
dtypes = ("bfloat16", "float16", "float32")
|
|
|
|
rtols = {
|
|
"float32": 1e-3,
|
|
"float16": 5e-3,
|
|
"bfloat16": 1e-2,
|
|
}
|
|
|
|
shapes = (
|
|
(2048, 2048, 10240),
|
|
(2048, 3072, 10240),
|
|
(3072, 3072, 10240),
|
|
(3072, 3072, 12288),
|
|
(3072, 4096, 12288),
|
|
(4096, 4096, 12288),
|
|
(4096, 4096, 18432),
|
|
(4096, 4096, 21504),
|
|
(4096, 6144, 21504),
|
|
(6144, 6144, 21504),
|
|
)
|
|
|
|
for dtype in dtypes:
|
|
print(f"\nPerformance ({dtype}):")
|
|
print(
|
|
f"{'M':>5s} {'N':>5s} {'K':>6s} "
|
|
f"{'MLX (ms)':>15s} {'Torch (ms)':>15s} {'Speedup':>10s}"
|
|
)
|
|
print("-" * 80)
|
|
|
|
for M, N, K in shapes:
|
|
mlx_mean, mlx_std, torch_mean, torch_std = bench_gemm(
|
|
M, N, K, dtype, rtols[dtype]
|
|
)
|
|
speedup = torch_mean / mlx_mean
|
|
|
|
print(
|
|
f"{M:5d} {N:5d} {K:6d} "
|
|
f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f} "
|
|
f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f} "
|
|
f"{speedup:8.2f}x"
|
|
)
|