import argparse
import time
import mlx.core as mx
import numpy as np
MLX_DTYPES = {
"float16": mx.float16,
"bfloat16": mx.bfloat16,
"float32": mx.float32,
}
def parse_cases(cases):
parsed = []
for spec in cases.split(","):
parts = spec.split("x")
m, n, k, bs = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
sparsity = float(parts[4]) if len(parts) > 4 else 0.5
parsed.append((m, n, k, bs, sparsity))
return parsed
def make_masks(m, n, k, block_size, sparsity, rng):
tm = (m + block_size - 1) // block_size
tn = (n + block_size - 1) // block_size
tk = (k + block_size - 1) // block_size
lhs_mask = (rng.random((tm, tk)) >= sparsity).astype(np.bool_)
rhs_mask = (rng.random((tk, tn)) >= sparsity).astype(np.bool_)
out_mask = (rng.random((tm, tn)) >= sparsity).astype(np.bool_)
return lhs_mask, rhs_mask, out_mask
def mlx_naive_block_masked_mm(a, b, block_size, out_mask, lhs_mask, rhs_mask):
M, K = a.shape[-2], a.shape[-1]
N = b.shape[-1]
def expand(mask, rows, cols):
e = mx.repeat(mx.repeat(mask, block_size, axis=-2), block_size, axis=-1)
return e[..., :rows, :cols]
a_masked = a * expand(lhs_mask, M, K)
b_masked = b * expand(rhs_mask, K, N)
c = a_masked @ b_masked
c = c * expand(out_mask, M, N)
return c
def bench_mlx(fn, warmup, iters):
for _ in range(warmup):
y = fn()
mx.eval(y)
mx.synchronize()
start = time.perf_counter()
for _ in range(iters):
y = fn()
mx.eval(y)
mx.synchronize()
return (time.perf_counter() - start) * 1e3 / iters
def print_table(headers, rows):
widths = [len(h) for h in headers]
for row in rows:
for i, cell in enumerate(row):
widths[i] = max(widths[i], len(cell))
def fmt_row(row):
return (
"| "
+ " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
+ " |"
)
sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
print(fmt_row(headers))
print(sep)
for row in rows:
print(fmt_row(row))
def main():
parser = argparse.ArgumentParser(
description="Benchmark block_masked_mm vs naive expand+matmul"
)
parser.add_argument(
"--cases",
default=(
"256x256x256x32x0.5,"
"512x512x512x32x0.5,"
"1024x1024x1024x32x0.5,"
"1024x1024x1024x64x0.5,"
"2048x2048x2048x64x0.5,"
"256x256x256x32x0.0,"
"1024x1024x1024x32x0.0,"
"1024x1024x1024x32x0.9"
),
help="Comma-separated MxNxKxBSxSparsity list. Sparsity=fraction of blocks zeroed.",
)
parser.add_argument(
"--dtype",
default="float32",
choices=["float16", "bfloat16", "float32"],
)
parser.add_argument("--warmup", type=int, default=10)
parser.add_argument("--iters", type=int, default=50)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--no-check", action="store_true")
args = parser.parse_args()
mlx_dtype = MLX_DTYPES[args.dtype]
print(f"dtype={args.dtype} warmup={args.warmup} iters={args.iters}")
headers = [
"Case (MxNxKxBS)",
"Sparsity",
"MLX ms",
"Naive ms",
"Speedup",
]
if not args.no_check:
headers.append("Max err")
rows = []
cases = parse_cases(args.cases)
for idx, (m, n, k, bs, sparsity) in enumerate(cases):
rng = np.random.default_rng(args.seed + idx)
a_np = rng.standard_normal((m, k)).astype(np.float32)
b_np = rng.standard_normal((k, n)).astype(np.float32)
lhs_mask_np, rhs_mask_np, out_mask_np = make_masks(m, n, k, bs, sparsity, rng)
a_mx = mx.array(a_np, dtype=mlx_dtype)
b_mx = mx.array(b_np, dtype=mlx_dtype)
lhs_mask_mx = mx.array(lhs_mask_np)
rhs_mask_mx = mx.array(rhs_mask_np)
out_mask_mx = mx.array(out_mask_np)
mx.eval(a_mx, b_mx, lhs_mask_mx, rhs_mask_mx, out_mask_mx)
err_str = ""
if not args.no_check:
y_op = mx.block_masked_mm(
a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
)
y_naive = mlx_naive_block_masked_mm(
a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
)
mx.eval(y_op, y_naive)
err = float(mx.max(mx.abs(y_op - y_naive)).item())
err_str = f"{err:.2e}"
t_mlx = bench_mlx(
lambda: mx.block_masked_mm(
a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
),
args.warmup,
args.iters,
)
t_naive = bench_mlx(
lambda: mlx_naive_block_masked_mm(
a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
),
args.warmup,
args.iters,
)
speedup = f"{t_naive / t_mlx:.2f}x" if t_mlx > 0 else "-"
row = [
f"{m}x{n}x{k}x{bs}",
f"{sparsity:.0%}",
f"{t_mlx:.3f}",
f"{t_naive:.3f}",
speedup,
]
if not args.no_check:
row.append(err_str)
rows.append(row)
print_table(headers, rows)
if not args.no_check:
print("err: max|block_masked_mm - naive_expand_matmul|")
if __name__ == "__main__":
main()