mlx 1.0.0

data/mlx/benchmarks/python/rope_bench.py

@@ -0,0 +1,35 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+from time_utils import time_fn
+
+
+def time_rope():
+    rope = nn.RoPE(64)
+
+    # vec
+    x = mx.random.uniform(shape=(1, 32, 1, 128)).astype(mx.float16)
+    mx.eval(x)
+
+    def rope_vec(x):
+        for _ in range(32):
+            x = rope(x, offset=100)
+        return x
+
+    time_fn(rope_vec, x)
+
+    # matrix
+    x = mx.random.uniform(shape=(1, 32, 1024, 128)).astype(mx.float16)
+    mx.eval(x)
+
+    def rope_mat(x):
+        for _ in range(32):
+            x = rope(x)
+        return x
+
+    time_fn(rope_mat, x)
+
+
+if __name__ == "__main__":
+    time_rope()

data/mlx/benchmarks/python/scatter_bench.py

@@ -0,0 +1,96 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+import torch
+from time_utils import measure_runtime
+
+
+def benchmark_scatter_mlx(dst_shape, x_shape, idx_shapes):
+    def scatter(dst, x, idx):
+        dst[tuple(idx)] = x
+        mx.eval(dst)
+
+    idx = []
+    for idx_shape in idx_shapes:
+        idx.append(mx.random.randint(0, dst_shape[0] - 1, idx_shape))
+    x = mx.random.normal(x_shape).astype(mx.float32)
+    dst = mx.random.normal(dst_shape).astype(mx.float32)
+
+    runtime = measure_runtime(scatter, dst=dst, x=x, idx=idx)
+    print(f"MLX: {runtime:.3f}ms")
+
+
+def benchmark_scatter_torch(dst_shape, x_shape, idx_shapes, device):
+    def scatter(dst, x, idx, device):
+        dst[tuple(idx)] = x
+        if device == torch.device("mps"):
+            torch.mps.synchronize()
+
+    idx = []
+    for idx_shape in idx_shapes:
+        idx.append(torch.randint(0, dst_shape[0] - 1, idx_shape).to(device))
+    x = torch.randn(x_shape, dtype=torch.float32).to(device)
+    dst = torch.randn(dst_shape, dtype=torch.float32).to(device)
+
+    runtime = measure_runtime(scatter, dst=dst, x=x, idx=idx, device=device)
+    print(f"PyTorch: {runtime:.3f}ms")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Gather benchmarks.")
+    parser.add_argument("--cpu", action="store_true", help="Use the CPU.")
+    args = parser.parse_args()
+
+    if args.cpu:
+        mx.set_default_device(mx.cpu)
+        device = torch.device("cpu")
+    else:
+        device = torch.device("mps")
+
+    dst_shapes = [
+        (10, 64),
+        (100_000, 64),
+        (1_000_000, 64),
+        (100_000,),
+        (200_000,),
+        (20_000_000,),
+        (10000, 64),
+        (100, 64),
+        (100, 10_000, 64),
+        (10, 100, 100, 21),
+        (1_000, 1_000, 10),
+    ]
+    idx_shapes = [
+        [(1_000_000,)],
+        [(1_000_000,)],
+        [(100_000,)],
+        [(1_000_000,)],
+        [(20_000_000,)],
+        [(20_000_000,)],
+        [(1000000,)],
+        [(10000000,)],
+        [(1_000,)],
+        [(10_000,)],
+        [(1_000,), (1_000,)],
+    ]
+    x_shapes = [
+        (1_000_000, 64),
+        (1_000_000, 64),
+        (100_000, 64),
+        (1_000_000,),
+        (20_000_000,),
+        (20_000_000,),
+        (1000000, 64),
+        (10000000, 64),
+        (1_000, 10_000, 64),
+        (10_000, 100, 100, 21),
+        (1_000, 10),
+    ]
+
+    for dst_shape, x_shape, idx_shape in zip(dst_shapes, x_shapes, idx_shapes):
+        print("=" * 20)
+        print(f"Dst: {dst_shape}, X {x_shape}, Indices {idx_shape}")
+        benchmark_scatter_mlx(dst_shape, x_shape, idx_shape)
+        benchmark_scatter_torch(dst_shape, x_shape, idx_shape, device=device)

