Refresh MLX fork onto current upstream cutoff

.github/actions/setup/action.yml

@@ -52,8 +52,10 @@ runs:
       shell: bash
       run: |
         echo "::group::Install macOS dependencies"
+        brew update
         brew install openmpi
         xcodebuild -showComponent MetalToolchain
+        sysctl -a | grep machdep.cpu
         echo "::endgroup::"
 
     - name: Setup Windows environment
@@ -90,15 +92,16 @@ runs:
         key: v7-${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.ccache-toolkit || inputs.toolkit }}
         max-size: |-
           ${{ case(inputs.ccache-key == 'release',
-                   case(startsWith(inputs.toolkit, 'cuda'), '150MB',
+                   case(startsWith(inputs.toolkit, 'cuda'),
+                        case(runner.os == 'Linux', '300MB',
+                             '160MB'),
                         '60MB'),
                    case(startsWith(inputs.toolkit, 'cuda'),
-                        case(runner.os == 'Linux' && runner.arch == 'x64', '500MB',
-                             runner.os == 'Linux', '240MB',
-                             '120MB'),
-                        runner.os == 'macOS' && inputs.toolkit == 'metal', '200MB',
-                        runner.os == 'macOS', '150MB',
-                        '60MB'))
+                        case(runner.os == 'Linux' && runner.arch == 'x64', '600MB',
+                             runner.os == 'Linux', '400MB',
+                             '320MB'),
+                        runner.os == 'macOS' && inputs.toolkit == 'metal', '300MB',
+                        '200MB'))
           }}
         save: ${{ !startsWith(github.ref, 'refs/pull/') && (inputs.ccache-save != 'false') }}
         # ccache-action bug: running "apt-get update" fails on large arm runner.

ACKNOWLEDGMENTS.md

@@ -20,6 +20,8 @@ MLX was developed with contributions from the following individuals:
 - Paul Paczuski: Improved stability of BCE loss calculation
 - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops.
 - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.
+- katlun-lgtm: Added `reflect` and `symmetric` padding modes.
+- Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting.
 
 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />

CMakeLists.txt

@@ -151,6 +151,7 @@ cmake_policy(SET CMP0135 NEW)
 
 add_library(mlx)
 
+target_compile_features(mlx INTERFACE cxx_std_20)
 target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
 target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
 
@@ -261,7 +262,8 @@ if(WIN32)
   FetchContent_MakeAvailable(dlfcn-win32)
   endblock()
   target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src")
-  target_link_libraries(mlx PRIVATE dl)
+  target_sources(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src/dlfcn.c")
+  target_link_libraries(mlx PRIVATE Psapi.lib)
 endif()
 
 if(MLX_BUILD_CPU)
@@ -296,10 +298,15 @@ if(MLX_BUILD_CPU)
       URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.33/${OPENBLAS_ZIP}"
     )
     FetchContent_MakeAvailable(openblas)
-    target_link_libraries(
-      mlx PRIVATE "${openblas_SOURCE_DIR}/lib/${OPENBLAS_LIB}.lib")
     target_include_directories(mlx
                                PRIVATE "${openblas_SOURCE_DIR}/${OPENBLAS_INC}")
+    # Make openblas importable by dependencies when built as static library.
+    set(OPENBLAS_IMPORT_LIBRARY
+        "${openblas_SOURCE_DIR}/lib/${OPENBLAS_LIB}.lib")
+    add_library(MLX::OpenBLAS UNKNOWN IMPORTED GLOBAL)
+    set_target_properties(MLX::OpenBLAS PROPERTIES IMPORTED_LOCATION
+                                                   "${OPENBLAS_IMPORT_LIBRARY}")
+    target_link_libraries(mlx PRIVATE MLX::OpenBLAS)
     # Make sure the DLL file is placed in the same dir with executables.
     set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/${OPENBLAS_LIB}.dll")
     add_custom_command(
@@ -417,6 +424,8 @@ if(WIN32)
   if(MLX_BUILD_CPU)
     # Install OpenBLAS.
     install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN)
+    install(FILES ${OPENBLAS_IMPORT_LIBRARY}
+            DESTINATION ${CMAKE_INSTALL_LIBDIR})
   endif()
 endif()
 

benchmarks/cpp/irregular_strides.cpp

@@ -9,6 +9,12 @@
 
 namespace mx = mlx::core;
 
