Merge branch 'ml-explore:main' into rocm-support

Author: Geramy

.github/actions/build-cuda-release/action.yml

@@ -18,7 +18,7 @@ runs:
       env:
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
       run: |
-        pip install auditwheel build patchelf setuptools
+        pip install auditwheel "build<=1.4.2" patchelf setuptools
         python setup.py clean --all
         MLX_DISABLE_SM90A_KERNELS=1 MLX_BUILD_STAGE=2 python -m build -w
 

.github/actions/build-linux-release/action.yml

@@ -25,7 +25,7 @@ runs:
     - name: Build Python package
       shell: bash
       run: |
-        pip install auditwheel patchelf build
+        pip install auditwheel patchelf "build<=1.4.2"
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
         auditwheel repair dist/mlx-*.whl \

.github/actions/setup-linux/action.yml

@@ -55,7 +55,7 @@ runs:
         echo "::endgroup::"
 
     - name: Set swap space
-      if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }}
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
       with:
         swap-size-gb: 16

CMakeLists.txt

@@ -361,6 +361,15 @@ FetchContent_MakeAvailable(json)
 target_include_directories(
   mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
 
+# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
+if(MLX_BUILD_CPU
+   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
+   AND DEFINED MACOS_SDK_VERSION
+   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
+                   ${CMAKE_BINARY_DIR}/jaccl)
+endif()
+
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(
@@ -388,7 +397,7 @@ if(MLX_BUILD_PYTHON_BINDINGS)
   FetchContent_Declare(
     nanobind
     GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v2.10.2
+    GIT_TAG v2.12.0
     GIT_SHALLOW TRUE
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(nanobind)

benchmarks/python/sdpa_bench.py

@@ -176,20 +176,26 @@ def get_gflop_count(B, M, N, K):
           (  1,  1024,  1024,       64,   32,     8),
           (  1,  2048,  2048,       64,   32,     8),
           (  1,  4096,  4096,       64,   32,     8),
+          (  1,  4096,  5000,       64,   32,     8),
+          (  1,  2048,  32121,      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),
+          (  1,  4096,  5000,       80,   32,     8),
+          (  1,  2048,  32121,      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),
+          (  1,  4096,  5000,      128,   32,     8),
+          (  1,  2048,  32121,     128,   32,     8),
     )
     # fmt: on
 

docs/src/python/devices_and_streams.rst

@@ -14,8 +14,10 @@ Devices and Streams
    set_default_device
    default_stream
    new_stream
+   new_thread_local_stream
    set_default_stream
    stream
    synchronize
+   clear_streams
    device_count
    device_info

examples/extensions/pyproject.toml

@@ -3,6 +3,6 @@ requires = [
   "setuptools>=42",
   "cmake>=3.25",
   "mlx>=0.18.0",
-  "nanobind==2.10.2",
+  "nanobind==2.12.0",
 ]
 build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
 mlx>=0.21.0
-nanobind==2.10.2
+nanobind==2.12.0

mlx/backend/cpu/sort.cpp

@@ -107,11 +107,11 @@ struct StridedIterator {
     return *this;
   }
 
-  StridedIterator operator+(difference_type diff) {
+  StridedIterator operator+(difference_type diff) const {
     return StridedIterator(ptr_, stride_, diff);
   }
 
-  StridedIterator operator-(difference_type diff) {
+  StridedIterator operator-(difference_type diff) const {
     return StridedIterator(ptr_, stride_, -diff);
   }
 

mlx/backend/cuda/CMakeLists.txt

@@ -31,6 +31,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/block_mask.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gather_gemm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
@@ -119,11 +120,11 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
-if(MSVC)
-  # Ignore warnings from CUTLASS.
-  target_compile_options(
-    mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=2908">)
-else()
+# Ignore warnings from CUTLASS.
+target_compile_options(
+  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=2908,2361">)
+
+if(NOT MSVC)
   # Required for generating optimized CUTLASS code.
   target_compile_options(
     mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
@@ -279,7 +280,7 @@ target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 FetchContent_Declare(
   cutlass
   GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-  GIT_TAG v4.3.5
+  GIT_TAG v4.4.2
   GIT_SHALLOW TRUE
   SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(cutlass)