data/mlx/benchmarks/python/sdpa_bench.py

@@ -0,0 +1,223 @@
+# Copyright © 2024 Apple Inc.
+
+import argparse
+import math
+import os
+import subprocess
+import time
+
+import mlx.core as mx
+import numpy as np
+
+device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+device_name = device_name.decode("utf-8").strip("\n")
+
+N_warmup = 5
+N_iter_bench = 40
+N_iter_func = 8
+
+
+def bench(f, *args):
+    for i in range(N_warmup):
+        f(*args)
+
+    s = time.perf_counter_ns()
+    for i in range(N_iter_bench):
+        f(*args)
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9
+
+
+def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype):
+    np_dtype = getattr(np, dtype)
+
+    shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D)
+    shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D)
+
+    scale = 1.0 / math.sqrt(D)
+
+    q_np = np.random.normal(0.0, 1.0, shape_q).astype(np_dtype)
+    k_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
+    v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
+
+    q_mx = mx.array(q_np)
+    k_mx = mx.array(k_np)
+    v_mx = mx.array(v_np)
+
+    if mask is not None:
+        if mask == "additive":
+            mask_np = np.random.normal(0.0, 1.0, (B, qH, qL, kL)).astype(np_dtype)
+            mask = mx.array(mask_np)
+        elif mask == "bool":
+            mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
+            mask = mx.array(mask_np)
+
+    return q_mx, k_mx, v_mx, scale, mask
+
+
+def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
+    q_dtype = q.dtype
+    q = q * mx.array(scale, q_dtype)
+    n_q_heads = q.shape[-3]
+    n_kv_heads = k.shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    B = q.shape[0]
+    L = q.shape[2]
+    kL = k.shape[2]
+
+    if n_repeats > 1:
+        q = mx.reshape(q, [B, n_kv_heads, n_repeats, L, -1])
+        k = mx.expand_dims(k, 2)
+        v = mx.expand_dims(v, 2)
+
+    scores = q @ mx.swapaxes(k, -1, -2)
+
+    if mask is not None:
+
+        if mask == "causal":
+            q_offset = max(0, kL - L)
+            q_indices = mx.arange(q_offset, q_offset + L)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+
+        if n_repeats > 1 and mask.ndim >= 3:
+            if mask.shape[-3] == 1:
+                mask = mx.expand_dims(mask, -3)
+            else:
+                mask = mx.unflatten(mask, -3, (n_kv_heads, n_repeats))
+
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, -np.float32(np.inf))
+        else:
+            scores += mask
+
+    scores = mx.softmax(scores, axis=-1, precise=True)
+
+    out = scores @ v
+    if n_repeats > 1:
+        out = mx.reshape(out, [B, n_q_heads, L, -1])
+
+    return out
+
+
+def mlx_fused_attn(q, k, v, scale, mask):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
+
+
+def do_attention(f, q, k, v, scale, mask=None, transpose=False):
+    if transpose:
+        q_t = mx.transpose(q, (0, 2, 1, 3))
+        k_t = mx.transpose(k, (0, 2, 1, 3))
+        v_t = mx.transpose(v, (0, 2, 1, 3))
+        o_t = f(q_t, k_t, v_t, scale=scale, mask=mask)
+        return mx.transpose(o_t, (0, 2, 1, 3))
+    else:
+        return f(q, k, v, scale=scale, mask=mask)
+
+
+def do_attention_bench(f, q, k, v, scale, mask=None, transpose=False):
+    q_out = q
+
+    for i in range(N_iter_func):
+        q_out = do_attention(f, q_out, k, v, scale, mask=mask, transpose=transpose)
+
+    mx.eval(q_out)
+    return q_out
+
+
+def bench_shape(
+    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None
+):
+    q_mx, k_mx, v_mx, scale, mask = prepare_inputs(
+        B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, mask_in, transpose, dtype
+    )
+
+    time_mlx_unfused = bench(
+        do_attention_bench, mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose
+    )
+    time_mlx_fused = bench(
+        do_attention_bench, mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
+    )
+
+    o_mlx_fused = do_attention(mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose)
+    o_mlx_unfused = do_attention(
+        mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
+    )
+
+    atol = 1e-5 if dtype == "float32" else 2e-4
+
+    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol, rtol=atol):
+        print(
+            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
+        )
+
+    return time_mlx_fused, time_mlx_unfused
+
+
+def get_gflop_count(B, M, N, K):
+    return float(2.0 * N_iter_bench * N_iter_func * B * M * N * K) / float(1024.0**3)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run gemm benchmarks")
+
+    dtypes = ("float16", "float32")[:1]
+    transposes = (False,)
+
+    # fmt: off
+    shapes_64 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,    32,    32,       64,   32,    32),
+          (  1,    64,    64,       64,   32,    32),
+          (  1,   128,   128,       64,   32,    32),
+          (  1,   256,   256,       64,   32,    32),
+          (  1,   512,   512,       64,   32,    32),
+          (  1,  1024,  1024,       64,   32,     8),
+          (  1,  2048,  2048,       64,   32,     8),
+          (  1,  4096,  4096,       64,   32,     8),
+    )
+
+    shapes_80 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,  1024,  1024,       80,   32,     8),
+          (  1,  2048,  2048,       80,   32,     8),
+          (  1,  4096,  4096,       80,   32,     8),
+    )
+
+    shapes_128 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,  1024,  1024,      128,   32,     8),
+          (  1,  2048,  2048,      128,   32,     8),
+          (  1,  4096,  4096,      128,   32,     8),
+    )
+    # fmt: on
+
+    shapes = shapes_64 + shapes_80 + shapes_128
+
+    masks = [None, "bool", "causal"]
+
+    print(
+        "  B,   qsl,   ksl, hdim, n_qh, n_kvh, t,   dtype,     mask, t_unfs, t_fuse, diff%"
+    )
+
+    for dtype in dtypes:
+        for transpose in transposes:
+            for B, qsl, ksl, head_dim, n_q_heads, n_kv_heads in shapes:
+                for mask_in in masks:
+                    time_mlx_fused, time_mlx_unfused = bench_shape(
+                        B,
+                        qsl,
+                        ksl,
+                        head_dim,
+                        n_q_heads,
+                        n_kv_heads,
+                        dtype,
+                        transpose,
+                        mask_in,
+                    )
+                    diff = time_mlx_unfused / time_mlx_fused - 1.0
+                    t_str = 1 if transpose else 0
+                    print(
+                        f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:1d}, {dtype}, {str(mask_in):>8}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
+                    )