+// mx::astype is overloaded, so it cannot be passed directly to the timing
+// helpers. Wrap the three argument form instead.
+auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) {
+  return mx::astype(a, dtype, s);
+};
+
 void time_irregular_binary_ops_1D() {
   auto device = mx::default_device();
   int size = 1000000;
@@ -164,7 +170,7 @@ void time_irregular_astype_1D() {
   int step = 2;
   auto a = mx::random::uniform({size});
   a = slice(a, {0}, {size}, {step});
-  TIMEM("1D strided", mx::astype, a, mx::int32, device);
+  TIMEM("1D strided", astype, a, mx::int32, device);
 }
 
 void time_irregular_astype_2D() {
@@ -173,16 +179,16 @@ void time_irregular_astype_2D() {
   mx::Shape shape = {size, size};
 
   auto a = mx::random::uniform(shape);
-  TIMEM("2D regular", mx::astype, a, mx::int32, device);
+  TIMEM("2D regular", astype, a, mx::int32, device);
 
   a = mx::transpose(a);
-  TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device);
+  TIMEM("2D mx::transpose", astype, a, mx::int32, device);
 
   a = mx::broadcast_to(mx::random::uniform({size}), shape);
-  TIMEM("2D broadcast dim 0", mx::astype, a, mx::int32, device);
+  TIMEM("2D broadcast dim 0", astype, a, mx::int32, device);
 
   a = mx::broadcast_to(mx::random::uniform({size, 1}), shape);
-  TIMEM("2D broadcast dim 1", mx::astype, a, mx::int32, device);
+  TIMEM("2D broadcast dim 1", astype, a, mx::int32, device);
 }
 
 int main(int argc, char** argv) {

benchmarks/cpp/single_ops.cpp

@@ -5,6 +5,12 @@
 
 namespace mx = mlx::core;
 
+// mx::astype is overloaded, so it cannot be passed directly to the timing
+// helpers. Wrap the three argument form instead.
+auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) {
+  return mx::astype(a, dtype, s);
+};
+
 void time_creation_ops() {
   int M = 2000;
   int N = 500;
@@ -28,18 +34,18 @@ void time_type_conversions() {
 
   auto a = mx::zeros(shape, mx::float32);
   mx::eval(a);
-  TIMEM("mx::float32 to mx::int32", mx::astype, a, mx::int32, device);
-  TIMEM("mx::float32 to mx::uint32", mx::astype, a, mx::uint32, device);
+  TIMEM("mx::float32 to mx::int32", astype, a, mx::int32, device);
+  TIMEM("mx::float32 to mx::uint32", astype, a, mx::uint32, device);
 
   a = mx::zeros(shape, mx::int32);
   mx::eval(a);
-  TIMEM("mx::int32 to mx::float32", mx::astype, a, mx::float32, device);
+  TIMEM("mx::int32 to mx::float32", astype, a, mx::float32, device);
 
   a = mx::zeros(shape, mx::bool_);
   mx::eval(a);
-  TIMEM("bool to mx::float32", mx::astype, a, mx::float32, device);
-  TIMEM("bool to mx::int32", mx::astype, a, mx::int32, device);
-  TIMEM("bool to mx::uint32", mx::astype, a, mx::uint32, device);
+  TIMEM("bool to mx::float32", astype, a, mx::float32, device);
+  TIMEM("bool to mx::int32", astype, a, mx::int32, device);
+  TIMEM("bool to mx::uint32", astype, a, mx::uint32, device);
 }
 
 void time_random_generation() {

benchmarks/python/comparative/bench_mlx.py

@@ -72,12 +72,17 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits):
 
 
 quant_matmul = {
+    "quant_matmul_32_1": partial(_quant_matmul, transpose=False, group_size=32, bits=1),
     "quant_matmul_32_2": partial(_quant_matmul, transpose=False, group_size=32, bits=2),
     "quant_matmul_32_4": partial(_quant_matmul, transpose=False, group_size=32, bits=4),
     "quant_matmul_32_8": partial(_quant_matmul, transpose=False, group_size=32, bits=8),
+    "quant_matmul_64_1": partial(_quant_matmul, transpose=False, group_size=64, bits=1),
     "quant_matmul_64_2": partial(_quant_matmul, transpose=False, group_size=64, bits=2),
     "quant_matmul_64_4": partial(_quant_matmul, transpose=False, group_size=64, bits=4),
     "quant_matmul_64_8": partial(_quant_matmul, transpose=False, group_size=64, bits=8),
+    "quant_matmul_128_1": partial(
+        _quant_matmul, transpose=False, group_size=128, bits=1
+    ),
     "quant_matmul_128_2": partial(
         _quant_matmul, transpose=False, group_size=128, bits=2
     ),
@@ -87,6 +92,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits):
     "quant_matmul_128_8": partial(
         _quant_matmul, transpose=False, group_size=128, bits=8
     ),
+    "quant_matmul_t_32_1": partial(
+        _quant_matmul, transpose=True, group_size=32, bits=1
+    ),
     "quant_matmul_t_32_2": partial(
         _quant_matmul, transpose=True, group_size=32, bits=2
     ),
@@ -96,6 +104,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits):
     "quant_matmul_t_32_8": partial(
         _quant_matmul, transpose=True, group_size=32, bits=8
     ),
+    "quant_matmul_t_64_1": partial(
+        _quant_matmul, transpose=True, group_size=64, bits=1
+    ),
     "quant_matmul_t_64_2": partial(
         _quant_matmul, transpose=True, group_size=64, bits=2
     ),
@@ -105,6 +116,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits):
     "quant_matmul_t_64_8": partial(
         _quant_matmul, transpose=True, group_size=64, bits=8
     ),
