Tentative JACCL examples

.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

mlx/backend/cuda/quantized/qmm/CMakeLists.txt

@@ -1,19 +1,35 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m16.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m32.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m64.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n128_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n256_m2.cu)
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu)
+
+foreach(TileN 16 32 64 128 256)
+  set(OUTPUT_FILE "qmm_sm90_impl_n${TileN}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm90.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  set(OUTPUT_FILE "qmm_sm80_impl_m${TileM}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm80.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  foreach(KMajor true false)
+    foreach(HasKResidue true false)
+      foreach(SM80 true false)
+        if(${KMajor} AND ${HasKResidue})
+          continue()
+        endif()
+        set(OUTPUT_FILE
+            "qmm_naive_impl_m${TileM}_${KMajor}_${HasKResidue}_${SM80}.cu")
+        configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_naive.cu"
+                       "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+        target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+      endforeach()
+    endforeach()
+  endforeach()
+endforeach()

mlx/backend/cuda/quantized/qmm/qmm.cu

@@ -17,9 +17,9 @@ inline bool is_last_2_dims_row_contiguous(const array& x) {
 } // namespace
 
 #if defined(MLX_CUDA_SM90A_ENABLED)
-// Defined in qmm_impl_sm90_xxx.cu files.
-template <typename TileShape, typename ClusterShape>
-void qmm_impl_sm90(
+// Defined in qmm_sm90.cu.
+template <int TileN>
+void qmm_sm90_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -83,24 +83,21 @@ void qmm_sm90(
     cu::CommandEncoder& encoder,
     Stream s) {
 #if defined(MLX_CUDA_SM90A_ENABLED)
-  auto dispatch = [&]<int tile_m, int tile_n, int cluster_m>() {
-    using cute::Int;
-    using TileShapeMN = cute::Shape<Int<tile_m>, Int<tile_n>>;
-    using ClusterShape = cute::Shape<Int<cluster_m>, Int<1>, Int<1>>;
-    qmm_impl_sm90<TileShapeMN, ClusterShape>(
+  auto dispatch = [&]<int TileN>() {
+    qmm_sm90_impl<TileN>(
         x, w, scales, biases, out, bits, group_size, encoder, s);
   };
   int m = out.ndim() > 1 ? out.shape(-2) : 1;
   if (m <= 16) {
-    dispatch.template operator()<128, 16, 1>();
+    dispatch.template operator()<16>();
   } else if (m <= 32) {
-    dispatch.template operator()<128, 32, 1>();
+    dispatch.template operator()<32>();
   } else if (m <= 64) {
-    dispatch.template operator()<128, 64, 2>();
+    dispatch.template operator()<64>();
   } else if (m <= 128) {
-    dispatch.template operator()<128, 128, 2>();
+    dispatch.template operator()<128>();
   } else {
-    dispatch.template operator()<128, 256, 2>();
+    dispatch.template operator()<256>();
   }
 #else
   throw std::runtime_error(
@@ -108,9 +105,9 @@ void qmm_sm90(
 #endif // defined(MLX_CUDA_SM90A_ENABLED)
 }
 
-// Defined in qmm_impl_sm80_xxx.cu files.
+// Defined in qmm_sm80.cu.
 template <int TileM>
-void qmm_impl_sm80(
+void qmm_sm80_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -174,7 +171,7 @@ void qmm_sm80(
     QuantizationMode mode,
     cu::CommandEncoder& encoder) {
   auto dispatch = [&]<int TileM>() {
-    qmm_impl_sm80<TileM>(
+    qmm_sm80_impl<TileM>(
         x,
         w,
         scales,
@@ -197,9 +194,9 @@ void qmm_sm80(
   }
 }
 
-// Defined in qmm_impl_naive_xxx.cu files.
-template <int TileM, bool KMajor>
-void qmm_impl_naive(
+// Defined in qmm_naive.cu.
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -250,8 +247,8 @@ void qmm_naive(
     int group_size,
     QuantizationMode mode,
     cu::CommandEncoder& encoder) {
-  auto dispatch = [&]<int TileM, bool KMajor>() {
-    qmm_impl_naive<TileM, KMajor>(
+  auto dispatch = [&]<int TileM, bool KMajor, bool HasKResidue, bool SM80>() {
+    qmm_naive_impl<TileM, KMajor, HasKResidue, SM80>(
         x,
         w,
         scales,
@@ -264,15 +261,37 @@ void qmm_naive(
         mode,
         encoder);
   };
-  dispatch_bool(transpose, [&](auto k_major) {
-    int m = out.ndim() > 1 ? out.shape(-2) : 1;
-    if (m <= 16) {
-      dispatch.template operator()<16, k_major.value>();
-    } else if (m <= 32) {
-      dispatch.template operator()<32, k_major.value>();
+  auto dispatch_k = [&](auto k_major, bool has_k_residue, auto&& f) {
+    if constexpr (k_major.value) {
+      if (has_k_residue) {
+        throw std::invalid_argument(
+            "[quantized_matmul] K must be multiples of group_size.");
+      }
+      f.template operator()<false>();
     } else {
-      dispatch.template operator()<64, k_major.value>();
+      dispatch_bool(has_k_residue, [&](auto has_k_residue) {
+        f.template operator()<has_k_residue.value>();
+      });
     }
+  };
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  int k = x.shape(-1);
+  bool has_k_residue = k % group_size != 0;
+  bool sm80 = encoder.device().compute_capability_major() >= 8;
+  dispatch_bool(transpose, [&](auto k_major) {
+    dispatch_k(k_major, has_k_residue, [&]<bool HasKResidue>() {
+      dispatch_bool(sm80, [&](auto sm80) {
+        constexpr bool KMajor = k_major.value;
+        constexpr bool SM80 = sm80.value;
+        if (m <= 16) {
+          dispatch.template operator()<16, KMajor, HasKResidue, SM80>();
+        } else if (m <= 32) {
+          dispatch.template operator()<32, KMajor, HasKResidue, SM80>();
+        } else {
+          dispatch.template operator()<64, KMajor, HasKResidue, SM80>();
+        }
+      });
+    });
   });
 }
 

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(16, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(16, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(32, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(32, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(64, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(64, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(16)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(32)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(64)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _128>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _16>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _256>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _32>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _64>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_naive.cu

@@ -316,7 +316,7 @@ inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size
   }
 }
 
-template <int TileM = 16, bool KMajor = true, bool SM80 = true, bool HasKResidue = false,
+template <int TileM = 16, bool KMajor = true, bool HasKResidue = false, bool SM80 = true,
           typename Element, typename Quant, typename Scale>
 void qmm_naive(
     const Element* A,
@@ -396,21 +396,6 @@ void qmm_naive(
 
 namespace mlx::core {
 
-template <bool KMajor, typename F>
-inline void dispatch_k(bool has_k_residue, const char* tag, F&& f) {
-  if constexpr (KMajor) {
-    if (has_k_residue) {
-      throw std::invalid_argument(
-          fmt::format("{} K must be multiples of group_size.", tag));
-    }
-    f.template operator()<false>();
-  } else {
-    dispatch_bool(has_k_residue, [&](auto has_k_residue) {
-      f.template operator()<has_k_residue.value>();
-    });
-  }
-}
-
 template <typename F>
 inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
   if (dtype == float32) {
@@ -474,8 +459,8 @@ inline void dispatch_quant_types(
   }
 }
 
-template <int TileM, bool KMajor>
-void qmm_impl_naive(
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -494,71 +479,65 @@ void qmm_impl_naive(
   int l = out.size() / (m * n);
   bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size());
 
-  bool is_sm80 = encoder.device().compute_capability_major() >= 8;
-  dispatch_bool(is_sm80, [&](auto sm80) {
-    dispatch_k<KMajor>(k % group_size != 0, tag, [&]<bool has_k_residue>() {
-      dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
-        dispatch_quant_types<Element>(
-            bits,
-            group_size,
-            mode,
-            tag,
-            [&]<typename Quant, typename Scale, int group_size>() {
-              encoder.set_input_array(x);
-              encoder.set_input_array(w);
-              encoder.set_input_array(scales);
-              if (biases) {
-                encoder.set_input_array(*biases);
-              }
-              if (lhs_indices) {
-                encoder.set_input_array(*lhs_indices);
-              }
-              if (rhs_indices) {
-                encoder.set_input_array(*rhs_indices);
-              }
-              encoder.set_output_array(out);
-              cutlass_gemm::qmm_naive<TileM, KMajor, sm80.value, has_k_residue>(
-                  gpu_ptr<Element>(x),
-                  gpu_ptr<Quant>(w),
-                  gpu_ptr<Scale>(scales),
-                  biases ? gpu_ptr<Element>(*biases) : nullptr,
-                  lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
-                  rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
-                  gpu_ptr<Element>(out),
-                  m,
-                  n,
-                  k,
-                  l,
-                  broadcast_b,
-                  cute::Int<group_size>{},
-                  [&](auto* kernel,
-                      dim3 num_blocks,
-                      dim3 block_dims,
-                      uint32_t smem_bytes,
-                      void** args) {
-                    encoder.add_kernel_node_raw(
-                        kernel, num_blocks, block_dims, {}, smem_bytes, args);
-                  });
-            });
-      });
-    });
+  dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
+    dispatch_quant_types<Element>(
+        bits,
+        group_size,
+        mode,
+        tag,
+        [&]<typename Quant, typename Scale, int group_size>() {
+          encoder.set_input_array(x);
+          encoder.set_input_array(w);
+          encoder.set_input_array(scales);
+          if (biases) {
+            encoder.set_input_array(*biases);
+          }
+          if (lhs_indices) {
+            encoder.set_input_array(*lhs_indices);
+          }
+          if (rhs_indices) {
+            encoder.set_input_array(*rhs_indices);
+          }
+          encoder.set_output_array(out);
+          cutlass_gemm::qmm_naive<TileM, KMajor, HasKResidue, SM80>(
+              gpu_ptr<Element>(x),
+              gpu_ptr<Quant>(w),
+              gpu_ptr<Scale>(scales),
+              biases ? gpu_ptr<Element>(*biases) : nullptr,
+              lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
+              rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
+              gpu_ptr<Element>(out),
+              m,
+              n,
+              k,
+              l,
+              broadcast_b,
+              cute::Int<group_size>{},
+              [&](auto* kernel,
+                  dim3 num_blocks,
+                  dim3 block_dims,
+                  uint32_t smem_bytes,
+                  void** args) {
+                encoder.add_kernel_node_raw(
+                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
+              });
+        });
   });
 }
 
-} // namespace mlx::core
+// clang-format off
+template void qmm_naive_impl<@TileM@, @KMajor@, @HasKResidue@, @SM80@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
 
-#define QMM_NAIVE_GPU(TileM, KMajor)           \
-  namespace mlx::core {                        \
-  template void qmm_impl_naive<TileM, KMajor>( \
-      const array& x,                          \
-      const array& w,                          \
-      const array& scales,                     \
-      const std::optional<array>& biases,      \
-      const std::optional<array>& lhs_indices, \
-      const std::optional<array>& rhs_indices, \
-      array& out,                              \
-      int bits,                                \
-      int group_size,                          \
-      QuantizationMode mode,                   \
-      cu::CommandEncoder& encoder);            \
-  }
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm80.cu

@@ -434,7 +434,7 @@ inline void dispatch_quant_types(
 }
 
 template <int TileM>
-void qmm_impl_sm80(
+void qmm_sm80_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -499,20 +499,19 @@ void qmm_impl_sm80(
   });
 }
 
-} // namespace mlx::core
+// clang-format off
+template void qmm_sm80_impl<@TileM@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
 
-#define QMM_SM80_GPU(TileM)                    \
-  namespace mlx::core {                        \
-  template void qmm_impl_sm80<TileM>(          \
-      const array& x,                          \
-      const array& w,                          \
-      const array& scales,                     \
-      const std::optional<array>& biases,      \
-      const std::optional<array>& lhs_indices, \
-      const std::optional<array>& rhs_indices, \
-      array& out,                              \
-      int bits,                                \
-      int group_size,                          \
-      QuantizationMode mode,                   \
-      cu::CommandEncoder& encoder);            \
-  }
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm90.cu

@@ -20,8 +20,7 @@ namespace cutlass_gemm {
 using namespace cute;
 
 template <
-    typename TileShapeMN = Shape<_128, _16>,
-    typename ClusterShape = Shape<_1, _1, _1>,
+    int TileN = 16,
     typename Element,
     typename Quant,
     typename GroupSize,
@@ -47,7 +46,8 @@ void qmm_sm90(
 
   using Arch = cutlass::arch::Sm90;
   using Accumulator = float;
-  using TileShape = decltype(append(TileShapeMN{}, Int<kTileShapeK>{}));
+  using TileShape = Shape<_128, Int<TileN>, Int<kTileShapeK>>;
+  using ClusterShape = Shape<Int<(TileN <= 32) ? 1 : 2>, _1, _1>;
 
   using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       Arch,
@@ -177,8 +177,8 @@ inline void dispatch_groups(int group_size, const char* tag, F&& f) {
   }
 }
 
-template <typename TileShapeMN, typename ClusterShape>
-void qmm_impl_sm90(
+template <int TileN>
+void qmm_sm90_impl(
     const array& x,
     const array& w,
     const array& scales_,
@@ -207,7 +207,7 @@ void qmm_impl_sm90(
         encoder.set_input_array(scales);
         encoder.set_input_array(biases);
         encoder.set_output_array(out);
-        cutlass_gemm::qmm_sm90(
+        cutlass_gemm::qmm_sm90<TileN>(
             gpu_ptr<Element>(x),
             gpu_ptr<Quant>(w),
             gpu_ptr<Element>(scales),
@@ -238,24 +238,19 @@ void qmm_impl_sm90(
   });
 }
 
+// clang-format off
+template void qmm_sm90_impl<@TileN@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s);
+// clang-format on
+
 } // namespace mlx::core
 
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)           \
-  namespace mlx::core {                                   \
-  template void qmm_impl_sm90<TileShapeMN, ClusterShape>( \
-      const array& x,                                     \
-      const array& w,                                     \
-      const array& scales,                                \
-      const array& biases,                                \
-      array& out,                                         \
-      int bits,                                           \
-      int group_size,                                     \
-      cu::CommandEncoder& encoder,                        \
-      Stream s);                                          \
-  }
-
-#else
-
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)
-
 #endif // defined(MLX_CUDA_SM90A_ENABLED)

mlx/distributed/jaccl/lib/examples/CMakeLists.txt

@@ -35,6 +35,8 @@ endfunction()
 # Examples
 build_example(minimal_env.cpp)
 build_example(minimal_cfg.cpp)
+build_example(monte_carlo_pi.cpp)
+build_example(file_broadcast.cpp)
 
 # Benchmarks
 build_example(allreduce_bench.cpp)

mlx/distributed/jaccl/lib/examples/file_broadcast.cpp

@@ -0,0 +1,360 @@
+// Copyright © 2025 Apple Inc.
+//
+// File Broadcast with JACCL
+//
+// This example demonstrates distributed file transfer using JACCL's all_sum
+// operation to broadcast a file from any rank to all other machines.
+//
+// The algorithm:
+// 1. The sender rank reads the file into memory
+// 2. All other ranks allocate zero-filled buffers of the same size
+// 3. Use all_sum to broadcast: sender has data, others have zeros
+// 4. After all_sum, all ranks have the file data
+// 5. All ranks write the file to disk
+//
+// For large files, the transfer is chunked to manage memory efficiently.
+//
+// Usage:
+//   Set environment variables (see README.md), then run:
+//
+//     ./jaccl_file_broadcast -f <file> [-s <sender_rank>] [-o <output_dir>]
+//
+//   Or with mlx.launch:
+//
+//     mlx.launch --hostfile hosts.json ./jaccl_file_broadcast -f myfile.bin
+//
+// Example output (4 ranks, sender rank 2):
+//   Rank 0 of 4: Received 10485760 bytes from rank 2 (982.5 MB/s)
+//   Rank 1 of 4: Received 10485760 bytes from rank 2 (985.2 MB/s)
+//   Rank 2 of 4: Sent 10485760 bytes (980.1 MB/s)
+//   Rank 3 of 4: Received 10485760 bytes from rank 2 (978.9 MB/s)
+
+#include <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <sys/stat.h>
+#include <atomic>
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <string>
+#include <thread>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr
+      << "Usage: " << prog << " [options]\n"
+      << "  -f <file>       File to broadcast (required)\n"
+      << "  -s <rank>       Sender rank (default: 0)\n"
+      << "  -o <dir>        Output directory (default: current dir)\n"
+      << "  -c <bytes>      Chunk size in bytes (default: 67108864 = 64MB)\n"
+      << "  -v              Verbose output\n"
+      << "  -h              Show this help\n";
+}
+
+static bool file_exists(const std::string& path) {
+  struct stat buffer;
+  return (stat(path.c_str(), &buffer) == 0);
+}
+
+static std::int64_t file_size(const std::string& path) {
+  struct stat buffer;
+  if (stat(path.c_str(), &buffer) != 0) {
+    return -1;
+  }
+  return static_cast<std::int64_t>(buffer.st_size);
+}
+
+static bool create_directory(const std::string& path) {
+  if (path.empty() || path == ".") {
+    return true;
+  }
+  return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST;
+}
+
+static std::string basename(const std::string& path) {
+  size_t pos = path.find_last_of("/\\");
+  return (pos == std::string::npos) ? path : path.substr(pos + 1);
+}
+
+struct BroadcastStats {
+  std::int64_t total_bytes;
+  std::int64_t chunks_sent;
+  std::int64_t chunks_received;
+  double total_time_ms;
+  int sender_rank;
+};
+
+int main(int argc, char** argv) {
+  std::string input_file;
+  std::string output_dir = ".";
+  int sender_rank = 0;
+  std::int64_t chunk_size = 67108864;
+  bool verbose = false;
+
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-h" || arg == "--help") {
+      usage(argv[0]);
+      return 0;
+    } else if (arg == "-f" && i + 1 < argc) {
+      input_file = argv[++i];
+    } else if (arg == "-s" && i + 1 < argc) {
+      sender_rank = std::atoi(argv[++i]);
+    } else if (arg == "-o" && i + 1 < argc) {
+      output_dir = argv[++i];
+    } else if (arg == "-c" && i + 1 < argc) {
+      chunk_size = std::atoll(argv[++i]);
+    } else if (arg == "-v" || arg == "--verbose") {
+      verbose = true;
+    } else {
+      std::cerr << "Unknown option: " << arg << "\n";
+      usage(argv[0]);
+      return 1;
+    }
+  }
+
+  if (input_file.empty()) {
+    std::cerr << "Error: Input file is required (-f <file>)\n";
+    usage(argv[0]);
+    return 1;
+  }
+
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int nranks = group->size();
+
+  if (sender_rank < 0 || sender_rank >= nranks) {
+    std::cerr << "Error: Sender rank " << sender_rank << " is out of range [0, "
+              << nranks << ")\n";
+    return 1;
+  }
+
+  std::int64_t total_file_size = 0;
+  if (rank == sender_rank) {
+    if (!file_exists(input_file)) {
+      std::cerr << "Error: File not found: " << input_file << "\n";
+      return 1;
+    }
+    total_file_size = file_size(input_file);
+    if (total_file_size < 0) {
+      std::cerr << "Error: Cannot read file size: " << input_file << "\n";
+      return 1;
+    }
+  }
+
+  group->all_sum(
+      &total_file_size, &total_file_size, sizeof(int64_t), jaccl::Int64);
+
+  if (!create_directory(output_dir)) {
+    std::cerr << "Error: Cannot create output directory: " << output_dir
+              << "\n";
+    return 1;
+  }
+
+  std::string output_file = output_dir == "."
+      ? basename(input_file)
+      : output_dir + "/" + basename(input_file);
+
+  if (verbose) {
+    std::printf(
+        "Rank %d of %d: Broadcasting '%s' (%ld bytes) from rank %d\n",
+        rank,
+        nranks,
+        input_file.c_str(),
+        static_cast<long>(total_file_size),
+        sender_rank);
+  }
+
+  auto t_start = std::chrono::high_resolution_clock::now();
+
+  std::int64_t num_chunks = (total_file_size + chunk_size - 1) / chunk_size;
+  if (num_chunks == 0) {
+    num_chunks = 1;
+  }
+
+  const int num_buffers = 4;
+  std::vector<std::vector<std::uint8_t>> buffers(
+      num_buffers, std::vector<std::uint8_t>(chunk_size, 0));
+
+  std::ifstream infile;
+  std::ofstream outfile;
+
+  if (rank == sender_rank) {
+    infile.open(input_file, std::ios::binary);
+    if (!infile.good()) {
+      std::cerr << "Error: Cannot open input file: " << input_file << "\n";
+      return 1;
+    }
+  }
+
+  outfile.open(output_file, std::ios::binary);
+  if (!outfile.good()) {
+    std::cerr << "Error: Cannot open output file: " << output_file << "\n";
+    return 1;
+  }
+
+  std::atomic<std::int64_t> next_read_chunk{0};
+  std::atomic<std::int64_t> next_comm_chunk{0};
+  std::atomic<std::int64_t> next_write_chunk{0};
+  std::atomic<bool> read_done{false};
+  std::atomic<bool> comm_done{false};
+
+  std::vector<std::atomic<bool>> buffer_ready(num_buffers);
+  std::vector<std::atomic<bool>> buffer_written(num_buffers);
+  for (int i = 0; i < num_buffers; i++) {
+    buffer_ready[i] = false;
+    buffer_written[i] = false;
+  }
+
+  std::vector<std::int64_t> chunk_sizes(num_chunks);
+  for (std::int64_t i = 0; i < num_chunks; i++) {
+    chunk_sizes[i] = std::min(chunk_size, total_file_size - i * chunk_size);
+  }
+
+  std::thread reader_thread;
+  if (rank == sender_rank) {
+    reader_thread = std::thread([&]() {
+      while (true) {
+        std::int64_t chunk_idx = next_read_chunk.fetch_add(1);
+        if (chunk_idx >= num_chunks) {
+          break;
+        }
+        std::int64_t offset = chunk_idx * chunk_size;
+        std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+        int buffer_idx = chunk_idx % num_buffers;
+
+        infile.seekg(offset, std::ios::beg);
+        infile.read(
+            reinterpret_cast<char*>(buffers[buffer_idx].data()),
+            this_chunk_size);
+
+        std::fill(
+            buffers[buffer_idx].begin() + this_chunk_size,
+            buffers[buffer_idx].end(),
+            0);
+
+        buffer_ready[buffer_idx] = true;
+      }
+      read_done = true;
+    });
+  } else {
+    read_done = true;
+  }
+
+  std::thread writer_thread([&]() {
+    while (true) {
+      std::int64_t chunk_idx = next_write_chunk.load();
+      if (chunk_idx >= num_chunks && comm_done) {
+        break;
+      }
+      if (chunk_idx >= num_chunks) {
+        std::this_thread::yield();
+        continue;
+      }
+
+      int buffer_idx = chunk_idx % num_buffers;
+      if (!buffer_written[buffer_idx]) {
+        std::this_thread::yield();
+        continue;
+      }
+
+      std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+      outfile.write(
+          reinterpret_cast<const char*>(buffers[buffer_idx].data()),
+          this_chunk_size);
+
+      buffer_written[buffer_idx] = false;
+      next_write_chunk.fetch_add(1);
+    }
+  });
+
+  for (std::int64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) {
+    std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+    int buffer_idx = chunk_idx % num_buffers;
+
+    if (rank == sender_rank) {
+      while (!buffer_ready[buffer_idx] && !read_done) {
+        std::this_thread::yield();
+      }
+    }
+
+    std::fill(
+        buffers[buffer_idx].begin() + this_chunk_size,
+        buffers[buffer_idx].end(),
+        0);
+
+    group->all_sum(
+        buffers[buffer_idx].data(),
+        buffers[buffer_idx].data(),
+        this_chunk_size,
+        jaccl::UInt8);
+
+    buffer_written[buffer_idx] = true;
+    next_comm_chunk.fetch_add(1);
+
+    if (verbose) {
+      double progress = 100.0 * (chunk_idx + 1) / num_chunks;
+      std::printf(
+          "Rank %d: Progress %.1f%% (%ld/%ld chunks)\n",
+          rank,
+          progress,
+          static_cast<long>(chunk_idx + 1),
+          static_cast<long>(num_chunks));
+    }
+  }
+
+  comm_done = true;
+
+  if (reader_thread.joinable()) {
+    reader_thread.join();
+  }
+  writer_thread.join();
+
+  infile.close();
+  outfile.close();
+
+  auto t_end = std::chrono::high_resolution_clock::now();
+  double elapsed_ms =
+      std::chrono::duration<double, std::milli>(t_end - t_start).count();
+  double elapsed_sec = elapsed_ms / 1000.0;
+  double bandwidth_mbps = (total_file_size / (1024.0 * 1024.0)) / elapsed_sec;
+
+  if (rank == sender_rank) {
+    std::printf(
+        "Rank %d of %d: Sent %ld bytes from '%s' (%.1f MB/s)\n",
+        rank,
+        nranks,
+        static_cast<long>(total_file_size),
+        input_file.c_str(),
+        bandwidth_mbps);
+  } else {
+    std::printf(
+        "Rank %d of %d: Received %ld bytes from rank %d to '%s' (%.1f MB/s)\n",
+        rank,
+        nranks,
+        static_cast<long>(total_file_size),
+        sender_rank,
+        output_file.c_str(),
+        bandwidth_mbps);
+  }
+
+  if (verbose) {
+    std::printf(
+        "Rank %d: Total time: %.2f ms, Chunks: %ld, Chunk size: %ld bytes\n",
+        rank,
+        elapsed_ms,
+        static_cast<long>(num_chunks),
+        static_cast<long>(chunk_size));
+  }
+
+  return 0;
+}

mlx/distributed/jaccl/lib/examples/monte_carlo_pi.cpp

@@ -0,0 +1,152 @@
+// Copyright © 2025 Apple Inc.
+//
+// Monte Carlo Pi Estimation with JACCL
+//
+// This example demonstrates distributed Monte Carlo simulation using JACCL
+// to estimate the value of π. Each rank generates random points independently
+// and uses all-reduce to aggregate the results across all machines.
+//
+// The algorithm:
+// 1. Each rank generates N random points in the unit square [0,1] x [0,1]
+// 2. Count how many fall inside the quarter circle (x² + y² ≤ 1)
+// 3. Use all_sum to aggregate hits and total points across all ranks
+// 4. π ≈ 4 × (hits / total)
+//
+// Usage:
+//   Set environment variables (see README.md), then run:
+//
+//     ./jaccl_monte_carlo_pi [-n <points_per_rank>]
+//
+//   Or with mlx.launch:
+//
+//     mlx.launch --hostfile hosts.json ./jaccl_monte_carlo_pi -n 10000000
+//
+// Example output (4 ranks, 10M points each):
+//   Rank 2 of 4
+//   Local: 7854321 hits out of 10000000 points
+//   Global: 31416789 hits out of 40000000 points
+//   Estimated π = 3.141679 (error: 0.000086)
+
+#include <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <random>
+#include <string>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr << "Usage: " << prog << " [options]\n"
+            << "  -n <points>  Points per rank (default: 1000000)\n"
+            << "  -s <seed>    Random seed base (default: 42)\n"
+            << "  -h           Show this help\n";
+}
+
+struct MonteCarloResult {
+  int64_t hits;
+  int64_t total;
+};
+
+MonteCarloResult estimate_pi_local(int64_t num_points, unsigned int seed) {
+  std::mt19937_64 rng(seed);
+  std::uniform_real_distribution<double> dist(0.0, 1.0);
+
+  int64_t hits = 0;
+  for (int64_t i = 0; i < num_points; i++) {
+    double x = dist(rng);
+    double y = dist(rng);
+    if (x * x + y * y <= 1.0) {
+      hits++;
+    }
+  }
+
+  return {hits, num_points};
+}
+
+int main(int argc, char** argv) {
+  int64_t points_per_rank = 1000000;
+  unsigned int seed_base = 42;
+
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-h" || arg == "--help") {
+      usage(argv[0]);
+      return 0;
+    } else if (arg == "-n" && i + 1 < argc) {
+      points_per_rank = std::atoll(argv[++i]);
+    } else if (arg == "-s" && i + 1 < argc) {
+      seed_base = static_cast<unsigned int>(std::atoi(argv[++i]));
+    } else {
+      std::cerr << "Unknown option: " << arg << "\n";
+      usage(argv[0]);
+      return 1;
+    }
+  }
+
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int nranks = group->size();
+
+  std::printf("Rank %d of %d\n", rank, nranks);
+  std::printf(
+      "Generating %ld random points (seed: %u)...\n",
+      static_cast<long>(points_per_rank),
+      seed_base + static_cast<unsigned int>(rank));
+
+  auto t0 = std::chrono::high_resolution_clock::now();
+
+  MonteCarloResult local = estimate_pi_local(
+      points_per_rank, seed_base + static_cast<unsigned int>(rank));
+
+  auto t1 = std::chrono::high_resolution_clock::now();
+  double local_time =
+      std::chrono::duration<double, std::milli>(t1 - t0).count();
+
+  std::printf(
+      "Rank %d: %ld hits out of %ld points (%.2f ms)\n",
+      rank,
+      static_cast<long>(local.hits),
+      static_cast<long>(local.total),
+      local_time);
+
+  MonteCarloResult global = {0, 0};
+
+  group->all_sum(&local.hits, &global.hits, sizeof(int64_t), jaccl::Int64);
+  group->all_sum(&local.total, &global.total, sizeof(int64_t), jaccl::Int64);
+
+  if (rank == 0) {
+    double pi_estimate = 4.0 * static_cast<double>(global.hits) /
+        static_cast<double>(global.total);
+    double error = std::abs(pi_estimate - M_PI);
+
+    std::printf("\n=== Results ===\n");
+    std::printf(
+        "Global: %ld hits out of %ld points\n",
+        static_cast<long>(global.hits),
+        static_cast<long>(global.total));
+    std::printf("Estimated π = %.10f\n", pi_estimate);
+    std::printf("True π      = %.10f\n", M_PI);
+    std::printf("Error       = %.10f (%.6f%%)\n", error, 100.0 * error / M_PI);
+
+    double total_time =
+        std::chrono::duration<double, std::milli>(t1 - t0).count();
+    std::printf("\nPerformance:\n");
+    std::printf("Total points: %ld\n", static_cast<long>(global.total));
+    std::printf("Time: %.2f ms\n", total_time);
+    std::printf(
+        "Points/sec: %.0f\n",
+        static_cast<double>(global.total) / (total_time / 1000.0));
+  }
+
+  return 0;
+}