data/mlx/benchmarks/python/sdpa_vector_bench.py

@@ -0,0 +1,95 @@
+import argparse
+import math
+
+import mlx.core as mx
+from time_utils import time_fn
+
+L = 16384
+H = 32
+H_k = H // 4
+D = 128
+V = 128
+dtype = mx.float16
+loops = 10
+
+
+def upproject(x, w):
+    if w is None:
+        return x
+    else:
+        return x @ w.T
+
+
+def attention(q, k, v, mask=None, w=None):
+    def _sdpa(q, k, v):
+        B, Hq, L, D = q.shape
+        _, Hk, S, _ = k.shape
+        _, _, _, V = v.shape
+        q = q.reshape(B, Hk, Hq // Hk, L, D)
+        k = k[:, :, None, :, :]
+        v = v[:, :, None, :, :]
+        s = q @ k.transpose(0, 1, 2, 4, 3)
+        if mask is not None:
+            m = mx.broadcast_to(mask, (B, Hq, L, S)).reshape(B, Hk, Hq // Hk, L, S)
+            s = mx.where(m, s, mx.finfo(s.dtype).min)
+        p = mx.softmax(s.astype(mx.float32), axis=-1).astype(s.dtype)
+        o = p @ v
+        return o.reshape(B, Hq, L, V)
+
+    for i in range(loops):
+        q = _sdpa(q, k, v)
+        q = upproject(q, w)
+    return q
+
+
+def sdpa(q, k, v, mask=None, w=None):
+    for i in range(loops):
+        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
+        q = upproject(q, w)
+    return q
+
+
+def time_self_attention_primitives():
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
+    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
+    mx.eval(q, k, v, w)
+    time_fn(attention, q, k, v, w=w)
+
+
+def time_self_attention_sdpa():
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
+    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
+    mx.eval(q, k, v, w)
+    time_fn(sdpa, q, k, v, w=w)
+
+
+def time_self_attention_sdpa_with_mask():
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
+    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
+    mask = mx.full((L,), True)
+    mask[L // 2 :] = False
+    mx.eval(q, k, v, mask, w)
+
+    def sdpa_mask(*args):
+        return sdpa(*args, mask=mask, w=w)
+
+    def attention_mask(*args):
+        return attention(*args, mask=mask, w=w)
+
+    time_fn(attention_mask, q, k, v)
+    time_fn(sdpa_mask, q, k, v)
+
+
+if __name__ == "__main__":
+    time_self_attention_sdpa()
+    time_self_attention_primitives()
+    time_self_attention_sdpa_with_mask()

data/mlx/benchmarks/python/single_ops.py

@@ -0,0 +1,132 @@
+# Copyright © 2023 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+from time_utils import time_fn
+
+
+def time_add():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    b = mx.random.uniform(shape=(32, 1024, 1024))
+    mx.eval(a, b)
+    time_fn(mx.add, a, b)
+
+    aT = mx.transpose(a, [0, 2, 1])
+    mx.eval(aT)
+
+    def transpose_add(a, b):
+        return mx.add(a, b)
+
+    time_fn(transpose_add, aT, b)
+
+    b = mx.random.uniform(shape=(1024,))
+    mx.eval(b)
+
+    def slice_add(a, b):
+        return mx.add(a, b)
+
+    time_fn(slice_add, a, b)
+
+    b = mx.reshape(b, (1, 1024, 1))
+    mx.eval(b)
+
+    def mid_slice_add(a, b):
+        return mx.add(a, b)
+
+    time_fn(mid_slice_add, a, b)
+
+
+def time_matmul():
+    a = mx.random.uniform(shape=(1024, 1024))
+    b = mx.random.uniform(shape=(1024, 1024))
+    mx.eval(a, b)
+    time_fn(mx.matmul, a, b)
+
+
+def time_maximum():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    b = mx.random.uniform(shape=(32, 1024, 1024))
+    mx.eval(a, b)
+    time_fn(mx.maximum, a, b)
+
+
+def time_max():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    a[1, 1] = mx.nan
+    mx.eval(a)
+    time_fn(mx.max, a, 0)
+
+
+def time_min():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    a[1, 1] = mx.nan
+    mx.eval(a)
+    time_fn(mx.min, a, 0)
+
+
+def time_negative():
+    a = mx.random.uniform(shape=(10000, 1000))
+    mx.eval(a)
+
+    def negative(a):
+        return -a
+
+    mx.eval(a)
+
+    time_fn(negative, a)
+
+
+def time_exp():
+    a = mx.random.uniform(shape=(1000, 100))
+    mx.eval(a)
+    time_fn(mx.exp, a)
+
+
+def time_logsumexp():
+    a = mx.random.uniform(shape=(64, 10, 10000))
+    mx.eval(a)
+    time_fn(mx.logsumexp, a, axis=-1)
+
+
+def time_take():
+    a = mx.random.uniform(shape=(10000, 500))
+    ids = mx.random.randint(low=0, high=10000, shape=(20, 10))
+    ids = [mx.reshape(idx, (-1,)) for idx in ids]
+    mx.eval(ids)
+
+    def random_take():
+        return [mx.take(a, idx, 0) for idx in ids]
+
+    time_fn(random_take)
+
+
+def time_reshape_transposed():
+    x = mx.random.uniform(shape=(256, 256, 128))
+    mx.eval(x)
+
+    def reshape_transposed():
+        return mx.reshape(mx.transpose(x, (1, 0, 2)), (-1,))
+
+    time_fn(reshape_transposed)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("MLX benchmarks.")
+    parser.add_argument("--gpu", action="store_true", help="Use the Metal back-end.")
+    args = parser.parse_args()
+    if args.gpu:
+        mx.set_default_device(mx.gpu)
+    else:
+        mx.set_default_device(mx.cpu)
+
+    time_add()
+    time_matmul()
+    time_min()
+    time_max()
+    time_maximum()
+    time_exp()
+    time_negative()
+    time_logsumexp()
+    time_take()
+    time_reshape_transposed()

data/mlx/benchmarks/python/synchronize_bench.py

@@ -0,0 +1,55 @@
+import time
+
+import mlx.core as mx
+
+rank = mx.distributed.init().rank()
+
+
+def timeit(fn, a):
+
+    # warmup
+    for _ in range(5):
+        mx.eval(fn(a))
+
+    its = 10
+    tic = time.perf_counter()
+    for _ in range(its):
+        mx.eval(fn(a))
+    toc = time.perf_counter()
+    ms = 1000 * (toc - tic) / its
+    return ms
+
+
+def all_reduce_benchmark():
+    a = mx.ones((5, 5), mx.int32)
+
+    its_per_eval = 100
+
+    def fn(x):
+        for _ in range(its_per_eval):
+            x = mx.distributed.all_sum(x)
+            x = x - 1
+        return x
+
+    ms = timeit(fn, a) / its_per_eval
+    if rank == 0:
+        print(f"All Reduce: time per iteration {ms:.6f} (ms)")
+
+
+def all_gather_benchmark():
+    a = mx.ones((5, 5), mx.int32)
+    its_per_eval = 100
+
+    def fn(x):
+        for _ in range(its_per_eval):
+            x = mx.distributed.all_gather(x)[0]
+        return x
+
+    ms = timeit(fn, a) / its_per_eval
+    if rank == 0:
+        print(f"All gather: time per iteration {ms:.6f} (ms)")
+
+
+if __name__ == "__main__":
+    all_reduce_benchmark()
+    all_gather_benchmark()

data/mlx/benchmarks/python/time_utils.py

@@ -0,0 +1,38 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import time
+
+import mlx.core as mx
+
+
+def time_fn(fn, *args, **kwargs):
+    msg = kwargs.pop("msg", None)
+    if msg:
+        print(f"Timing {msg} ...", end=" ")
+    else:
+        print(f"Timing {fn.__name__} ...", end=" ")
+
+    # warmup
+    for _ in range(5):
+        mx.eval(fn(*args, **kwargs))
+
+    num_iters = 100
+    tic = time.perf_counter()
+    for _ in range(num_iters):
+        x = mx.eval(fn(*args, **kwargs))
+    toc = time.perf_counter()
+
+    msec = 1e3 * (toc - tic) / num_iters
+    print(f"{msec:.5f} msec")
+
+
+def measure_runtime(fn, **kwargs):
+    # Warmup
+    for _ in range(5):
+        fn(**kwargs)
+
+    tic = time.perf_counter()
+    iters = 100
+    for _ in range(iters):
+        fn(**kwargs)
+    return (time.perf_counter() - tic) * 1000 / iters