+    "quant_matmul_t_128_1": partial(
+        _quant_matmul, transpose=True, group_size=128, bits=1
+    ),
     "quant_matmul_t_128_2": partial(
         _quant_matmul, transpose=True, group_size=128, bits=2
     ),
@@ -420,7 +434,22 @@ def selu(x):
         print(bench(matmul, *xs))
 
     elif args.benchmark.startswith("quant_matmul"):
-        print(bench(quant_matmul[args.benchmark], *xs))
+        # Parse group_size and bits from the benchmark name, e.g.
+        # "quant_matmul_128_4" or "quant_matmul_t_128_4"
+        fn = quant_matmul[args.benchmark]
+        gs = fn.keywords["group_size"]
+        bits = fn.keywords["bits"]
+        transpose = fn.keywords["transpose"]
+
+        # xs[0] = activation x, xs[1] = original (float) weight matrix
+        # Quantize the weight internally so the caller only needs:
+        #   --size MxK --size NxK  (transpose=True)  or  --size MxK --size KxN
+        w_float = xs[1].astype(mx.float16)
+        w_q, scales, biases = mx.quantize(w_float, group_size=gs, bits=bits)
+        mx.eval(w_q, scales, biases)
+        x_input = xs[0].astype(mx.float16)
+        mx.eval(x_input)
+        print(bench(_quant_matmul, x_input, w_q, scales, biases, transpose, gs, bits))
 
     elif args.benchmark == "linear":
         if args.fused:

benchmarks/python/comparative/compare.py

@@ -4,6 +4,7 @@
 
 import argparse
 import re
+import sys
 from pathlib import Path
 from subprocess import run
 
@@ -22,15 +23,27 @@ def run_or_raise(*args, **kwargs):
 
 
 def compare(args):
-    t_mlx = run_or_raise(["python", BENCH_MLX] + args)
-    t_torch = run_or_raise(["python", BENCH_TORCH] + args)
+    t_mlx = run_or_raise([sys.executable, BENCH_MLX] + args)
+    t_torch = run_or_raise([sys.executable, BENCH_TORCH] + args)
 
     print((t_torch - t_mlx) / t_torch, " ".join(args), sep="\t")
 
 
+def compare_mlx_quant(args_base, bits_list):
+    """Compare quantized matmul across bit widths (MLX only, no PyTorch)."""
+    results = {}
+    for bits in bits_list:
+        bench_args = args_base.replace("{bits}", str(bits)).split()
+        results[bits] = run_or_raise(["python", BENCH_MLX] + bench_args)
+    baseline = max(results.values())
+    for bits in bits_list:
+        speedup = (baseline - results[bits]) / baseline if baseline > 0 else 0
+        print(f"{speedup:.4f}\t{args_base.replace('{bits}', str(bits))}")
+
+
 def compare_mlx_dtypes(args, dt1, dt2):
-    t_mlx_dt1 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt1])
-    t_mlx_dt2 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt2])
+    t_mlx_dt1 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt1])
+    t_mlx_dt2 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt2])
 
     print((t_mlx_dt2 - t_mlx_dt1) / t_mlx_dt2, " ".join(args), sep="\t")
 
@@ -282,3 +295,26 @@ def predicate(x):
     compare_filtered("topk --size 32768x128 --axis 1")
     compare_filtered("topk --size 128x128 --axis 0 --cpu")
     compare_filtered("topk --size 128x128 --axis 1 --cpu")
+
+    # Quantized matmul ops (MLX only — compare across bit widths)
+    # qmv path (M=1, token generation, memory-bandwidth bound)
+    for gs in [64, 128]:
+        compare_mlx_quant(
+            f"quant_matmul_t_{gs}_{{bits}} --size 1x4096 --size 4096x4096",
+            [1, 2, 4, 8],
+        )
+        compare_mlx_quant(
+            f"quant_matmul_t_{gs}_{{bits}} --size 1x4096 --size 11008x4096",
+            [1, 2, 4, 8],
+        )
+    # qmm path (prompt processing, more compute bound)
+    for gs in [64, 128]:
+        for M in [32, 512]:
+            compare_mlx_quant(
+                f"quant_matmul_t_{gs}_{{bits}} --size {M}x4096 --size 4096x4096",
+                [1, 2, 4, 8],
+            )
+            compare_mlx_quant(
+                f"quant_matmul_t_{gs}_{{bits}} --size {M}x4096 --size 11008x4096",
+                [1, 2, 4, 8],
+            )

benchmarks/python/gather_qmm_bench.py

@@ -80,5 +80,26 @@ def equivalent_matmul(x, w1, w2):
     time_fn(equivalent_matmul, x, w1, w2)
 
 
+def time_gather_qmm_short_runs():
+    # Many experts and few tokens, so each expert gets N * I / E = 16 rows.
+    N, E, I = 512, 256, 8
+    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
+    w1 = mx.random.normal((E, M, D)) / 1024**0.5
+    w2 = mx.random.normal((E, D, M)) / 1024**0.5
+    w1 = mx.quantize(w1)
+    w2 = mx.quantize(w2)
+    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
+    mx.eval(x, w1, w2, indices)
+
+    def gather_mm(x, w1, w2, indices):
+        x, idx, inv_order = gather_sort(x, indices)
+        x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=True)
+        x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=True)
+        return scatter_unsort(x, inv_order, indices.shape)
+
+    time_fn(gather_mm, x, w1, w2, indices)
+
+
 if __name__ == "__main__":
     time_gather_qmm()
+    time_gather_qmm_short_runs()

benchmarks/python/sdpa_bench.py

@@ -189,6 +189,15 @@ def get_gflop_count(B, M, N, K):
           (  1,  2048,  32121,      80,   32,     8),
     )
 
+    shapes_96 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,  1024,  1024,       96,   32,     8),
+          (  1,  2048,  2048,       96,   32,     8),
+          (  1,  4096,  4096,       96,   32,     8),
+          (  1,  4096,  5000,       96,   32,     8),
+          (  1,  2048,  32121,      96,   32,     8),
+    )
+
     shapes_128 = (
         # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
           (  1,  1024,  1024,      128,   32,     8),
@@ -199,7 +208,7 @@ def get_gflop_count(B, M, N, K):
     )
     # fmt: on
 
-    shapes = shapes_64 + shapes_80 + shapes_128
+    shapes = shapes_64 + shapes_80 + shapes_96 + shapes_128
 
     masks = [None, "bool", "causal"]
 

benchmarks/python/synchronize_bench.py

@@ -42,7 +42,7 @@ def all_gather_benchmark():
 
     def fn(x):
         for _ in range(its_per_eval):
-            x = mx.distributed.all_gather(x)[0]
+            x = mx.distributed.all_gather(x)[: a.shape[0]]
         return x
 
     ms = timeit(fn, a) / its_per_eval

cmake/extension.cmake

@@ -26,7 +26,8 @@ macro(mlx_build_metallib)
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
   # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions
+                             -Wmetal-addr-spaces)
   if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
     set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
                                -frecord-sources)

docs/src/conf.py

@@ -28,7 +28,10 @@
 
 python_use_unqualified_type_names = True
 autosummary_generate = True
-autosummary_filename_map = {"mlx.core.Stream": "stream_class"}
+autosummary_filename_map = {
+    "mlx.core.Stream": "stream_class",
+    "mlx.core.PrintOptions": "printoptions_class",
+}
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),