Tentative JACCL examples

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

@@ -20,7 +20,7 @@ runs:
       run: |
         pip install auditwheel "build<=1.4.2" patchelf setuptools
         python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
+        MLX_DISABLE_SM90A_KERNELS=1 MLX_BUILD_STAGE=2 python -m build -w
 
         auditwheel repair dist/mlx_cuda*.whl \
           --plat manylinux_2_35_${{ inputs.arch }} \

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

@@ -14,9 +14,6 @@ inputs:
     description: 'Whether to enable ccache'
     required: false
     default: 'true'
-  ccache-key:
-    required: false
-    default: 'ccache'
 
 runs:
   using: "composite"
@@ -36,7 +33,7 @@ runs:
       if: ${{ inputs.use-ccache == 'true' }}
       uses: hendrikmuhs/ccache-action@v1.2
       with:
-        key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
+        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
         max-size: 1GB
         # ccache-action bug: running "apt-get update" fails on large arm runner.
         update-package-index: false

.github/workflows/build-wheels.yml

@@ -1,94 +0,0 @@
-name: Build macOS arm64 wheels
-
-on:
-  push:
-    branches:
-      - main
-      - 'metal-*'
-      - 'q-*'
-      - attn-mask-fix
-      - fix-rope
-  workflow_dispatch:
-    inputs:
-      branch_to_build:
-        description: 'Branch to build (optional, defaults to current ref)'
-        required: false
-        default: ''
-
-concurrency:
-  group: build-${{ github.ref }}-${{ github.event.inputs.branch_to_build }}
-  cancel-in-progress: true
-
-jobs:
-  build:
-    name: Build wheel (Python ${{ matrix.python }})
-    runs-on: macos-14
-    timeout-minutes: 60
-    strategy:
-      fail-fast: false
-      matrix:
-        python: ['3.11', '3.12']
-    env:
-      CMAKE_BUILD_PARALLEL_LEVEL: '4'
-    steps:
-      - name: Determine target branch
-        id: branch
-        run: |
-          NAME="${{ github.event.inputs.branch_to_build }}"
-          if [ -z "$NAME" ]; then
-            NAME="${{ github.ref_name }}"
-          fi
-          # Sanitize for artifact naming (replace / with -)
-          SAFE_NAME=$(echo "$NAME" | tr '/' '-')
-          echo "name=$NAME" >> $GITHUB_OUTPUT
-          echo "safe_name=$SAFE_NAME" >> $GITHUB_OUTPUT
-          echo "Target branch: $NAME (safe: $SAFE_NAME)"
-
-      - name: Checkout
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ steps.branch.outputs.name }}
-          submodules: recursive
-
-      - name: Set up Python ${{ matrix.python }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python }}
-
-      - name: Cache pip
-        uses: actions/cache@v4
-        with:
-          path: |
-            ~/.cache/pip
-            ~/Library/Caches/pip
-          key: pip-${{ runner.os }}-py${{ matrix.python }}-${{ hashFiles('CMakeLists.txt', 'setup.py', 'pyproject.toml') }}
-          restore-keys: |
-            pip-${{ runner.os }}-py${{ matrix.python }}-
-
-      - name: Install build dependencies
-        run: |
-          python -m pip install -U pip wheel build setuptools cmake nanobind
-
-      - name: Build wheel
-        run: |
-          mkdir -p ./wheels
-          pip wheel --no-deps . -w ./wheels
-
-      - name: List built wheels
-        run: ls -lh ./wheels
-
-      - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: mlx-${{ steps.branch.outputs.safe_name }}-py${{ matrix.python }}-wheels
-          path: ./wheels/*.whl
-          retention-days: 30
-          if-no-files-found: error
-
-      - name: Create GitHub Release (on tag)
-        if: startsWith(github.ref, 'refs/tags/v')
-        uses: softprops/action-gh-release@v2
-        with:
-          files: ./wheels/*.whl
-          fail_on_unmatched_files: false
-          generate_release_notes: true

.github/workflows/nightly.yml

@@ -85,24 +85,20 @@ jobs:
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        arch: ['x86_64', 'aarch64']
-        toolkit: ['cuda-12.9', 'cuda-13.0']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
+    runs-on: ubuntu-22-large
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
-          toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
+          toolkit: 'cuda-12.9'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          arch: ${{ matrix.arch }}
+          toolkit: 'cuda-12.9'
+          arch: 'x86_64'
       - name: Upload artifacts
         uses: actions/upload-artifact@v7
         with:
-          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          name: mlx-cuda
           path: wheelhouse/mlx_cuda_*.whl
           retention-days: 7

.github/workflows/release.yml

@@ -146,7 +146,7 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
+          use-ccache: false
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:

docs/src/python/linalg.rst

@@ -14,7 +14,6 @@ Linear Algebra
     cholesky
     cholesky_inv
     cross
-    det
     qr
     svd
     eigvals
@@ -24,6 +23,5 @@ Linear Algebra
     lu
     lu_factor
     pinv
-    slogdet
     solve
     solve_triangular

mlx/backend/common/common.cpp

@@ -19,28 +19,27 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
         "AsStrided must be used with row contiguous arrays only.");
   }
 
-  auto [no_bsx_size, row_contiguous, col_contiguous] =
-      check_contiguity(shape_, strides_);
-
-  int64_t l = 0, h = 0;
-  bool has_negative_stride = false;
-  for (int i = 0; i < strides_.size(); i++) {
-    auto delta = strides_[i] * (shape_[i] - 1);
-    if (strides_[i] >= 0) {
-      h += delta;
-    } else {
-      l += delta;
-      has_negative_stride |= shape_[i] > 1;
-    }
+  // Compute the flags given the shape and strides
+  bool row_contiguous = true, col_contiguous = true;
+  size_t r = 1, c = 1;
+  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
   }
-  size_t data_size = out.size() == 0 ? 0 : (h - l) + 1;
-
   auto flags = in.flags();
-  flags.contiguous =
-      out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size);
+  // TODO: Compute the contiguous flag in a better way cause now we are
+  //       unnecessarily strict.
+  flags.contiguous = row_contiguous || col_contiguous;
   flags.row_contiguous = row_contiguous;
   flags.col_contiguous = col_contiguous;
 
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
+
   return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 

mlx/backend/cpu/compiled.cpp

@@ -10,6 +10,7 @@
 #include <fmt/format.h>
 
 #include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cpu/compiled_preamble.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/jit_compiler.h"
 #include "mlx/device.h"
@@ -315,9 +316,7 @@ void Compiled::eval_cpu(
   // Get the function
   auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
     std::ostringstream kernel;
-    kernel << std::get<2>(JitCompiler::get_preamble()) << std::endl;
-    kernel << "using namespace mlx::core;" << std::endl;
-    kernel << "using namespace mlx::core::detail;" << std::endl;
+    kernel << get_kernel_preamble() << std::endl;
     kernel << "extern \"C\"  {" << std::endl;
     build_kernel(
         kernel,

mlx/backend/cpu/compiled_preamble.h

@@ -9,4 +9,4 @@
 #include "mlx/backend/cpu/binary_ops.h"
 // clang-format on
 
-const char* get_prebuilt_preamble();
+const char* get_kernel_preamble();

mlx/backend/cpu/jit_compiler.cpp

@@ -1,8 +1,6 @@
 // Copyright © 2024 Apple Inc.
 
 #include "mlx/backend/cpu/jit_compiler.h"
-#include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/compiled_preamble.h"
 
 #include <algorithm>
 #include <sstream>
@@ -88,61 +86,30 @@ const VisualStudioInfo& GetVisualStudioInfo() {
 
 #endif // _MSC_VER
 
-const std::tuple<bool, std::string, std::string>& JitCompiler::get_preamble() {
-  static auto preamble = []() -> std::tuple<bool, std::string, std::string> {
-    // Check whether the headers are shipped with the binary, if so use the
-    // preamble from the headers, otherwise use the prebuilt one embeded in
-    // binary, which may not work with all compilers.
-    auto root_dir = current_binary_dir();
-#if !defined(_WIN32)
-    root_dir = root_dir.parent_path();
-#endif
-    auto include_dir = root_dir / "include";
-    if (std::filesystem::exists(include_dir / "mlx")) {
-      return std::make_tuple(
-          true,
-          include_dir.string(),
-          "#include \"mlx/backend/cpu/compiled_preamble.h\"\n");
-    } else {
-      return std::make_tuple(false, "", get_prebuilt_preamble());
-    }
-  }();
-  return preamble;
-}
-
 std::string JitCompiler::build_command(
     const std::filesystem::path& dir,
     const std::string& source_file_name,
     const std::string& shared_lib_name) {
-  auto& [use_include, include_dir, preamble] = get_preamble();
 #ifdef _MSC_VER
-  std::string extra_flags;
-  if (use_include) {
-    extra_flags += fmt::format("/I \"{}\"", include_dir);
-  }
   const VisualStudioInfo& info = GetVisualStudioInfo();
+  std::string libpaths;
   for (const std::string& lib : info.libpaths) {
-    extra_flags += fmt::format(" /libpath:\"{}\"", lib);
+    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
   }
   return fmt::format(
       "\""
-      "cd /D \"{}\" && "
-      "\"{}\" /LD /EHsc /MD /Ox /nologo /std:c++17 {} \"{}\" "
-      "/link /out:\"{}\" 2>&1"
+      "cd /D \"{0}\" && "
+      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
+      "/link /out:\"{3}\" {4} 2>&1"
       "\"",
       dir.string(),
       info.cl_exe,
-      extra_flags,
       source_file_name,
-      shared_lib_name);
+      shared_lib_name,
+      libpaths);
 #else
-  std::string extra_flags;
-  if (use_include) {
-    extra_flags = fmt::format("-I \"{}\"", include_dir);
-  }
   return fmt::format(
-      "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1",
-      extra_flags,
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1",
       (dir / source_file_name).string(),
       (dir / shared_lib_name).string());
 #endif

mlx/backend/cpu/jit_compiler.h

@@ -7,9 +7,6 @@ namespace mlx::core {
 
 class JitCompiler {
  public:
-  // Return the includes that should be prepended to the source code.
-  static const std::tuple<bool, std::string, std::string>& get_preamble();
-
   // Build a shell command that compiles a source code file to a shared library.
   static std::string build_command(
       const std::filesystem::path& dir,

mlx/backend/cpu/luf.cpp

@@ -67,10 +67,11 @@ void luf_impl(
               /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
               /* info */ &info);
 
-          if (info < 0) {
+          if (info != 0) {
             std::stringstream ss;
             ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-               << " because argument had an illegal value";
+               << ((info > 0) ? " because matrix is singular"
+                              : " because argument had an illegal value");
             throw std::runtime_error(ss.str());
           }
 

mlx/backend/cpu/make_compiled_preamble.ps1

@@ -15,6 +15,13 @@ $CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
 # Concatenate to string.
 $CONTENT = $CONTENT -join "`n"
 
+# Append extra content.
+$CONTENT = @"
+$($CONTENT)
+using namespace mlx::core;
+using namespace mlx::core::detail;
+"@
+
 # Convert each char to ASCII code.
 # Unlike the unix script that outputs string literal directly, the output from
 # MSVC is way too large to be embedded as string and compilation will fail, so
@@ -22,7 +29,7 @@ $CONTENT = $CONTENT -join "`n"
 $CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
 
 $OUTPUT = @"
-const char* get_prebuilt_preamble() {
+const char* get_kernel_preamble() {
   static char preamble[] = { $CHARCODES };
   return preamble;
 }

mlx/backend/cpu/make_compiled_preamble.sh

@@ -30,10 +30,12 @@ fi
 CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
-const char* get_prebuilt_preamble() {
+const char* get_kernel_preamble() {
 return R"preamble(
 $INCLUDES
 $CONTENT
+using namespace mlx::core;
+using namespace mlx::core::detail;
 )preamble";
 }
 EOF

mlx/backend/cuda/CMakeLists.txt

@@ -168,8 +168,9 @@ set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
 # Skip Hopper-only kernels when not building for sm90a.
-if(("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
-                                              MLX_CUDA_ARCHITECTURES))
+if(NOT DEFINED ENV{MLX_DISABLE_SM90A_KERNELS}
+   AND (("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
+                                                   MLX_CUDA_ARCHITECTURES)))
   target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED)
 endif()
 

mlx/backend/cuda/gemms/gather_gemm.cu

@@ -43,12 +43,6 @@ class GatherGemm {
   using ElementD = typename CollectiveEpilogue::ElementD;
   using StrideD = typename CollectiveEpilogue::StrideD;
 
-  static_assert(
-      cute::is_same_v<
-          ElementAccumulator,
-          typename CollectiveEpilogue::ElementAccumulator>,
-      "Mainloop and epilogue do not agree on accumulator value type.");
-
   static constexpr int SharedStorageSize = static_cast<int>(cute::max(
       sizeof(typename CollectiveMainloop::SharedStorage),
       sizeof(typename CollectiveEpilogue::SharedStorage)));
@@ -104,9 +98,7 @@ class GatherGemm {
 
   CUTLASS_DEVICE void operator()(const Params& params, char* smem_buf) {
     int thread_idx = int(threadIdx.x);
-    int m_coord = int(blockIdx.x);
-    int n_coord = int(blockIdx.y);
-    int l_coord = int(blockIdx.z);
+    auto [m_coord, n_coord, l_coord] = uint3(blockIdx);
 
     auto shape_MNKL = append<4>(params.problem_shape, Int<1>{});
     auto cta_tile = TileShape{};
@@ -228,7 +220,7 @@ void gather_mm(
       using TileShape = Shape<_128, _128, _8>;
       using DispatchPolicy = cutlass::gemm::MainloopSm70TwoStage;
       using TiledMma = TiledMMA<
-          MMA_Atom<UniversalFMA<Accumulator, Element, Element, Accumulator>>,
+          MMA_Atom<UniversalFMA<Accumulator, Element, Element, Element>>,
           Layout<Shape<_16, _16, _1>>>;
 
       using CopyTraitsA = SimtCopyTraits<Element, k_major_a.value>;
@@ -304,6 +296,9 @@ void cutlass_gather_mm(
   int n = out.shape(-1);
   int k = a.shape(-1);
   int l = out.size() / (m * n);
+  if (m < 16 || n < 16) {
+    throw std::invalid_argument("[gather_mm] M/N is too small.");
+  }
 
   encoder.set_input_array(a);
   encoder.set_input_array(b);

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -245,7 +245,7 @@ void grouped_gemm_v2(
           LayoutB,
           cutlass::ComplexTransform::kNone,
           GemmConfiguration::kAlignmentAB,
-          typename GemmConfiguration::Accumulator,
+          typename GemmConfiguration::Element,
           cutlass::layout::RowMajor,
           typename GemmConfiguration::Accumulator,
           typename GemmConfiguration::OpClass,

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

@@ -52,7 +52,7 @@ bool supports_qmm_sm90(
   if (!biases) {
     return false;
   }
-  if (!is_last_2_dims_row_contiguous(w) ||
+  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales) ||
       !is_last_2_dims_row_contiguous(*biases)) {
     return false;
@@ -139,7 +139,7 @@ bool supports_qmm_sm80(
   if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) {
     return false;
   }
-  if (!is_last_2_dims_row_contiguous(w) ||
+  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }
@@ -224,7 +224,7 @@ bool supports_qmm_naive(
   if (transpose && (k % std::max(64, group_size) != 0)) {
     return false;
   }
-  if (!is_last_2_dims_row_contiguous(w) ||
+  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }
@@ -265,7 +265,7 @@ void qmm_naive(
     if constexpr (k_major.value) {
       if (has_k_residue) {
         throw std::invalid_argument(
-            "[quantized_matmul] K must be multiples of max(64, group_size).");
+            "[quantized_matmul] K must be multiples of group_size.");
       }
       f.template operator()<false>();
     } else {
@@ -276,8 +276,7 @@ void qmm_naive(
   };
   int m = out.ndim() > 1 ? out.shape(-2) : 1;
   int k = x.shape(-1);
-  int tile_k = std::max(64, group_size);
-  bool has_k_residue = k % tile_k != 0;
+  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>() {
@@ -343,7 +342,7 @@ bool supports_qmv(
   if (k % 8 != 0) {
     return false;
   }
-  if (!is_last_2_dims_row_contiguous(w) ||
+  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }

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

@@ -1,8 +1,9 @@
 // Copyright © 2026 Apple Inc.
 
 #include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
 #include "mlx/backend/cuda/quantized/qmm/qmm.h"
-#include "mlx/backend/cuda/quantized/qmm/qmm_naive.cuh"
+#include "mlx/dtype_utils.h"
 
 // clang-format off
 
@@ -11,34 +12,69 @@ namespace cutlass_gemm {
 
 using namespace cute;
 
-template <bool KMajor, bool HasKResidue, bool SM80,
+template <typename Element, typename SmemLayoutA, typename SmemLayoutB>
+struct SharedStorage {
+  ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
+  ArrayEngine<Element, cosize_v<SmemLayoutB>> B;
+};
+
+__device__ __forceinline__ void
+cute_naive_dequant(auto w, auto s, auto z, auto out) {
+  using Element = typename decltype(out)::value_type;
+  using Quant = typename decltype(w)::value_type;
+  using Scale = typename decltype(s)::value_type;
+  transform(w, out, [](Quant q) { return Element(q); } );
+  transform(out, s, out, [](Element e, Scale s) { return e * Element(s); });
+  if constexpr (quant_has_bias_v<Quant>) {
+    transform(out, z, out, plus{});
+  }
+}
+
+__device__ __forceinline__ void
+cute_dequant(auto w, auto s, auto z, auto out) {
+  if constexpr (stride(coalesce(w.layout())) == Int<1>{} &&
+                is_static_v<decltype(s.layout())>) {
+    cute_vectorized_dequant(w, s, z, out);
+  } else {
+    cute_naive_dequant(w, s, z, out);
+  }
+}
+
+template <bool HasKResidue, typename ProblemShape, typename CtaTiler,
           typename Element, typename Quant, typename Scale,
-          typename ProblemShape,
-          typename CtaTiler,
-          typename StrideA,
-          typename StrideB,
-          typename LayoutS,
-          typename StrideC,
-          typename TiledMma>
-__global__
-__launch_bounds__(decltype(size(TiledMma{}))::value)
-void qmm_naive_kernel(
-    ProblemShape shape_MNKL,
-    CtaTiler cta_tiler,
-    const Element* A, StrideA dA,
-    const Quant* B, StrideB dB,
+          typename StrideA, typename SmemLayoutA, typename TiledCopyA,
+          typename StrideB, typename SmemLayoutB, typename TiledCopyB,
+          typename StrideC, typename LayoutS, typename TiledMma>
+__global__ void qmm_naive_kernel(
+    ProblemShape shape_MNKL, CtaTiler cta_tiler,
+    const Element* A, StrideA dA, SmemLayoutA sA_layout, TiledCopyA copy_a,
+    const Quant*   B, StrideB dB, SmemLayoutB sB_layout, TiledCopyB copy_b,
+          Element* C, StrideC dC,
     const Scale* S, const Element* Z, LayoutS S_layout,
     const uint32_t* lhs_indices, const uint32_t* rhs_indices,
-    Element* C, StrideC dC,
     TiledMma mma) {
+  CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma));
+  CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma));
   CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA));
   CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB));
   CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
 
   int thread_idx = int(threadIdx.x);
-  int m_coord = int(blockIdx.x);
-  int n_coord = int(blockIdx.y);
-  int l_coord = int(blockIdx.z);
+  auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
+
+  auto m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord
+  auto n_max_coord = size<1>(shape_MNKL) - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord
+
+  // Shift tensor so we handle residue of K in the 0th tile.
+  auto shape_K = size<2>(shape_MNKL);
+  auto bK = size<2>(cta_tiler);
+  auto k_residue = shape_K - bK * ceil_div(shape_K, bK);
+  if constexpr (HasKResidue) {
+    A += k_residue * get<1>(dA);
+    B += k_residue * get<1>(dB) * cuda::std::min(8, sizeof_bits_v<Quant>) / 8;
+    S += k_residue * stride<1>(S_layout);
+    Z += k_residue * stride<1>(S_layout);
+  }
 
   // Represent the full tensors.
   Tensor mA_mkl = make_tensor(make_gmem_ptr(A),        select<0,2,3>(shape_MNKL), dA); // (M,K,L)
@@ -69,24 +105,218 @@ void qmm_naive_kernel(
   Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
   Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
 
-  // Compute tile residues for predication.
-  int m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord
-  int n_max_coord = size<1>(shape_MNKL) - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord
-  int k_residue = size<2>(shape_MNKL) - size<1>(gA) * size<2>(gA);
+  // Shared memory buffers.
+  extern __shared__ char shared_memory[];
+  using SharedStorage = SharedStorage<Element, SmemLayoutA, SmemLayoutB>;
+  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(shared_memory);
+  Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K)
+  Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K)
 
-  qmm_naive_mainloop<KMajor, HasKResidue, SM80>(
-      cta_tiler,
-      gA,
-      gB,
-      gS,
-      gZ,
-      gC,
-      mma,
-      m_max_coord, n_max_coord, k_residue,
-      thread_idx);
+  // Partition the copying of A/B/C tiles across the threads.
+  ThrCopy thr_copy_a = copy_a.get_slice(thread_idx);
+  Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
+  Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K)
+  Tensor tArA = make_fragment_like(tAsA);   // (ACPY,ACPY_M,ACPY_K)
+
+  ThrCopy thr_copy_b = copy_b.get_slice(thread_idx);
+  Tensor tBgB = thr_copy_b.partition_S(gB);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBsB = thr_copy_b.partition_D(sB);        // (BCPY,BCPY_N,BCPY_K)
+  Tensor tBrB = make_fragment_like<Quant>(tBsB);   // (BCPY,BCPY_M,BCPY_K)
+  Tensor tBrB_dq = make_fragment_like(tBsB);       // (BCPY,BCPY_M,BCPY_K)
+  Tensor tBgS = thr_copy_b.partition_S(gS);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBrS = make_fragment_like(tBgS(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
+  Tensor tBgZ = thr_copy_b.partition_S(gZ);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBrZ = make_fragment_like(tBgZ(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
+
+  // MMA.
+  ThrMMA thr_mma = mma.get_slice(thread_idx);
+  Tensor tCsA = thr_mma.partition_A(sA);       // (MMA,MMA_M,MMA_K)
+  Tensor tCsB = thr_mma.partition_B(sB);       // (MMA,MMA_N,MMA_K)
+  Tensor tCgC = thr_mma.partition_C(gC);       // (MMA,MMA_M,MMA_N)
+  Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
+
+  // Predicates for m/n bounds.
+  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K)
+  Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K)
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
+  Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K)
+  Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N)
+  Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K)
+  Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K)
+  Tensor tCcC = thr_mma.partition_C(cC);    // (MMA,MMA_M,MMA_N)
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tApA); ++m) {
+    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
+  }
+  CUTE_UNROLL
+  for (int n = 0; n < size<0>(tBpB); ++n) {
+    tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord;
+  }
+
+  // GMEM => RMEM.
+  auto fetch_gmem = [&](int tile) {
+    copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA);
+    copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB);
+    copy(tBgS(_,_,_,tile), tBrS);
+    copy(tBgZ(_,_,_,tile), tBrZ);
+  };
+  // RMEM => SMEM.
+  auto store_smem = [&]() {
+    __syncthreads();
+    copy(tArA, tAsA);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tBrB); ++k) {
+      CUTE_UNROLL
+      for (int n = 0; n < size<1>(tBrB); ++n) {
+        cute_dequant(tBrB(_,n,k), tBrS(_,n,k), tBrZ(_,n,k), tBrB_dq(_,n,k));
+      }
+    }
+    copy(tBrB_dq, tBsB);
+    __syncthreads();
+  };
+
+  // Clear the rmem tiles to account for predicated off loads.
+  if constexpr (HasKResidue) {
+    clear(tArA);
+    clear(tBrB);
+    clear(tBrS);
+    clear(tBrZ);
+  }
+
+  // Prefetch first tile.
+  if constexpr (HasKResidue) {
+    Tensor tAgA_k = tAgA(_,_,_,0);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tArA); ++k) {
+      if (get<1>(tAcA(0,0,k)) >= -k_residue) {
+        copy_if(copy_a, tApA(_,k), tAgA_k(_,_,k), tArA(_,_,k));
+      }
+    }
+    Tensor tBgB_k = tBgB(_,_,_,0);
+    Tensor tBgS_k = tBgS(_,_,_,0);
+    Tensor tBgZ_k = tBgZ(_,_,_,0);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tBrB); ++k) {
+      if (get<1>(tBcB(0,0,k)) >= -k_residue) {
+        copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k));
+        copy(tBgS_k(_,_,k), tBrS(_,_,k));
+        copy(tBgZ_k(_,_,k), tBrZ(_,_,k));
+      }
+    }
+  } else {
+    fetch_gmem(0);
+  }
+
+  // Clear accumulators.
+  clear(tCrC);
+
+  // Loop over CTA tiles.
+  auto K_TILE_MAX = size<3>(tAgA);
+  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
+    store_smem();
+    if constexpr (HasKResidue) {
+      // Avoid fetching full 0th-tile when there is residue.
+      if (K_TILE_MAX > 1) {
+        fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+      }
+    } else {
+      fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+    }
+    gemm(mma, tCsA, tCsB, tCrC);
+  }
+
+  // Epilogue.
+  CUTE_UNROLL
+  for (int i = 0; i < size(tCrC); ++i) {
+    if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) {
+      tCgC(i) = Element(tCrC(i));
+    }
+  }
 }
 
-template <int TileM, bool KMajor, bool HasKResidue, bool SM80,
+template <bool KMajor>
+inline constexpr auto make_matrix_stride(auto m, auto k) {
+  if constexpr (KMajor) {
+    return cute::make_stride(k, cute::Int<1>{}, m * k);
+  } else {
+    return cute::make_stride(cute::Int<1>{}, m, m * k);
+  }
+}
+
+template <bool KMajor = true>
+inline constexpr auto make_smem_layout(auto bM, auto bK) {
+  // TODO: Calculate swizzle based on tile shape.
+  if constexpr (KMajor) {
+    auto swizzle = composition(Swizzle<3,3,3>{},
+                               Layout<Shape <_8,Shape <_8, _8>>,
+                                      Stride<_8,Stride<_1,_64>>>{});
+    return tile_to_shape(swizzle, make_shape(bM, bK));
+  } else {
+    auto swizzle = composition(Swizzle<3,3,3>{},
+                               Layout<Shape<_64,_1>, Stride<_1,_64>>{});
+    return tile_to_shape(swizzle, make_shape(bM, bK));
+  }
+}
+
+template <int TileM, bool SM80, typename Element>
+inline constexpr auto make_tiled_mma() {
+  using Atom = std::conditional_t<
+      SM80,
+      std::conditional_t<
+          std::is_same_v<Element, half_t>,
+          SM80_16x8x16_F32F16F16F32_TN,
+          std::conditional_t<
+              std::is_same_v<Element, bfloat16_t>,
+              SM80_16x8x16_F32BF16BF16F32_TN,
+              UniversalFMA<float>
+          >
+      >,
+      UniversalFMA<float, Element, Element>>;
+  if constexpr (!SM80 || std::is_same_v<Element, float>) {
+    return make_tiled_mma(Atom{}, Layout<Shape<_16,_8,_1>>{});
+  } else {
+    if constexpr (TileM >= 32) {
+      return make_tiled_mma(Atom{}, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
+    } else {
+      return make_tiled_mma(Atom{}, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
+    }
+  }
+}
+
+template <typename T, bool KMajor = true, bool HasKResidue = false>
+inline auto make_tiled_copy(auto num_threads, auto bM, auto bK) {
+  // TODO: Only do 1-element read for the tile of residue.
+  auto n_read = Int<HasKResidue ? 1 : 8>{};
+  auto atom = Copy_Atom<UniversalCopy<uint_bit_t<n_read * sizeof_bits_v<T>>>, T>{};
+  if constexpr (KMajor) {
+    auto k_threads = bK / n_read;
+    return make_tiled_copy(
+        atom,
+        make_layout(make_shape(Int<num_threads / k_threads>{}, k_threads), LayoutRight{}),
+        make_layout(make_shape(Int<1>{}, n_read)));
+  } else {
+    auto m_threads = bM / n_read;
+    return make_tiled_copy(
+        atom,
+        make_layout(make_shape(m_threads, Int<num_threads / m_threads>{}), LayoutLeft{}),
+        make_layout(make_shape(n_read, Int<1>{})));
+  }
+}
+
+template <bool KMajor>
+inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) {
+  if constexpr (KMajor) {
+    return make_layout(
+        make_shape(n, make_shape(group_size, k / group_size), l),
+        make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size));
+  } else {
+    return make_layout(
+        make_shape(make_shape(group_size, n / group_size), k, l),
+        make_stride(Stride<_0,_1>{}, n / group_size, n * k / group_size));
+  }
+}
+
+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,
@@ -101,12 +331,14 @@ void qmm_naive(
     auto group_size,
     auto&& launch_kernel) {
   // Define shapes (dynamic).
-  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)
+  auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L)
 
-  // Define layouts (mixed).
+  // Define TN strides (mixed).
   auto dA = make_stride(k, Int<1>{}, m * k);  // (dM,dK,dL)
   auto dB = make_matrix_stride<KMajor>(n, k); // (dN,dK,dL)
   auto dC = make_stride(n, Int<1>{}, m * n);  // (dM,dN,dL)
+
+  // Define layout of scales/biases (mixed).
   auto S_layout = make_scales_layout<KMajor>(n, k, l, group_size);
 
   // Handle broadcasting.
@@ -115,41 +347,45 @@ void qmm_naive(
     get<2>(stride(S_layout)) = 0;
   }
 
-  // Define CTA tile size (static).
-  auto cta_tiler = make_cta_tiler<TileM, SM80>(group_size);
+  // Define CTA tile sizes (static).
+  auto bM = Int<TileM>{};
+  auto bN = Int<(!SM80 && group_size > 64) ? 64 : 128>{};
+  auto bK = Int<max(64, group_size)>{};
+  auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M,BLK_N,BLK_K)
 
   // Define MMA.
-  auto mma = make_tiled_mma<SM80, Element>(cta_tiler);
+  TiledMMA mma = make_tiled_mma<TileM, SM80, Element>();
   auto num_threads = size(mma);
 
-  // Shared memory size.
-  auto [sA_layout, sB_layout] = make_smem_layouts<KMajor>(cta_tiler);
-  size_t smem_bytes = sizeof(SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>);
+  // Define the A/B smem layouts (static).
+  auto sA_layout = make_smem_layout(bM, bK);
+  auto sB_layout = make_smem_layout<KMajor>(bN, bK);
+
+  // Atoms.
+  TiledCopy copy_a = make_tiled_copy<Element, true, HasKResidue>(num_threads, bM, bK);
+  TiledCopy copy_b = make_tiled_copy<Quant, KMajor>(num_threads, bN, bK);
 
   auto* kernel = &qmm_naive_kernel<
-      KMajor, HasKResidue, SM80,
+      HasKResidue, decltype(prob_shape), decltype(cta_tiler),
       Element, Quant, Scale,
-      decltype(shape_MNKL),
-      decltype(cta_tiler),
-      decltype(dA),
-      decltype(dB),
-      decltype(S_layout),
-      decltype(dC),
-      decltype(mma)>;
-  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+      decltype(dA), decltype(sA_layout), decltype(copy_a),
+      decltype(dB), decltype(sB_layout), decltype(copy_b),
+      decltype(dC), decltype(S_layout), decltype(mma)>;
 
-  dim3 num_blocks{uint32_t(ceil_div(m, size<0>(cta_tiler))),
-                  uint32_t(ceil_div(n, size<1>(cta_tiler))),
-                  uint32_t(l)};
-  dim3 block_dims{num_threads};
+  // Set L1 to be SMEM only.
+  size_t smem_bytes = sizeof(SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>);
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+  cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+
+  dim3 num_blocks(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l);
+  dim3 block_dims(num_threads);
   void* args[] = {
-      &shape_MNKL,
-      &cta_tiler,
-      &A, &dA,
-      &B, &dB,
+      &prob_shape, &cta_tiler,
+      &A, &dA, &sA_layout, &copy_a,
+      &B, &dB, &sB_layout, &copy_b,
+      &C, &dC,
       &S, &Z, &S_layout,
       &lhs_indices, &rhs_indices,
-      &C, &dC,
       &mma};
   launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
 }
@@ -160,6 +396,69 @@ void qmm_naive(
 
 namespace mlx::core {
 
+template <typename F>
+inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
+  if (dtype == float32) {
+    f.template operator()<float>();
+  } else if (dtype == float16) {
+    f.template operator()<cutlass::half_t>();
+  } else if (dtype == bfloat16) {
+    f.template operator()<cutlass::bfloat16_t>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+template <typename F>
+inline void dispatch_groups(int group_size, const char* tag, F&& f) {
+  if (group_size == 32) {
+    f.template operator()<32>();
+  } else if (group_size == 64) {
+    f.template operator()<64>();
+  } else if (group_size == 128) {
+    f.template operator()<128>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Group size {} is not supported.", tag, group_size));
+  }
+}
+
+template <typename T, typename F>
+inline void dispatch_quant_types(
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    const char* tag,
+    F&& f) {
+  if (mode == QuantizationMode::Mxfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Mxfp8) {
+    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Nvfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
+  } else {
+    dispatch_groups(group_size, tag, [&]<int group_size>() {
+      if (bits == 2) {
+        f.template operator()<cutlass::uint2b_t, T, group_size>();
+      } else if (bits == 3) {
+        f.template operator()<cutlass::uint3b_t, T, group_size>();
+      } else if (bits == 4) {
+        f.template operator()<cutlass::uint4b_t, T, group_size>();
+      } else if (bits == 5) {
+        f.template operator()<cutlass::uint5b_t, T, group_size>();
+      } else if (bits == 6) {
+        f.template operator()<cutlass::uint6b_t, T, group_size>();
+      } else if (bits == 8) {
+        f.template operator()<uint8_t, T, group_size>();
+      } else {
+        throw std::invalid_argument(
+            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
+      }
+    });
+  }
+}
+
 template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
 void qmm_naive_impl(
     const array& x,
@@ -217,7 +516,7 @@ void qmm_naive_impl(
               [&](auto* kernel,
                   dim3 num_blocks,
                   dim3 block_dims,
-                  size_t smem_bytes,
+                  uint32_t smem_bytes,
                   void** args) {
                 encoder.add_kernel_node_raw(
                     kernel, num_blocks, block_dims, {}, smem_bytes, args);

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

@@ -1,381 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
-#include "mlx/dtype_utils.h"
-
-// clang-format off
-
-// We can't put kernel code in mlx::core due to name conflicts of "Shape".
-namespace cutlass_gemm {
-
-using namespace cute;
-
-template <typename Element, typename SmemLayoutA, typename SmemLayoutB>
-struct SharedStorage {
-  ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
-  ArrayEngine<Element, cosize_v<SmemLayoutB>> B;
-};
-
-template <bool KMajor = true>
-inline constexpr auto make_smem_layout(auto bM, auto bK) {
-  // TODO: Calculate swizzle based on tile shape.
-  if constexpr (KMajor) {
-    auto swizzle = composition(Swizzle<3,3,3>{},
-                               Layout<Shape <_8,Shape <_8, _8>>,
-                                      Stride<_8,Stride<_1,_64>>>{});
-    return tile_to_shape(swizzle, make_shape(bM, bK));
-  } else {
-    auto swizzle = composition(Swizzle<3,3,3>{},
-                               Layout<Shape<_64,_1>, Stride<_1,_64>>{});
-    return tile_to_shape(swizzle, make_shape(bM, bK));
-  }
-}
-
-template <bool KMajor = true>
-inline constexpr auto make_smem_layouts(auto cta_tiler) {
-  auto [bM, bN, bK] = cta_tiler;
-  auto sA_layout = make_smem_layout(bM, bK);
-  auto sB_layout = make_smem_layout<KMajor>(bN, bK);
-  return std::make_tuple(sA_layout, sB_layout);
-}
-
-template <typename T, bool KMajor = true, bool HasKResidue = false>
-inline constexpr auto make_tiled_copy(auto num_threads, auto bM, auto bK) {
-  // TODO: Only do 1-element read for the tile of residue.
-  auto n_read = Int<HasKResidue ? 1 : 8>{};
-  auto atom = Copy_Atom<UniversalCopy<uint_bit_t<n_read * sizeof_bits_v<T>>>, T>{};
-  if constexpr (KMajor) {
-    auto k_threads = bK / n_read;
-    return make_tiled_copy(
-        atom,
-        make_layout(make_shape(Int<num_threads / k_threads>{}, k_threads), LayoutRight{}),
-        make_layout(make_shape(Int<1>{}, n_read)));
-  } else {
-    auto m_threads = bM / n_read;
-    return make_tiled_copy(
-        atom,
-        make_layout(make_shape(m_threads, Int<num_threads / m_threads>{}), LayoutLeft{}),
-        make_layout(make_shape(n_read, Int<1>{})));
-  }
-}
-
-
-__device__ __forceinline__ void
-cute_naive_dequant(auto w, auto s, auto z, auto out) {
-  using Element = typename decltype(out)::value_type;
-  using Quant = typename decltype(w)::value_type;
-  using Scale = typename decltype(s)::value_type;
-  transform(w, out, [](Quant q) { return Element(q); } );
-  transform(out, s, out, [](Element e, Scale s) { return e * Element(s); });
-  if constexpr (quant_has_bias_v<Quant>) {
-    transform(out, z, out, plus{});
-  }
-}
-
-__device__ __forceinline__ void
-cute_dequant(auto w, auto s, auto z, auto out) {
-  if constexpr (stride(coalesce(w.layout())) == Int<1>{} &&
-                is_static_v<decltype(s.layout())>) {
-    cute_vectorized_dequant(w, s, z, out);
-  } else {
-    cute_naive_dequant(w, s, z, out);
-  }
-}
-
-template <bool KMajor, bool HasKResidue, bool SM80,
-          typename CtaTiler,
-          typename TensorA,
-          typename TensorB,
-          typename TensorS,
-          typename TensorZ,
-          typename TensorC,
-          typename TiledMma>
-CUTE_DEVICE void qmm_naive_mainloop(
-    CtaTiler cta_tiler,
-    TensorA gA,
-    TensorB gB,
-    TensorS gS,
-    TensorZ gZ,
-    TensorC gC,
-    TiledMma mma,
-    int m_max_coord,
-    int n_max_coord,
-    int k_residue,
-    int thread_idx) {
-  // Get the types of operands.
-  using Element = decltype(gA)::value_type;
-  using Quant = decltype(gB)::value_type;
-
-  // Shift tensor so we handle residue of K in the 0th tile.
-  gA = domain_offset(make_coord(0, k_residue, 0), gA);
-  if constexpr (sizeof_bits_v<Quant> % 8 == 0) {
-    gB = domain_offset(make_coord(0, k_residue, 0), gB);
-  } else {
-    gB.data() = recast_ptr<Quant>(raw_pointer_cast(gB.data()) + gB.layout()(0, k_residue, 0) * cuda::std::min(8, sizeof_bits_v<Quant>) / 8);
-  }
-  gS = domain_offset(make_coord(0, k_residue, 0), gS);
-  gZ = domain_offset(make_coord(0, k_residue, 0), gZ);
-
-  // Define smem layouts.
-  auto [sA_layout, sB_layout] = make_smem_layouts(cta_tiler);
-
-  // Shared memory buffer.
-  extern __shared__ char smem_buf[];
-  using SharedStorage = SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>;
-  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(smem_buf);
-  Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K)
-  Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K)
-
-  // Define copy atoms.
-  auto num_threads = size(mma);
-  auto [bM, bN, bK] = cta_tiler;
-  TiledCopy copy_a = make_tiled_copy<Element, true, HasKResidue>(num_threads, bM, bK);
-  TiledCopy copy_b = make_tiled_copy<Quant, KMajor>(num_threads, bN, bK);
-
-  // Partition the copying of A/B/C tiles across the threads.
-  ThrCopy thr_copy_a = copy_a.get_slice(thread_idx);
-  Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
-  Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K)
-  Tensor tArA = make_fragment_like(tAsA);   // (ACPY,ACPY_M,ACPY_K)
-
-  ThrCopy thr_copy_b = copy_b.get_slice(thread_idx);
-  Tensor tBgB = thr_copy_b.partition_S(gB);        // (BCPY,BCPY_N,BCPY_K,k)
-  Tensor tBsB = thr_copy_b.partition_D(sB);        // (BCPY,BCPY_N,BCPY_K)
-  Tensor tBrB = make_fragment_like<Quant>(tBsB);   // (BCPY,BCPY_M,BCPY_K)
-  Tensor tBrB_dq = make_fragment_like(tBsB);       // (BCPY,BCPY_M,BCPY_K)
-  Tensor tBgS = thr_copy_b.partition_S(gS);        // (BCPY,BCPY_N,BCPY_K,k)
-  Tensor tBrS = make_fragment_like(tBgS(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
-  Tensor tBgZ = thr_copy_b.partition_S(gZ);        // (BCPY,BCPY_N,BCPY_K,k)
-  Tensor tBrZ = make_fragment_like(tBgZ(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
-
-  // MMA.
-  ThrMMA thr_mma = mma.get_slice(thread_idx);
-  Tensor tCsA = thr_mma.partition_A(sA);       // (MMA,MMA_M,MMA_K)
-  Tensor tCsB = thr_mma.partition_B(sB);       // (MMA,MMA_N,MMA_K)
-  Tensor tCgC = thr_mma.partition_C(gC);       // (MMA,MMA_M,MMA_N)
-  Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
-
-  // Predicates for m/n bounds.
-  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K)
-  Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K)
-  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
-  Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K)
-  Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N)
-  Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K)
-  Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K)
-  Tensor tCcC = thr_mma.partition_C(cC);    // (MMA,MMA_M,MMA_N)
-  CUTE_UNROLL
-  for (int m = 0; m < size<0>(tApA); ++m) {
-    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
-  }
-  CUTE_UNROLL
-  for (int n = 0; n < size<0>(tBpB); ++n) {
-    tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord;
-  }
-
-  // GMEM => RMEM.
-  auto fetch_gmem = [&](int tile) {
-    copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA);
-    copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB);
-    copy(tBgS(_,_,_,tile), tBrS);
-    copy(tBgZ(_,_,_,tile), tBrZ);
-  };
-  // RMEM => SMEM.
-  auto store_smem = [&]() {
-    __syncthreads();
-    copy(tArA, tAsA);
-    CUTE_UNROLL
-    for (int k = 0; k < size<2>(tBrB); ++k) {
-      CUTE_UNROLL
-      for (int n = 0; n < size<1>(tBrB); ++n) {
-        cute_dequant(tBrB(_,n,k), tBrS(_,n,k), tBrZ(_,n,k), tBrB_dq(_,n,k));
-      }
-    }
-    copy(tBrB_dq, tBsB);
-    __syncthreads();
-  };
-
-  // Clear the rmem tiles to account for predicated off loads.
-  if constexpr (HasKResidue) {
-    clear(tArA);
-    clear(tBrB);
-    clear(tBrS);
-    clear(tBrZ);
-  }
-
-  // Prefetch first tile.
-  if constexpr (HasKResidue) {
-    Tensor tAgA_k = tAgA(_,_,_,0);
-    CUTE_UNROLL
-    for (int k = 0; k < size<2>(tArA); ++k) {
-      if (get<1>(tAcA(0,0,k)) >= -k_residue) {
-        copy_if(copy_a, tApA(_,k), tAgA_k(_,_,k), tArA(_,_,k));
-      }
-    }
-    Tensor tBgB_k = tBgB(_,_,_,0);
-    Tensor tBgS_k = tBgS(_,_,_,0);
-    Tensor tBgZ_k = tBgZ(_,_,_,0);
-    CUTE_UNROLL
-    for (int k = 0; k < size<2>(tBrB); ++k) {
-      if (get<1>(tBcB(0,0,k)) >= -k_residue) {
-        copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k));
-        copy(tBgS_k(_,_,k), tBrS(_,_,k));
-        copy(tBgZ_k(_,_,k), tBrZ(_,_,k));
-      }
-    }
-  } else {
-    fetch_gmem(0);
-  }
-
-  // Clear accumulators.
-  clear(tCrC);
-
-  // Loop over CTA tiles.
-  auto K_TILE_MAX = size<3>(tAgA);
-  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
-    store_smem();
-    if constexpr (HasKResidue) {
-      // Avoid fetching full 0th-tile when there is residue.
-      if (K_TILE_MAX > 1) {
-        fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
-      }
-    } else {
-      fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
-    }
-    gemm(mma, tCsA, tCsB, tCrC);
-  }
-
-  // Epilogue.
-  CUTE_UNROLL
-  for (int i = 0; i < size(tCrC); ++i) {
-    if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) {
-      tCgC(i) = Element(tCrC(i));
-    }
-  }
-}
-
-template <bool KMajor>
-inline constexpr auto make_matrix_stride(auto m, auto k) {
-  if constexpr (KMajor) {
-    return cute::make_stride(k, cute::Int<1>{}, m * k);
-  } else {
-    return cute::make_stride(cute::Int<1>{}, m, m * k);
-  }
-}
-
-template <bool KMajor>
-inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) {
-  if constexpr (KMajor) {
-    return make_layout(
-        make_shape(n, make_shape(group_size, k / group_size), l),
-        make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size));
-  } else {
-    return make_layout(
-        make_shape(make_shape(group_size, n / group_size), k, l),
-        make_stride(Stride<_0,_1>{}, n / group_size, n * k / group_size));
-  }
-}
-
-template <int TileM, bool SM80>
-inline constexpr auto make_cta_tiler(auto group_size) {
-  auto bM = Int<TileM>{};
-  auto bN = Int<(!SM80 && group_size > 64) ? 64 : 128>{};
-  auto bK = Int<max(64, group_size)>{};
-  return make_shape(bM, bN, bK);
-}
-
-template <bool SM80, typename Element>
-inline constexpr auto make_tiled_mma(auto cta_tiler) {
-  using Atom = std::conditional_t<
-      SM80,
-      std::conditional_t<
-          std::is_same_v<Element, half_t>,
-          SM80_16x8x16_F32F16F16F32_TN,
-          std::conditional_t<
-              std::is_same_v<Element, bfloat16_t>,
-              SM80_16x8x16_F32BF16BF16F32_TN,
-              UniversalFMA<float>
-          >
-      >,
-      UniversalFMA<float, Element, Element>>;
-  if constexpr (!SM80 || std::is_same_v<Element, float>) {
-    return make_tiled_mma(Atom{}, Layout<Shape<_16,_8,_1>>{});
-  } else {
-    if constexpr (size<0>(cta_tiler) >= 32) {
-      return make_tiled_mma(Atom{}, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
-    } else {
-      return make_tiled_mma(Atom{}, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
-    }
-  }
-}
-
-} // namespace cutlass_gemm
-
-// clang-format on
-
-namespace mlx::core {
-
-template <typename F>
-inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
-  if (dtype == float32) {
-    f.template operator()<float>();
-  } else if (dtype == float16) {
-    f.template operator()<cutlass::half_t>();
-  } else if (dtype == bfloat16) {
-    f.template operator()<cutlass::bfloat16_t>();
-  } else {
-    throw std::invalid_argument(
-        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
-  }
-}
-
-template <typename F>
-inline void dispatch_groups(int group_size, const char* tag, F&& f) {
-  if (group_size == 32) {
-    f.template operator()<32>();
-  } else if (group_size == 64) {
-    f.template operator()<64>();
-  } else if (group_size == 128) {
-    f.template operator()<128>();
-  } else {
-    throw std::invalid_argument(
-        fmt::format("{} Group size {} is not supported.", tag, group_size));
-  }
-}
-
-template <typename T, typename F>
-inline void dispatch_quant_types(
-    int bits,
-    int group_size,
-    QuantizationMode mode,
-    const char* tag,
-    F&& f) {
-  if (mode == QuantizationMode::Mxfp4) {
-    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
-  } else if (mode == QuantizationMode::Mxfp8) {
-    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
-  } else if (mode == QuantizationMode::Nvfp4) {
-    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
-  } else {
-    dispatch_groups(group_size, tag, [&]<int group_size>() {
-      if (bits == 2) {
-        f.template operator()<cutlass::uint2b_t, T, group_size>();
-      } else if (bits == 3) {
-        f.template operator()<cutlass::uint3b_t, T, group_size>();
-      } else if (bits == 4) {
-        f.template operator()<cutlass::uint4b_t, T, group_size>();
-      } else if (bits == 5) {
-        f.template operator()<cutlass::uint5b_t, T, group_size>();
-      } else if (bits == 6) {
-        f.template operator()<cutlass::uint6b_t, T, group_size>();
-      } else if (bits == 8) {
-        f.template operator()<uint8_t, T, group_size>();
-      } else {
-        throw std::invalid_argument(
-            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
-      }
-    });
-  }
-}
-
-} // namespace mlx::core

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

@@ -1,7 +1,8 @@
 // Copyright © 2026 Apple Inc.
 
+#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
 #include "mlx/backend/cuda/quantized/qmm/qmm.h"
-#include "mlx/backend/cuda/quantized/qmm/qmm_sm80.cuh"
+#include "mlx/dtype_utils.h"
 
 // clang-format off
 
@@ -10,32 +11,44 @@ namespace cutlass_gemm {
 
 using namespace cute;
 
-template <typename Element, typename Quant, typename Scale,
-          typename ProblemShape,
-          typename CtaTiler,
-          typename StrideA,
-          typename StrideB,
-          typename LayoutS,
-          typename StrideC,
-          typename TiledMma>
-__global__
-__launch_bounds__(decltype(size(TiledMma{}))::value)
-void qmm_sm80_kernel(
+template <typename Element,
+          typename Quant,
+          typename SmemLayoutA,
+          typename SmemLayoutB,
+          typename SmemLayoutC>
+union SharedStorage {
+  struct {
+    ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
+    ArrayEngine<Quant,   cosize_v<SmemLayoutB>> B;
+  } mainloop;
+  struct {
+    ArrayEngine<Element, cosize_v<SmemLayoutC>> C;
+  } epilogue;
+};
+
+template <typename ProblemShape, typename CtaTiler,
+          typename Element, typename Quant, typename Scale,
+          typename StrideA, typename SmemLayoutA, typename TiledCopyA, typename S2RAtomA,
+          typename StrideB, typename SmemLayoutB, typename TiledCopyB, typename S2RAtomB,
+          typename StrideC, typename SmemLayoutC, typename TiledCopyC, typename R2SAtomC,
+          typename LayoutS, typename G2RAtomS, typename TiledMma>
+__global__ void qmm_sm80_kernel(
     ProblemShape shape_MNKL, CtaTiler cta_tiler,
-    const Element* A, StrideA dA,
-    const Quant* B, StrideB dB,
-    const Scale* S, const Element* Z, LayoutS S_layout,
+    const Element* A, StrideA dA, SmemLayoutA sA_layout, TiledCopyA g2s_copy_a, S2RAtomA s2r_atom_a,
+    const Quant*   B, StrideB dB, SmemLayoutB sB_layout, TiledCopyB g2s_copy_b, S2RAtomB s2r_atom_b,
+          Element* C, StrideC dC, SmemLayoutC sC_layout, TiledCopyC s2g_copy_c, R2SAtomC r2s_atom_c,
+    const Scale* S, const Element* Z, LayoutS S_layout, G2RAtomS g2r_atom_s,
     const uint32_t* lhs_indices, const uint32_t* rhs_indices,
-    Element* C, StrideC dC,
     TiledMma mma) {
+  CUTE_STATIC_ASSERT_V(size(g2s_copy_a) == size(mma));
+  CUTE_STATIC_ASSERT_V(size(g2s_copy_b) == size(mma));
+  CUTE_STATIC_ASSERT_V(size(s2g_copy_c) == size(mma));
   CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA));
   CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB));
   CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
 
   int thread_idx = int(threadIdx.x);
-  int m_coord = int(blockIdx.x);
-  int n_coord = int(blockIdx.y);
-  int l_coord = int(blockIdx.z);
+  auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
 
   // For gather, use index lookup for input batch slicing.
   uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
@@ -66,23 +79,201 @@ void qmm_sm80_kernel(
   Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
   Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
 
-  // Compute tile residues for predication.
-  auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord
+  // Shared memory buffers.
+  extern __shared__ char shared_memory[];
+  using SharedStorage = SharedStorage<Element, Quant,
+                                      SmemLayoutA,
+                                      SmemLayoutB,
+                                      SmemLayoutC>;
+  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(shared_memory);
+  Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K)
+  Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K)
+  Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N)
 
-  qmm_sm80_mainloop(
-      cta_tiler,
-      gA,
-      gB,
-      gS,
-      gZ,
-      gC,
-      mma,
-      m_max_coord,
-      thread_idx);
+  // Partition the copying of A/B/C tiles across the threads.
+  ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx);
+  Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
+  Tensor tAsA = g2s_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
+
+  ThrCopy g2s_thr_copy_b = g2s_copy_b.get_slice(thread_idx);
+  Tensor tBgB = g2s_thr_copy_b.partition_S(gB);  // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBsB = g2s_thr_copy_b.partition_D(sB);  // (BCPY,BCPY_N,BCPY_K,PIPE)
+
+  ThrCopy s2g_thr_copy_c = s2g_copy_c.get_slice(thread_idx);
+  Tensor s2g_tCsC = s2g_thr_copy_c.partition_S(sC); // (CCPY,CCPY_M,CCPY_N)
+  Tensor s2g_tCgC = s2g_thr_copy_c.partition_D(gC); // (CCPY,CCPY_M,CCPY_N)
+
+  // MMA.
+  ThrMMA thr_mma = mma.get_slice(thread_idx);
+  Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K)
+  Tensor tCsB = thr_mma.partition_B(sB(_,_,0));          // (MMA,MMA_N,MMA_K)
+  Tensor tCrB = make_fragment_like<Quant>(tCsB);         // (MMA,MMA_N,MMA_K)
+  Tensor tCrB_dq = make_fragment_like<Element>(tCsB);    // (MMA,MMA_N,MMA_K)
+  Tensor tCgC = thr_mma.partition_C(gC);                 // (MMA,MMA_M,MMA_N)
+  Tensor tCrC_accu = make_fragment_like<float>(tCgC);    // (MMA,MMA_M,MMA_N)
+  Tensor tCrC = make_fragment_like<Element>(tCgC);       // (MMA,MMA_M,MMA_N)
+
+  Tensor tCgS = thr_mma.partition_B(gS);         // (MMA,MMA_N,MMA_K,k)
+  Tensor tCrS = make_tensor_like(tCgS(_,_,_,0)); // (MMA,MMA_N,MMA_K)
+  Tensor tCgZ = thr_mma.partition_B(gZ);         // (MMA,MMA_N,MMA_K,k)
+  Tensor tCrZ = make_tensor_like(tCgZ(_,_,_,0)); // (MMA,MMA_N,MMA_K)
+
+  // Copy Atom retiling.
+  TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma);
+  ThrCopy s2r_thr_copy_a = s2r_copy_a.get_slice(thread_idx);
+  Tensor s2r_tCsA = s2r_thr_copy_a.partition_S(sA); // (ACPY,MMA_M,MMA_K,PIPE)
+  Tensor s2r_tCrA = s2r_thr_copy_a.retile_D(tCrA);  // (ACPY,MMA_M,MMA_K)
+
+  TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_b, mma);
+  ThrCopy s2r_thr_copy_b = s2r_copy_b.get_slice(thread_idx);
+  Tensor s2r_tCsB = s2r_thr_copy_b.partition_S(sB); // (BCPY,MMA_N,MMA_K,PIPE)
+  Tensor s2r_tCrB = s2r_thr_copy_b.retile_D(tCrB);  // (BCPY,MMA_N,MMA_K)
+
+  TiledCopy r2s_copy_c = make_tiled_copy_C(r2s_atom_c, mma);
+  ThrCopy r2s_thr_copy_c = r2s_copy_c.get_slice(thread_idx);
+  Tensor r2s_tCrC = r2s_thr_copy_c.retile_S(tCrC);  // (CCPY,MMA_M,MMA_N)
+  Tensor r2s_tCsC = r2s_thr_copy_c.partition_D(sC); // (CCPY,MMA_M,MMA_N)
+
+  TiledCopy g2r_copy_s = make_tiled_copy_B(g2r_atom_s, mma);
+  ThrCopy g2r_thr_copy_s = g2r_copy_s.get_slice(thread_idx);
+  Tensor g2r_tCgS = g2r_thr_copy_s.partition_S(gS); // (BCPY,MMA_N,MMA_K,k)
+  Tensor g2r_tCrS = g2r_thr_copy_s.retile_D(tCrS);  // (BCPY,MMA_N,MMA_K)
+  Tensor g2r_tCgZ = g2r_thr_copy_s.partition_S(gZ); // (BCPY,MMA_N,MMA_K,k)
+  Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ);  // (BCPY,MMA_N,MMA_K)
+
+  // Predicates for m bound.
+  auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord
+  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});         // (CPY_M,CPY_K)
+  Tensor tCpC = make_tensor<bool>(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N)
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
+  Tensor cC = make_identity_tensor(make_shape(size<0>(sC), size<1>(sC))); // (BLK_M,BLK_N)
+  Tensor tAcA = g2s_thr_copy_a.partition_D(cA); // (CPY,CPY_M,CPY_K)
+  Tensor tCcC = s2g_thr_copy_c.partition_D(cC); // (CPY,CPY_M,CPY_N)
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tApA); ++m) {
+    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
+  }
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tCpC); ++m) {
+    tCpC(m,0) = get<0>(tCcC(0,m,0)) < m_max_coord;
+  }
+
+  auto K_PIPE_MAX = size<3>(tAsA);
+  int smem_pipe_read = 0;
+  int smem_pipe_write = 0;
+
+  // Copy A/B: GMEM => SMEM.
+  auto fetch_gmem = [&](int tile) {
+    copy_if(g2s_copy_a, tApA, tAgA(_,_,_,tile), tAsA(_,_,_,smem_pipe_write));
+    copy(g2s_copy_b, tBgB(_,_,_,tile), tBsB(_,_,_,smem_pipe_write));
+    cp_async_fence();
+    smem_pipe_write = (smem_pipe_write + 1) % K_PIPE_MAX;
+  };
+  // Copy S/Z: GMEM => RMEM.
+  auto fetch_scales = [&](int tile) {
+    copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS);
+    if constexpr (quant_has_bias_v<Quant>) {
+      copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ);
+    }
+  };
+  // Copy A/B: SMEM => RMEM.
+  auto fetch_smem = [&](auto block) {
+    copy(s2r_atom_a, s2r_tCsA(_,_,block,smem_pipe_read), s2r_tCrA(_,_,block));
+    copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block));
+    CUTE_UNROLL
+    for (int n = 0; n < size<1>(tCrB); ++n) {
+      cute_vectorized_dequant(
+          tCrB(_,n,block),
+          tCrS(_,n,block),
+          tCrZ(_,n,block),
+          tCrB_dq(_,n,block));
+    }
+  };
+
+  auto K_TILE_MAX = size<3>(tAgA);
+  auto K_BLOCK_MAX = size<2>(tCrA);
+
+  // Prefetch beginning tiles.
+  int tile_pipe = 0;
+  CUTE_UNROLL
+  for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) {
+    fetch_gmem(tile_pipe);
+  }
+
+  // Clear accumulators.
+  clear(tCrC_accu);
+
+  // Prefetch first block.
+  if constexpr (K_BLOCK_MAX > 1) {
+    cp_async_wait<K_PIPE_MAX - 2>();
+    __syncthreads();
+    fetch_scales(0);
+    fetch_smem(Int<0>{});
+  }
+
+  // Loop over CTA tiles.
+  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
+    // Unroll MMA blocks.
+    CUTE_UNROLL
+    for (int block = 0; block < K_BLOCK_MAX; ++block) {
+      // Wait for last tile.
+      if (block == K_BLOCK_MAX - 1) {
+        smem_pipe_read = (smem_pipe_read + 1) % K_PIPE_MAX;
+        cp_async_wait<K_PIPE_MAX - 2>();
+        __syncthreads();
+        fetch_scales((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+      }
+      // Prefetch next block.
+      fetch_smem((block + 1) % K_BLOCK_MAX);
+      // Prefetch next tile.
+      if (block == 0) {
+        fetch_gmem(tile_pipe);
+        tile_pipe = (tile_pipe + 1 < K_TILE_MAX) ? tile_pipe + 1 : tile_pipe;
+      }
+      // MMA.
+      gemm(mma, tCrA(_,_,block), tCrB_dq(_,_,block), tCrC_accu);
+    }
+  }
+
+  // Epilogue.
+  CUTE_UNROLL
+  for (int i = 0; i < size(tCrC_accu); i++) {
+    tCrC(i) = Element(tCrC_accu(i));
+  }
+  copy(r2s_copy_c, r2s_tCrC, r2s_tCsC);
+  __syncthreads();
+  copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC);
 }
 
-template <int TileM,
-          typename Element, typename Quant, typename Scale>
+template <typename Element>
+inline constexpr auto make_mma_atom() {
+  if constexpr (std::is_same_v<Element, half_t>) {
+    return SM80_16x8x16_F32F16F16F32_TN{};
+  }
+  if constexpr (std::is_same_v<Element, bfloat16_t>) {
+    return SM80_16x8x16_F32BF16BF16F32_TN{};
+  }
+}
+
+template <int TileM, typename Element>
+inline constexpr auto make_tiled_mma() {
+  constexpr auto atom = make_mma_atom<Element>();
+  if constexpr (TileM >= 32) {
+    return make_tiled_mma(atom, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
+  } else {
+    return make_tiled_mma(atom, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
+  }
+}
+
+template <typename T, int bits, template <typename U> typename Atom, typename NumThreads>
+inline auto make_tiled_copy(NumThreads num_threads) {
+  return make_tiled_copy(
+      Copy_Atom<Atom<uint_bit_t<bits>>, T>{},
+      make_layout(make_shape(Int<num_threads / 8>{}, Int<8>{}), LayoutRight{}),
+      make_layout(make_shape(Int<1>{}, Int<bits / sizeof_bits_v<T>>{})));
+}
+
+template <int TileM = 16, typename Element, typename Quant, typename Scale, typename GroupSize>
 void qmm_sm80(
     const Element* A,
     const Quant*   B,
@@ -93,16 +284,20 @@ void qmm_sm80(
     Element* C,
     int m, int n, int k, int l,
     bool broadcast_b,
-    auto group_size,
+    GroupSize group_size,
     auto&& launch_kernel) {
   // Define shapes (dynamic).
-  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)
+  auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L)
 
-  // Define layouts (mixed).
+  // Define TN strides (mixed).
   auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL)
   auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL)
   auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL)
-  auto S_layout = make_scales_layout(n, k, l, group_size);
+
+  // Define layout of scales/biases (mixed).
+  auto S_layout = make_layout(
+      make_shape(n, make_shape(group_size, k / group_size), l),
+      make_stride(k / group_size, Stride<_0, _1>{}, n * k / group_size));
 
   // Handle broadcasting.
   if (broadcast_b) {
@@ -111,41 +306,70 @@ void qmm_sm80(
   }
 
   // Define CTA tile sizes (static).
-  auto cta_tiler = make_cta_tiler<TileM>(group_size);
+  auto bM = Int<TileM>{};
+  auto bN = Int<128>{};
+  auto bK = Int<max(64, group_size)>{};
+  auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M,BLK_N,BLK_K)
 
   // Define MMA.
   TiledMMA mma = make_tiled_mma<TileM, Element>();
   auto num_threads = size(mma);
 
-  // Shared memory size.
-  auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler);
+  // Define the A/B smem layouts (static).
+  auto swizzle_ab = composition(Swizzle<3,3,3>{},
+                                Layout<Shape <_8,Shape <_8, _8>>,
+                                       Stride<_8,Stride<_1,_64>>>{});
+  auto bP = Int<3>{}; // pipeline
+  auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP));
+  auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP));
+
+  // Define the C smem layouts (static).
+  // TODO: Find a better swizzle.
+  auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN));
+
+  // Define the scales/biases smem layouts (static).
+  auto bS = ceil_div(bK, group_size);
+  auto sS_layout = make_layout(make_shape(bN, make_shape(group_size, bS)),
+                               make_stride(bS, Stride<_0, _1>{}));
+
+  // Atoms.
+  constexpr int element_bits = sizeof_bits_v<Element>;
+  constexpr int quant_bits = sizeof_bits_v<Quant>;
+  constexpr int qload = 128 / (element_bits / quant_bits);
+  TiledCopy g2s_copy_a = make_tiled_copy<Element, 128, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
+  TiledCopy g2s_copy_b = make_tiled_copy<Quant, qload, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
+  TiledCopy s2g_copy_c = make_tiled_copy<Element, 128, UniversalCopy>(num_threads);
+
+  Copy_Atom<SM75_U32x4_LDSM_N, Element> s2r_atom_a;
+  Copy_Atom<UniversalCopy<uint_bit_t<2 * quant_bits>>, Quant> s2r_atom_b;
+  Copy_Atom<UniversalCopy<uint_bit_t<2 * element_bits>>, Element> r2s_atom_c;
+  Copy_Atom<UniversalCopy<Scale>, Scale> g2r_atom_s;
+
+  auto* kernel = &qmm_sm80_kernel<
+      decltype(prob_shape), decltype(cta_tiler),
+      Element, Quant, Scale,
+      decltype(dA), decltype(sA_layout), decltype(g2s_copy_a), decltype(s2r_atom_a),
+      decltype(dB), decltype(sB_layout), decltype(g2s_copy_b), decltype(s2r_atom_b),
+      decltype(dC), decltype(sC_layout), decltype(s2g_copy_c), decltype(r2s_atom_c),
+      decltype(S_layout), decltype(g2r_atom_s), decltype(mma)>;
+
+  // Set L1 to be SMEM only.
   size_t smem_bytes = sizeof(SharedStorage<Element, Quant,
                                            decltype(sA_layout),
                                            decltype(sB_layout),
                                            decltype(sC_layout)>);
-
-  auto* kernel = &qmm_sm80_kernel<
-      Element, Quant, Scale,
-      decltype(shape_MNKL),
-      decltype(cta_tiler),
-      decltype(dA),
-      decltype(dB),
-      decltype(S_layout),
-      decltype(dC),
-      decltype(mma)>;
   cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+  cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
 
-  dim3 num_blocks{uint32_t(ceil_div(m, size<0>(cta_tiler))),
-                  uint32_t(ceil_div(n, size<1>(cta_tiler))),
-                  uint32_t(l)};
-  dim3 block_dims{num_threads};
+  dim3 num_blocks(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l);
+  dim3 block_dims(num_threads);
   void* args[] = {
-      &shape_MNKL, &cta_tiler,
-      &A, &dA,
-      &B, &dB,
-      &S, &Z, &S_layout,
+      &prob_shape, &cta_tiler,
+      &A, &dA, &sA_layout, &g2s_copy_a, &s2r_atom_a,
+      &B, &dB, &sB_layout, &g2s_copy_b, &s2r_atom_b,
+      &C, &dC, &sC_layout, &s2g_copy_c, &r2s_atom_c,
+      &S, &Z, &S_layout, &g2r_atom_s,
       &lhs_indices, &rhs_indices,
-      &C, &dC,
       &mma};
   launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
 }
@@ -156,6 +380,59 @@ void qmm_sm80(
 
 namespace mlx::core {
 
+template <typename F>
+inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
+  if (dtype == float16) {
+    f.template operator()<cutlass::half_t>();
+  } else if (dtype == bfloat16) {
+    f.template operator()<cutlass::bfloat16_t>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+template <typename F>
+inline void dispatch_groups(int group_size, const char* tag, F&& f) {
+  if (group_size == 32) {
+    f.template operator()<32>();
+  } else if (group_size == 64) {
+    f.template operator()<64>();
+  } else if (group_size == 128) {
+    f.template operator()<128>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Group size {} is not supported.", tag, group_size));
+  }
+}
+
+template <typename T, typename F>
+inline void dispatch_quant_types(
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    const char* tag,
+    F&& f) {
+  if (mode == QuantizationMode::Mxfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Mxfp8) {
+    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Nvfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
+  } else {
+    dispatch_groups(group_size, tag, [&]<int group_size>() {
+      if (bits == 4) {
+        f.template operator()<cutlass::uint4b_t, T, group_size>();
+      } else if (bits == 8) {
+        f.template operator()<uint8_t, T, group_size>();
+      } else {
+        throw std::invalid_argument(
+            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
+      }
+    });
+  }
+}
+
 template <int TileM>
 void qmm_sm80_impl(
     const array& x,
@@ -213,7 +490,7 @@ void qmm_sm80_impl(
               [&](auto* kernel,
                   dim3 num_blocks,
                   dim3 block_dims,
-                  size_t smem_bytes,
+                  uint32_t smem_bytes,
                   void** args) {
                 encoder.add_kernel_node_raw(
                     kernel, num_blocks, block_dims, {}, smem_bytes, args);

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

@@ -1,346 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
-#include "mlx/dtype_utils.h"
-
-// clang-format off
-
-// We can't put kernel code in mlx::core due to name conflicts of "Shape".
-namespace cutlass_gemm {
-
-using namespace cute;
-
-template <typename Element,
-          typename Quant,
-          typename SmemLayoutA,
-          typename SmemLayoutB,
-          typename SmemLayoutC>
-union SharedStorage {
-  struct {
-    ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
-    ArrayEngine<Quant,   cosize_v<SmemLayoutB>> B;
-  } mainloop;
-  struct {
-    ArrayEngine<Element, cosize_v<SmemLayoutC>> C;
-  } epilogue;
-};
-
-inline constexpr auto make_smem_layouts(auto cta_tiler) {
-  // Define the A/B smem layouts (static).
-  auto swizzle_ab = composition(Swizzle<3,3,3>{},
-                                Layout<Shape <_8,Shape <_8, _8>>,
-                                       Stride<_8,Stride<_1,_64>>>{});
-  auto [bM, bN, bK] = cta_tiler;
-  auto bP = Int<3>{}; // pipeline
-  auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP));
-  auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP));
-
-  // Define the C smem layouts (static).
-  // TODO: Find a better swizzle.
-  auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN));
-
-  return std::make_tuple(sA_layout, sB_layout, sC_layout);
-}
-
-template <typename T, int bits, template <typename U> typename Atom>
-inline constexpr auto make_tiled_copy(auto num_threads) {
-  return make_tiled_copy(
-      Copy_Atom<Atom<uint_bit_t<bits>>, T>{},
-      make_layout(make_shape(Int<num_threads / 8>{}, Int<8>{}), LayoutRight{}),
-      make_layout(make_shape(Int<1>{}, Int<bits / sizeof_bits_v<T>>{})));
-}
-
-template <typename CtaTiler,
-          typename TensorA,
-          typename TensorB,
-          typename TensorS,
-          typename TensorZ,
-          typename TensorC,
-          typename TiledMma>
-CUTE_DEVICE void qmm_sm80_mainloop(
-    CtaTiler cta_tiler,
-    TensorA gA,
-    TensorB gB,
-    TensorS gS,
-    TensorZ gZ,
-    TensorC gC,
-    TiledMma mma,
-    int m_max_coord,
-    int thread_idx) {
-  // Get the types of operands.
-  using Element = decltype(gA)::value_type;
-  using Quant = decltype(gB)::value_type;
-  using Scale = decltype(gS)::value_type;
-
-  // Define smem layouts.
-  auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler);
-
-  // Shared memory buffer.
-  extern __shared__ char smem_buf[];
-  using SharedStorage = SharedStorage<Element, Quant,
-                                      decltype(sA_layout),
-                                      decltype(sB_layout),
-                                      decltype(sC_layout)>;
-  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(smem_buf);
-  Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K)
-  Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K)
-  Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N)
-
-  // Define copy atoms.
-  constexpr int element_bits = sizeof_bits_v<Element>;
-  constexpr int quant_bits = sizeof_bits_v<Quant>;
-  constexpr int qload = 128 / (element_bits / quant_bits);
-  auto num_threads = size(mma);
-  TiledCopy g2s_copy_a = make_tiled_copy<Element, 128, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
-  TiledCopy g2s_copy_b = make_tiled_copy<Quant, qload, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
-  TiledCopy s2g_copy_c = make_tiled_copy<Element, 128, UniversalCopy>(num_threads);
-
-  Copy_Atom<SM75_U32x4_LDSM_N, Element> s2r_atom_a;
-  Copy_Atom<UniversalCopy<uint_bit_t<2 * quant_bits>>, Quant> s2r_atom_b;
-  Copy_Atom<UniversalCopy<uint_bit_t<2 * element_bits>>, Element> r2s_atom_c;
-  Copy_Atom<UniversalCopy<Scale>, Scale> g2r_atom_s;
-
-  // Partition the copying of A/B/C tiles across the threads.
-  ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx);
-  Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
-  Tensor tAsA = g2s_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
-
-  ThrCopy g2s_thr_copy_b = g2s_copy_b.get_slice(thread_idx);
-  Tensor tBgB = g2s_thr_copy_b.partition_S(gB);  // (BCPY,BCPY_N,BCPY_K,k)
-  Tensor tBsB = g2s_thr_copy_b.partition_D(sB);  // (BCPY,BCPY_N,BCPY_K,PIPE)
-
-  ThrCopy s2g_thr_copy_c = s2g_copy_c.get_slice(thread_idx);
-  Tensor s2g_tCsC = s2g_thr_copy_c.partition_S(sC); // (CCPY,CCPY_M,CCPY_N)
-  Tensor s2g_tCgC = s2g_thr_copy_c.partition_D(gC); // (CCPY,CCPY_M,CCPY_N)
-
-  // MMA.
-  ThrMMA thr_mma = mma.get_slice(thread_idx);
-  Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K)
-  Tensor tCsB = thr_mma.partition_B(sB(_,_,0));          // (MMA,MMA_N,MMA_K)
-  Tensor tCrB = make_fragment_like<Quant>(tCsB);         // (MMA,MMA_N,MMA_K)
-  Tensor tCrB_dq = make_fragment_like<Element>(tCsB);    // (MMA,MMA_N,MMA_K)
-  Tensor tCgC = thr_mma.partition_C(gC);                 // (MMA,MMA_M,MMA_N)
-  Tensor tCrC_accu = make_fragment_like<float>(tCgC);    // (MMA,MMA_M,MMA_N)
-  Tensor tCrC = make_fragment_like<Element>(tCgC);       // (MMA,MMA_M,MMA_N)
-
-  Tensor tCgS = thr_mma.partition_B(gS);         // (MMA,MMA_N,MMA_K,k)
-  Tensor tCrS = make_tensor_like(tCgS(_,_,_,0)); // (MMA,MMA_N,MMA_K)
-  Tensor tCgZ = thr_mma.partition_B(gZ);         // (MMA,MMA_N,MMA_K,k)
-  Tensor tCrZ = make_tensor_like(tCgZ(_,_,_,0)); // (MMA,MMA_N,MMA_K)
-
-  // Copy Atom retiling.
-  TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma);
-  ThrCopy s2r_thr_copy_a = s2r_copy_a.get_slice(thread_idx);
-  Tensor s2r_tCsA = s2r_thr_copy_a.partition_S(sA); // (ACPY,MMA_M,MMA_K,PIPE)
-  Tensor s2r_tCrA = s2r_thr_copy_a.retile_D(tCrA);  // (ACPY,MMA_M,MMA_K)
-
-  TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_b, mma);
-  ThrCopy s2r_thr_copy_b = s2r_copy_b.get_slice(thread_idx);
-  Tensor s2r_tCsB = s2r_thr_copy_b.partition_S(sB); // (BCPY,MMA_N,MMA_K,PIPE)
-  Tensor s2r_tCrB = s2r_thr_copy_b.retile_D(tCrB);  // (BCPY,MMA_N,MMA_K)
-
-  TiledCopy r2s_copy_c = make_tiled_copy_C(r2s_atom_c, mma);
-  ThrCopy r2s_thr_copy_c = r2s_copy_c.get_slice(thread_idx);
-  Tensor r2s_tCrC = r2s_thr_copy_c.retile_S(tCrC);  // (CCPY,MMA_M,MMA_N)
-  Tensor r2s_tCsC = r2s_thr_copy_c.partition_D(sC); // (CCPY,MMA_M,MMA_N)
-
-  TiledCopy g2r_copy_s = make_tiled_copy_B(g2r_atom_s, mma);
-  ThrCopy g2r_thr_copy_s = g2r_copy_s.get_slice(thread_idx);
-  Tensor g2r_tCgS = g2r_thr_copy_s.partition_S(gS); // (BCPY,MMA_N,MMA_K,k)
-  Tensor g2r_tCrS = g2r_thr_copy_s.retile_D(tCrS);  // (BCPY,MMA_N,MMA_K)
-  Tensor g2r_tCgZ = g2r_thr_copy_s.partition_S(gZ); // (BCPY,MMA_N,MMA_K,k)
-  Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ);  // (BCPY,MMA_N,MMA_K)
-
-  // Predicates for m bound.
-  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});         // (CPY_M,CPY_K)
-  Tensor tCpC = make_tensor<bool>(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N)
-  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
-  Tensor cC = make_identity_tensor(make_shape(size<0>(sC), size<1>(sC))); // (BLK_M,BLK_N)
-  Tensor tAcA = g2s_thr_copy_a.partition_D(cA); // (CPY,CPY_M,CPY_K)
-  Tensor tCcC = s2g_thr_copy_c.partition_D(cC); // (CPY,CPY_M,CPY_N)
-  CUTE_UNROLL
-  for (int m = 0; m < size<0>(tApA); ++m) {
-    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
-  }
-  CUTE_UNROLL
-  for (int m = 0; m < size<0>(tCpC); ++m) {
-    tCpC(m,0) = get<0>(tCcC(0,m,0)) < m_max_coord;
-  }
-
-  auto K_PIPE_MAX = size<3>(tAsA);
-  int smem_pipe_read = 0;
-  int smem_pipe_write = 0;
-
-  // Copy A/B: GMEM => SMEM.
-  auto fetch_gmem = [&](int tile) {
-    copy_if(g2s_copy_a, tApA, tAgA(_,_,_,tile), tAsA(_,_,_,smem_pipe_write));
-    copy(g2s_copy_b, tBgB(_,_,_,tile), tBsB(_,_,_,smem_pipe_write));
-    cp_async_fence();
-    smem_pipe_write = (smem_pipe_write + 1) % K_PIPE_MAX;
-  };
-  // Copy S/Z: GMEM => RMEM.
-  auto fetch_scales = [&](int tile) {
-    copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS);
-    if constexpr (quant_has_bias_v<Quant>) {
-      copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ);
-    }
-  };
-  // Copy A/B: SMEM => RMEM.
-  auto fetch_smem = [&](auto block) {
-    copy(s2r_atom_a, s2r_tCsA(_,_,block,smem_pipe_read), s2r_tCrA(_,_,block));
-    copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block));
-    CUTE_UNROLL
-    for (int n = 0; n < size<1>(tCrB); ++n) {
-      cute_vectorized_dequant(
-          tCrB(_,n,block),
-          tCrS(_,n,block),
-          tCrZ(_,n,block),
-          tCrB_dq(_,n,block));
-    }
-  };
-
-  auto K_TILE_MAX = size<3>(tAgA);
-  auto K_BLOCK_MAX = size<2>(tCrA);
-
-  // Prefetch beginning tiles.
-  int tile_pipe = 0;
-  CUTE_UNROLL
-  for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) {
-    fetch_gmem(tile_pipe);
-  }
-
-  // Clear accumulators.
-  clear(tCrC_accu);
-
-  // Prefetch first block.
-  if constexpr (K_BLOCK_MAX > 1) {
-    cp_async_wait<K_PIPE_MAX - 2>();
-    __syncthreads();
-    fetch_scales(0);
-    fetch_smem(Int<0>{});
-  }
-
-  // Loop over CTA tiles.
-  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
-    // Unroll MMA blocks.
-    CUTE_UNROLL
-    for (int block = 0; block < K_BLOCK_MAX; ++block) {
-      // Wait for last tile.
-      if (block == K_BLOCK_MAX - 1) {
-        smem_pipe_read = (smem_pipe_read + 1) % K_PIPE_MAX;
-        cp_async_wait<K_PIPE_MAX - 2>();
-        __syncthreads();
-        fetch_scales((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
-      }
-      // Prefetch next block.
-      fetch_smem((block + 1) % K_BLOCK_MAX);
-      // Prefetch next tile.
-      if (block == 0) {
-        fetch_gmem(tile_pipe);
-        tile_pipe = (tile_pipe + 1 < K_TILE_MAX) ? tile_pipe + 1 : tile_pipe;
-      }
-      // MMA.
-      gemm(mma, tCrA(_,_,block), tCrB_dq(_,_,block), tCrC_accu);
-    }
-  }
-
-  // Epilogue.
-  CUTE_UNROLL
-  for (int i = 0; i < size(tCrC_accu); i++) {
-    tCrC(i) = Element(tCrC_accu(i));
-  }
-  copy(r2s_copy_c, r2s_tCrC, r2s_tCsC);
-  __syncthreads();
-  copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC);
-}
-
-inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) {
-  return make_layout(
-      make_shape(n, make_shape(group_size, k / group_size), l),
-      make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size));
-}
-
-template <int TileM>
-inline constexpr auto make_cta_tiler(auto group_size) {
-  auto bM = Int<TileM>{};
-  auto bN = Int<128>{};
-  auto bK = Int<max(64, group_size)>{};
-  return make_shape(bM, bN, bK);
-}
-
-template <int TileM, typename Element>
-inline constexpr auto make_tiled_mma() {
-  using Atom = std::conditional_t<
-      std::is_same_v<Element, half_t>,
-      SM80_16x8x16_F32F16F16F32_TN,
-      std::conditional_t<
-          std::is_same_v<Element, bfloat16_t>,
-          SM80_16x8x16_F32BF16BF16F32_TN,
-          UniversalFMA<float>>>;
-  if constexpr (TileM >= 32) {
-    return make_tiled_mma(Atom{}, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
-  } else {
-    return make_tiled_mma(Atom{}, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
-  }
-}
-
-} // namespace cutlass_gemm
-
-// clang-format on
-
-namespace mlx::core {
-
-template <typename F>
-inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
-  if (dtype == float16) {
-    f.template operator()<cutlass::half_t>();
-  } else if (dtype == bfloat16) {
-    f.template operator()<cutlass::bfloat16_t>();
-  } else {
-    throw std::invalid_argument(
-        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
-  }
-}
-
-template <typename F>
-inline void dispatch_groups(int group_size, const char* tag, F&& f) {
-  if (group_size == 32) {
-    f.template operator()<32>();
-  } else if (group_size == 64) {
-    f.template operator()<64>();
-  } else if (group_size == 128) {
-    f.template operator()<128>();
-  } else {
-    throw std::invalid_argument(
-        fmt::format("{} Group size {} is not supported.", tag, group_size));
-  }
-}
-
-template <typename T, typename F>
-inline void dispatch_quant_types(
-    int bits,
-    int group_size,
-    QuantizationMode mode,
-    const char* tag,
-    F&& f) {
-  if (mode == QuantizationMode::Mxfp4) {
-    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
-  } else if (mode == QuantizationMode::Mxfp8) {
-    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
-  } else if (mode == QuantizationMode::Nvfp4) {
-    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
-  } else {
-    dispatch_groups(group_size, tag, [&]<int group_size>() {
-      if (bits == 4) {
-        f.template operator()<cutlass::uint4b_t, T, group_size>();
-      } else if (bits == 8) {
-        f.template operator()<uint8_t, T, group_size>();
-      } else {
-        throw std::invalid_argument(
-            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
-      }
-    });
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/quantized/quantized.cpp

@@ -17,7 +17,7 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  array x = ensure_row_contiguous(inputs[0], encoder, s);
+  const array& x = inputs[0];
   const array& w = inputs[1];
   const array& scales = inputs[2];
   std::optional<array> biases;
@@ -146,17 +146,15 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  array x = ensure_row_contiguous(inputs[0], encoder, s);
+  const array& x = inputs[0];
   const array& w = inputs[1];
   const array& scales = inputs[2];
   std::optional<array> biases;
   if (inputs.size() == 6) {
     biases = inputs[3];
   }
-  array lhs_indices =
-      ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s);
-  array rhs_indices =
-      ensure_row_contiguous(inputs[inputs.size() - 1], encoder, s);
+  array lhs_indices = ensure_contiguous(inputs[inputs.size() - 2], encoder, s);
+  array rhs_indices = ensure_contiguous(inputs[inputs.size() - 1], encoder, s);
 
   int M = out.ndim() > 1 ? out.shape(-2) : 1;
   int N = out.shape(-1);

mlx/backend/metal/device_info.cpp

@@ -13,12 +13,7 @@ bool is_available() {
 }
 
 int device_count() {
-  try {
-    metal::device(Device::gpu);
-    return 1;
-  } catch (...) {
-    return 0;
-  }
+  return 1;
 }
 
 const std::unordered_map<std::string, std::variant<std::string, size_t>>&

mlx/backend/metal/kernels/indexing/scatter.h

@@ -80,9 +80,7 @@ template <
     uint3 gsize [[threads_per_grid]]) {
   Op op;
 
-  IdxT idx =
-      (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK +
-      IdxT(gid.x) * NWORK;
+  IdxT idx = IdxT(gid.z) * gsize.y + gid.y * gsize.x + gid.x * NWORK;
   IdxT out_idx;
   IdxT update_idx;
 

mlx/distributed/jaccl/lib/README.md

@@ -29,8 +29,6 @@ in macOS 26.2.
 - **Point-to-Point Operations**:
   - `send`: Send data to a specific node
   - `recv`: Receive data from a specific node
-- **Synchronization**:
-  - `barrier`: Block until all nodes in the group reach this point
 - **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32,
   Float64, Complex64
 
@@ -288,9 +286,6 @@ class Group {
   // Simple send/recv primitives.
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
-
-  // Block until every rank reaches this point.
-  virtual void barrier() = 0;
 };
 ```
 

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

@@ -35,7 +35,8 @@ endfunction()
 # Examples
 build_example(minimal_env.cpp)
 build_example(minimal_cfg.cpp)
-build_example(minimal_barrier.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/minimal_barrier.cpp

@@ -1,42 +0,0 @@
-// Copyright © 2026 Apple Inc.
-//
-// Exercises Group::barrier(). Ranks arrive at the barrier at staggered times;
-// after the barrier returns we do a small all_sum to confirm the group is
-// healthy and that barrier() carried the correct fence semantics.
-
-#include <chrono>
-#include <iostream>
-#include <thread>
-
-#include <jaccl/jaccl.h>
-
-int main() {
-  auto group = jaccl::init();
-  if (!group) {
-    std::cerr << "Failed to initialize JACCL" << std::endl;
-    return 1;
-  }
-
-  int rank = group->rank();
-  int size = group->size();
-
-  std::this_thread::sleep_for(std::chrono::milliseconds(100 * rank));
-  std::cout << "rank " << rank << " entering barrier" << std::endl;
-
-  group->barrier();
-
-  std::cout << "rank " << rank << " exited barrier" << std::endl;
-
-  int in = rank + 1;
-  int out = 0;
-  group->all_sum(&in, &out, sizeof(in), jaccl::Int32);
-  int expected = size * (size + 1) / 2;
-  if (out != expected) {
-    std::cerr << "rank " << rank << ": post-barrier all_sum mismatch (got "
-              << out << ", expected " << expected << ")" << std::endl;
-    return 1;
-  }
-  std::cout << "rank " << rank << ": post-barrier all_sum OK (" << out << ")"
-            << std::endl;
-  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;
+}

mlx/distributed/jaccl/lib/jaccl/group.h

@@ -30,7 +30,6 @@ class Group {
 
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
-  virtual void barrier() = 0;
 };
 
 /**

mlx/distributed/jaccl/lib/jaccl/mesh.cpp

@@ -184,11 +184,6 @@ void MeshGroup::recv(void* output, size_t n_bytes, int src) {
   mesh_.recv(static_cast<char*>(output), n_bytes, src);
 }
 
-void MeshGroup::barrier() {
-  uint8_t b = 0;
-  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
-}
-
 template <typename T, typename ReduceOp>
 void MeshGroup::all_reduce(
     const void* input,

mlx/distributed/jaccl/lib/jaccl/mesh.h

@@ -47,8 +47,6 @@ class MeshGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
-  void barrier() override;
-
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(

mlx/distributed/jaccl/lib/jaccl/ring.cpp

@@ -190,11 +190,6 @@ void RingGroup::recv(void* output, size_t n_bytes, int src) {
   ring_.recv(static_cast<char*>(output), n_bytes, src, n_conns_);
 }
 
-void RingGroup::barrier() {
-  uint8_t b = 0;
-  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
-}
-
 template <typename T, typename ReduceOp>
 void RingGroup::all_reduce(
     const void* input,

mlx/distributed/jaccl/lib/jaccl/ring.h

@@ -48,8 +48,6 @@ class RingGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
-  void barrier() override;
-
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(

mlx/io/CMakeLists.txt

@@ -17,11 +17,6 @@ if(MLX_BUILD_GGUF)
                              PRIVATE $<BUILD_INTERFACE:${gguflib_SOURCE_DIR}>)
   add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c
                              ${gguflib_SOURCE_DIR}/gguflib.c)
-  # gguflib uses assert() to reject malformed tensor headers (e.g. ndim > 8).
-  # Those checks are otherwise compiled out by -DNDEBUG in release builds, which
-  # leaves out-of-bounds reads/writes unguarded when loading untrusted GGUF
-  # files. Force NDEBUG off for this target so the asserts stay live.
-  target_compile_options(gguflib PRIVATE -UNDEBUG)
   target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:gguflib>)
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp
                              ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp)

mlx/io/safetensors.cpp

@@ -28,7 +28,6 @@ using json = nlohmann::json;
 #define ST_U32 "U32"
 #define ST_U64 "U64"
 #define ST_F8_E4M3 "F8_E4M3"
-#define ST_F8_E8M0 "F8_E8M0"
 
 // Note: Complex numbers aren't in the spec yet so this could change -
 // https://github.com/huggingface/safetensors/issues/389
@@ -98,8 +97,6 @@ Dtype dtype_from_safetensor_str(std::string_view str) {
     return complex64;
   } else if (str == ST_F8_E4M3) {
     return uint8;
-  } else if (str == ST_F8_E8M0) {
-    return uint8;
   } else {
     std::ostringstream msg;
     msg << "[safetensor] unsupported dtype" << str;

mlx/linalg.cpp

@@ -705,170 +705,4 @@ array solve_triangular(
   return matmul(a_inv, b, s);
 }
 
-void validate_det(
-    const array& a,
-    const StreamOrDevice& stream,
-    const std::string& fname) {
-  check_cpu_stream(stream, fname);
-  if (issubdtype(a.dtype(), complexfloating)) {
-    throw std::invalid_argument(fname + " Complex inputs are not supported.");
-  }
-  if (a.ndim() < 2) {
-    std::ostringstream msg;
-    msg << fname
-        << " Arrays must have >= 2 dimensions. Received array "
-           "with "
-        << a.ndim() << " dimensions.";
-    throw std::invalid_argument(msg.str());
-  }
-
-  if (a.shape(-1) != a.shape(-2)) {
-    throw std::invalid_argument(fname + " Only defined for square matrices.");
-  }
-}
-
-array det_raw_small(const array& a, StreamOrDevice s) {
-  int n = a.shape(-1);
-
-  // Empty 0x0 matrix: determinant is the empty product = 1
-  if (n == 0) {
-    Shape out_shape(a.shape().begin(), a.shape().end() - 2);
-    return broadcast_to(array(1.0f, a.dtype()), std::move(out_shape), s);
-  }
-
-  // Helper to extract a[..., i, j] from the last two dims
-  auto elem = [&](int i, int j) {
-    auto starts = Shape(a.ndim(), 0);
-    auto stops = a.shape();
-    starts[a.ndim() - 2] = i;
-    stops[a.ndim() - 2] = i + 1;
-    starts[a.ndim() - 1] = j;
-    stops[a.ndim() - 1] = j + 1;
-    return squeeze(squeeze(slice(a, starts, stops, s), -1, s), -1, s);
-  };
-
-  if (n == 1) {
-    return elem(0, 0);
-  } else if (n == 2) {
-    return subtract(
-        multiply(elem(0, 0), elem(1, 1), s),
-        multiply(elem(0, 1), elem(1, 0), s),
-        s);
-  } else {
-    // 3x3: a00*(a11*a22 - a12*a21) - a01*(a10*a22 - a12*a20) + a02*(a10*a21 -
-    // a11*a20)
-    auto a00 = elem(0, 0), a01 = elem(0, 1), a02 = elem(0, 2);
-    auto a10 = elem(1, 0), a11 = elem(1, 1), a12 = elem(1, 2);
-    auto a20 = elem(2, 0), a21 = elem(2, 1), a22 = elem(2, 2);
-    return add(
-        subtract(
-            multiply(
-                a00,
-                subtract(multiply(a11, a22, s), multiply(a12, a21, s), s),
-                s),
-            multiply(
-                a01,
-                subtract(multiply(a10, a22, s), multiply(a12, a20, s), s),
-                s),
-            s),
-        multiply(
-            a02, subtract(multiply(a10, a21, s), multiply(a11, a20, s), s), s),
-        s);
-  }
-}
-
-std::pair<array, array> slogdet_impl(const array& input, StreamOrDevice s) {
-  int n = input.shape(-1);
-  auto dtype = input.dtype();
-
-  // Small-matrix fast path
-  if (n <= 3) {
-    auto raw = det_raw_small(input, s);
-    auto abs_raw = abs(raw, s);
-    auto sgn = sign(raw, s);
-    auto logabs = log(abs_raw, s);
-    return std::make_pair(sgn, logabs);
-  }
-
-  // General LU-based path
-  auto [LU, pivots] = lu_factor(input, s);
-
-  // Extract diagonal of U
-  auto diag = diagonal(LU, 0, -2, -1, s);
-
-  // Permutation parity: count positions where pivot[i] != i
-  int k = std::min(input.shape(-2), input.shape(-1));
-  auto iota = arange(0, k, uint32, s);
-  auto parity = astype(
-      sum(not_equal(pivots, iota, s),
-          /* axis = */ -1,
-          /* keepdims = */ false,
-          s),
-      int32,
-      s);
-
-  // Count negative diagonal elements
-  auto num_neg = astype(
-      sum(less(diag, array(0.0f, dtype), s),
-          /* axis = */ -1,
-          /* keepdims = */ false,
-          s),
-      int32,
-      s);
-
-  // sign = (-1)^(parity + num_neg)
-  auto total = add(parity, num_neg, s);
-  auto sign_val = astype(
-      subtract(
-          array(1, int32),
-          multiply(array(2, int32), remainder(total, array(2, int32), s), s),
-          s),
-      dtype,
-      s);
-
-  // logabsdet = sum(log(abs(diag)))
-  auto logabsdet =
-      sum(log(abs(diag, s), s), /* axis = */ -1, /* keepdims = */ false, s);
-
-  // Handle singular matrices: any zero on diagonal
-  auto is_zero =
-      any(equal(diag, array(0.0f, dtype), s),
-          /* axis = */ -1,
-          /* keepdims = */ false,
-          s);
-  sign_val = where(is_zero, array(0.0f, dtype), sign_val, s);
-  logabsdet = where(
-      is_zero,
-      array(-std::numeric_limits<float>::infinity(), dtype),
-      logabsdet,
-      s);
-
-  return std::make_pair(sign_val, logabsdet);
-}
-
-std::pair<array, array> slogdet(const array& a, StreamOrDevice s /* = {} */) {
-  validate_det(a, s, "[linalg::slogdet]");
-
-  auto dtype = at_least_float(a.dtype());
-  auto input = astype(a, dtype, s);
-  return slogdet_impl(input, s);
-}
-
-array det(const array& a, StreamOrDevice s /* = {} */) {
-  validate_det(a, s, "[linalg::det]");
-
-  auto dtype = at_least_float(a.dtype());
-  auto input = astype(a, dtype, s);
-  int n = input.shape(-1);
-
-  // Small-matrix fast path: compute directly, skip log/exp round-trip
-  if (n <= 3) {
-    return det_raw_small(input, s);
-  }
-
-  // General case: det = sign * exp(logabsdet)
-  auto [sign_val, logabsdet] = slogdet_impl(input, s);
-  return multiply(sign_val, exp(logabsdet, s), s);
-}
-
-} // namespace mlx::core::linalg
+} // namespace mlx::core::linalg

mlx/linalg.h

@@ -112,8 +112,4 @@ eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 MLX_API std::pair<array, array>
 eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 
-MLX_API array det(const array& a, StreamOrDevice s = {});
-
-MLX_API std::pair<array, array> slogdet(const array& a, StreamOrDevice s = {});
-
 } // namespace mlx::core::linalg

mlx/version.h

@@ -5,8 +5,8 @@
 #include "mlx/api.h"
 
 #define MLX_VERSION_MAJOR 0
-#define MLX_VERSION_MINOR 32
-#define MLX_VERSION_PATCH 0
+#define MLX_VERSION_MINOR 31
+#define MLX_VERSION_PATCH 2
 #define MLX_VERSION_NUMERIC \
   (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
 

python/src/convert.cpp

@@ -1,8 +1,5 @@
 // Copyright © 2024 Apple Inc.
 
-#include <limits>
-#include <sstream>
-
 #include <nanobind/stl/complex.h>
 
 #include "python/src/convert.h"
@@ -18,15 +15,9 @@ enum PyScalarT {
 };
 
 int check_shape_dim(int64_t dim) {
-  if (dim > std::numeric_limits<int>::max() ||
-      dim < std::numeric_limits<int>::min()) {
-    std::ostringstream msg;
-    msg << "Shape dimension " << dim << " is outside the supported range ["
-        << std::numeric_limits<int>::min() << ", "
-        << std::numeric_limits<int>::max()
-        << "]. MLX currently uses 32-bit integers for shape dimensions.";
-    PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
-    nb::detail::raise_python_error();
+  if (dim > std::numeric_limits<int>::max()) {
+    throw std::invalid_argument(
+        "Shape dimension falls outside supported `int` range.");
   }
   return static_cast<int>(dim);
 }

python/src/convert.h

@@ -76,7 +76,3 @@ nb::object tolist(mx::array& a);
 mx::array create_array(nb::object v, std::optional<mx::Dtype> t);
 mx::array array_from_list(nb::list pl, std::optional<mx::Dtype> dtype);
 mx::array array_from_list(nb::tuple pl, std::optional<mx::Dtype> dtype);
-
-// Narrow a Python-side shape dimension (int64) to a C++ mx::ShapeElem (int32),
-// raising a clear error if the value would overflow.
-int check_shape_dim(int64_t dim);

python/src/linalg.cpp

@@ -660,77 +660,4 @@ void init_linalg(nb::module_& parent_module) {
         Returns:
             array: The unique solution to the system ``AX = B``.
       )pbdoc");
-
-  m.def(
-      "det",
-      &mx::linalg::det,
-      "a"_a,
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def det(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        Compute the determinant of a square matrix.
-
-        This function supports arrays with at least 2 dimensions. When the
-        input has more than two dimensions, the determinant is computed for
-        each matrix in the last two dimensions.
-
-        Args:
-            a (array): Input array.
-            stream (Stream, optional): Stream or device. Defaults to ``None``
-              in which case the default stream of the default device is used.
-
-        Returns:
-            array: The determinant(s) of the input matrix (matrices).
-
-        Example:
-            >>> A = mx.array([[1., 2.], [3., 4.]])
-            >>> mx.linalg.det(A, stream=mx.cpu)
-            array(-2, dtype=float32)
-      )pbdoc");
-
-  m.def(
-      "slogdet",
-      [](const mx::array& a, mx::StreamOrDevice s) {
-        auto result = mx::linalg::slogdet(a, s);
-        return nb::make_tuple(result.first, result.second);
-      },
-      "a"_a,
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def slogdet(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
-      R"pbdoc(
-        Compute the sign and natural log of the absolute value of the
-        determinant of a square matrix.
-
-        This function supports arrays with at least 2 dimensions. When the
-        input has more than two dimensions, the sign and log-absolute-determinant
-        are computed for each matrix in the last two dimensions.
-
-        For a singular matrix, ``sign`` is 0 and ``logabsdet`` is ``-inf``.
-
-        The determinant can be reconstructed as ``det = sign * exp(logabsdet)``.
-        This is more numerically stable than computing the determinant directly
-        for matrices with large or small determinants.
-
-        Args:
-            a (array): Input array.
-            stream (Stream, optional): Stream or device. Defaults to ``None``
-              in which case the default stream of the default device is used.
-
-        Returns:
-            tuple(array, array): The ``sign`` and ``logabsdet`` of the
-              determinant. ``sign`` is -1, 0, or +1. ``logabsdet`` is the
-              natural log of the absolute value of the determinant.
-
-        Example:
-            >>> A = mx.array([[1., 2.], [3., 4.]])
-            >>> sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-            >>> sign
-            array(-1, dtype=float32)
-            >>> logabsdet
-            array(0.693147, dtype=float32)
-      )pbdoc");
 }

python/src/ops.cpp

@@ -15,7 +15,6 @@
 #include "mlx/einsum.h"
 #include "mlx/ops.h"
 #include "mlx/utils.h"
-#include "python/src/convert.h"
 #include "python/src/load.h"
 #include "python/src/small_vector.h"
 #include "python/src/utils.h"
@@ -46,13 +45,6 @@ double scalar_to_double(Scalar s) {
   }
 }
 
-mx::Shape to_shape(const nb::object& shape) {
-  if (nb::isinstance<nb::int_>(shape)) {
-    return {check_shape_dim(nb::cast<int64_t>(shape))};
-  }
-  return nb::cast<mx::Shape>(shape);
-}
-
 void init_ops(nb::module_& m) {
   m.def(
       "reshape",
@@ -1710,11 +1702,15 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "full",
-      [](const nb::object& shape,
+      [](const std::variant<int, mx::Shape>& shape,
          const ScalarOrArray& vals,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
-        return mx::full(to_shape(shape), to_array(vals, dtype), s);
+        if (auto pv = std::get_if<int>(&shape); pv) {
+          return mx::full({*pv}, to_array(vals, dtype), s);
+        } else {
+          return mx::full(std::get<mx::Shape>(shape), to_array(vals, dtype), s);
+        }
       },
       "shape"_a,
       "vals"_a,
@@ -1740,11 +1736,15 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "zeros",
-      [](const nb::object& shape,
+      [](const std::variant<int, mx::Shape>& shape,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
         auto t = dtype.value_or(mx::float32);
-        return mx::zeros(to_shape(shape), t, s);
+        if (auto pv = std::get_if<int>(&shape); pv) {
+          return mx::zeros({*pv}, t, s);
+        } else {
+          return mx::zeros(std::get<mx::Shape>(shape), t, s);
+        }
       },
       "shape"_a,
       "dtype"_a.none() = mx::float32,
@@ -1802,11 +1802,15 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "ones",
-      [](const nb::object& shape,
+      [](const std::variant<int, mx::Shape>& shape,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
         auto t = dtype.value_or(mx::float32);
-        return mx::ones(to_shape(shape), t, s);
+        if (auto pv = std::get_if<int>(&shape); pv) {
+          return mx::ones({*pv}, t, s);
+        } else {
+          return mx::ones(std::get<mx::Shape>(shape), t, s);
+        }
       },
       "shape"_a,
       "dtype"_a.none() = mx::float32,

python/src/small_vector.h

@@ -2,11 +2,6 @@
 
 #pragma once
 
-#include <cstdint>
-#include <limits>
-#include <sstream>
-#include <type_traits>
-
 #include "mlx/small_vector.h"
 
 #include <nanobind/stl/detail/nb_list.h>
@@ -19,19 +14,11 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
   using List = mlx::core::SmallVector<Type, Size, Alloc>;
   using Caster = make_caster<Type>;
 
-  // For narrow integer element types we fetch each element through a wider
-  // integer caster so we can emit a clean OverflowError on overflow instead of
-  // nanobind's generic "incompatible function arguments" TypeError.
-  static constexpr bool kNarrowInt = std::is_integral_v<Type> &&
-      !std::is_same_v<Type, bool> && (sizeof(Type) < sizeof(int64_t));
-
   NB_TYPE_CASTER(
       List,
       const_name("tuple[") + make_caster<Type>::Name + const_name(", ...]"))
 
-  // Not noexcept: on overflow of a narrow integer element we raise
-  // OverflowError so nanobind surfaces a clean error to the user.
-  bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) {
+  bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
     size_t size;
     PyObject* temp;
 
@@ -42,39 +29,19 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
     value.clear();
     value.reserve(size);
 
+    Caster caster;
     bool success = o != nullptr;
 
     flags = flags_for_local_caster<Type>(flags);
 
     for (size_t i = 0; i < size; ++i) {
-      if constexpr (kNarrowInt) {
-        make_caster<int64_t> wide;
-        if (!wide.from_python(o[i], flags, cleanup) ||
-            !wide.template can_cast<int64_t>()) {
-          success = false;
-          break;
-        }
-        int64_t v = wide.operator cast_t<int64_t>();
-        if (v > std::numeric_limits<Type>::max() ||
-            v < std::numeric_limits<Type>::min()) {
-          std::ostringstream msg;
-          msg << "Integer value " << v << " is outside the supported range ["
-              << static_cast<int64_t>(std::numeric_limits<Type>::min()) << ", "
-              << static_cast<int64_t>(std::numeric_limits<Type>::max()) << "].";
-          Py_XDECREF(temp);
-          PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
-          raise_python_error();
-        }
-        value.push_back(static_cast<Type>(v));
-      } else {
-        Caster caster;
-        if (!caster.from_python(o[i], flags, cleanup) ||
-            !caster.template can_cast<Type>()) {
-          success = false;
-          break;
-        }
-        value.push_back(caster.operator cast_t<Type>());
+      if (!caster.from_python(o[i], flags, cleanup) ||
+          !caster.template can_cast<Type>()) {
+        success = false;
+        break;
       }
+
+      value.push_back(caster.operator cast_t<Type>());
     }
 
     Py_XDECREF(temp);

python/tests/cuda_skip.py

@@ -1,5 +1,22 @@
 cuda_skip = {
+    # Lapack ops NYI
+    "TestLinalg.test_cholesky",
+    "TestLinalg.test_cholesky_inv",
+    "TestLinalg.test_eig",
+    "TestLinalg.test_eigh",
+    "TestLinalg.test_inverse",
+    "TestVmap.test_vmap_inverse",
+    "TestLinalg.test_lu",
+    "TestLinalg.test_lu_factor",
+    "TestLinalg.test_pseudo_inverse",
+    "TestLinalg.test_qr_factorization",
+    "TestInit.test_orthogonal",
+    "TestLinalg.test_svd_decomposition",
+    "TestVmap.test_vmap_svd",
+    "TestLinalg.test_tri_inverse",
     # Quantization NYI
     "TestQuantized.test_gather_matmul_grad",
     "TestQuantized.test_gather_qmm",
+    "TestQuantized.test_gather_qmm_sorted",
+    "TestQuantized.test_gather_qmm_grad",
 }

python/tests/test_array.py

@@ -767,13 +767,10 @@ class TestArray(mlx_tests.MLXTestCase):
 
     def test_array_np_shape_dim_check(self):
         a_npy = np.empty(2**31, dtype=np.bool_)
-        with self.assertRaises(OverflowError) as e:
+        with self.assertRaises(ValueError) as e:
             mx.array(a_npy)
         self.assertEqual(
-            str(e.exception),
-            "Shape dimension 2147483648 is outside the supported range "
-            "[-2147483648, 2147483647]. MLX currently uses 32-bit integers "
-            "for shape dimensions.",
+            str(e.exception), "Shape dimension falls outside supported `int` range."
         )
 
     def test_dtype_promotion(self):
@@ -1550,30 +1547,6 @@ class TestArray(mlx_tests.MLXTestCase):
         a = a.at[:, :, :].add(update)
         self.assertEqualArray(a, update)
 
-    def test_slice_update_contiguous_2d(self):
-        for shape in [(32, 32), (64, 64), (17, 33), (128, 128)]:
-            for upd_shape in [(16, 16), (8, 8), (4, 4)]:
-                if upd_shape[0] > shape[0] or upd_shape[1] > shape[1]:
-                    continue
-                y = mx.zeros(shape)
-                x = mx.random.normal(upd_shape)
-                z = y.at[: upd_shape[0], : upd_shape[1]].add(x)
-                diff = z[: upd_shape[0], : upd_shape[1]] - x
-                self.assertTrue(mx.allclose(diff, mx.zeros_like(diff)))
-
-        # Test non-zero offset slice update
-        y = mx.zeros((32, 32))
-        x = mx.random.normal((16, 16))
-        z = y.at[16:, 16:].add(x)
-        self.assertTrue(mx.allclose(z[16:, 16:] - x, mx.zeros_like(x)))
-
-        # Test with size divisible by 4, 2, and odd
-        for cols in [32, 18, 15]:
-            y = mx.zeros((32, cols))
-            x = mx.random.normal((16, cols))
-            z = y.at[:16, :].add(x)
-            self.assertTrue(mx.allclose(z[:16, :] - x, mx.zeros_like(x)))
-
     def test_slice_negative_step(self):
         a_np = np.arange(20)
         a_mx = mx.array(a_np)

python/tests/test_linalg.py

@@ -520,19 +520,6 @@ class TestLinalg(mlx_tests.MLXTestCase):
         P, L, U = mx.linalg.lu(a, stream=mx.cpu)
         self.assertTrue(mx.allclose(L[P, :] @ U, a))
 
-        # Test singular matrix (should not throw)
-        a = mx.array(
-            [
-                [1.0, 2.0, 3.0, 4.0],
-                [2.0, 4.0, 6.0, 8.0],
-                [0.0, 1.0, 1.0, 0.0],
-                [1.0, 0.0, 0.0, 1.0],
-            ]
-        )
-        P, L, U = mx.linalg.lu(a, stream=mx.cpu)
-        L_permuted = mx.take_along_axis(L, P[..., None], axis=-2)
-        self.assertTrue(mx.allclose(L_permuted @ U, a))
-
     def test_lu_factor(self):
         mx.random.seed(7)
 
@@ -629,248 +616,6 @@ class TestLinalg(mlx_tests.MLXTestCase):
         expected = np.linalg.solve(a, b)
         self.assertTrue(np.allclose(result, expected))
 
-    def test_det(self):
-        # 1x1 fast path
-        A = mx.array([[5.0]])
-        self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 5.0))
-
-        # 2x2 fast path
-        A = mx.array([[1.0, 2.0], [3.0, 4.0]])
-        d = mx.linalg.det(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(d, -2.0))
-
-        # 3x3 fast path
-        A = mx.array([[1.0, 2.0, 3.0], [0.0, 1.0, 4.0], [5.0, 6.0, 0.0]])
-        d = mx.linalg.det(A, stream=mx.cpu)
-        expected = np.linalg.det(np.array(A))
-        self.assertTrue(np.allclose(d, expected, atol=1e-5))
-
-        # 4x4 LU path: compare with numpy
-        np.random.seed(42)
-        A_np = np.random.randn(4, 4).astype(np.float32)
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
-
-        # 5x5 LU path
-        A_np = np.random.randn(5, 5).astype(np.float32)
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
-
-        # Identity matrix
-        A = mx.eye(5)
-        self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 1.0))
-
-        # Batched: (3, 4, 4)
-        A_np = np.random.randn(3, 4, 4).astype(np.float32)
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
-
-        # Multi-batch: (2, 3, 3, 3)
-        A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
-
-        # Integer input auto-promotes to float
-        A = mx.array([[1, 2], [3, 4]])
-        d = mx.linalg.det(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(d, -2.0))
-
-        # float64
-        A_np = np.random.randn(4, 4).astype(np.float64)
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-10))
-
-        # Singular 4x4 matrix (LU path): det should be 0
-        A = mx.array(
-            [
-                [1.0, 2.0, 3.0, 4.0],
-                [2.0, 4.0, 6.0, 8.0],
-                [0.0, 1.0, 1.0, 0.0],
-                [1.0, 0.0, 0.0, 1.0],
-            ]
-        )
-        d = mx.linalg.det(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
-
-        # Singular 5x5 matrix (LU path)
-        A_np = np.ones((5, 5), dtype=np.float32)
-        A_mx = mx.array(A_np)
-        d = mx.linalg.det(A_mx, stream=mx.cpu)
-        self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
-
-        # Batched singular matrices (LU path)
-        A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
-        A_mx = mx.array(A_np)
-        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        d_np = np.linalg.det(A_np)
-        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-5))
-
-        # Empty 0x0 matrix: det is the empty product = 1
-        d = mx.linalg.det(mx.zeros((0, 0)), stream=mx.cpu)
-        self.assertEqual(d.shape, ())
-        self.assertEqual(float(d), 1.0)
-
-        # Batched empty matrices: shape preserves batch dims
-        d = mx.linalg.det(mx.zeros((3, 0, 0)), stream=mx.cpu)
-        self.assertTrue(np.allclose(d, np.linalg.det(np.zeros((3, 0, 0)))))
-
-        # Error: non-square
-        with self.assertRaises(ValueError):
-            mx.linalg.det(mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu)
-
-        # Error: 1D
-        with self.assertRaises(ValueError):
-            mx.linalg.det(mx.array([1.0, 2.0]), stream=mx.cpu)
-
-        # Error: complex unsupported (small-matrix path)
-        with self.assertRaises(ValueError):
-            mx.linalg.det(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
-
-        # Error: complex unsupported (LU path)
-        with self.assertRaises(ValueError):
-            mx.linalg.det(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
-
-    def test_slogdet(self):
-        # 2x2: det = -2 => sign = -1, logabsdet = log(2)
-        A = mx.array([[1.0, 2.0], [3.0, 4.0]])
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(sign, -1.0))
-        self.assertTrue(np.allclose(logabsdet, np.log(2.0), atol=1e-5))
-
-        # Identity: sign = 1, logabsdet = 0
-        A = mx.eye(4)
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(sign, 1.0))
-        self.assertTrue(np.allclose(logabsdet, 0.0, atol=1e-6))
-
-        # Compare with numpy for random matrices
-        np.random.seed(42)
-        for n in [1, 2, 3, 4, 5]:
-            A_np = np.random.randn(n, n).astype(np.float32)
-            A_mx = mx.array(A_np)
-            sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-            sign_np, logabs_np = np.linalg.slogdet(A_np)
-            with self.subTest(n=n):
-                self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
-                self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
-
-        # Singular matrix 2x2 (fast path): sign = 0, logabsdet = -inf
-        A = mx.array([[1.0, 2.0], [2.0, 4.0]])
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertEqual(float(sign), 0.0)
-        self.assertEqual(float(logabsdet), float("-inf"))
-
-        # Singular 4x4 matrix (LU path): sign = 0, logabsdet = -inf
-        A = mx.array(
-            [
-                [1.0, 2.0, 3.0, 4.0],
-                [2.0, 4.0, 6.0, 8.0],
-                [0.0, 1.0, 1.0, 0.0],
-                [1.0, 0.0, 0.0, 1.0],
-            ]
-        )
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertEqual(float(sign), 0.0)
-        self.assertEqual(float(logabsdet), float("-inf"))
-
-        # Singular 5x5 matrix (LU path): all-ones matrix
-        A = mx.array(np.ones((5, 5), dtype=np.float32))
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertEqual(float(sign), 0.0)
-        self.assertEqual(float(logabsdet), float("-inf"))
-
-        # Batched with mix of singular and non-singular (LU path)
-        A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
-        A_mx = mx.array(A_np)
-        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-        sign_np, logabs_np = np.linalg.slogdet(A_np)
-        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
-        # Check -inf for singular, 0.0 for identity
-        self.assertEqual(float(logabs_mx[0]), float("-inf"))
-        self.assertTrue(np.allclose(logabs_mx[1], 0.0, atol=1e-6))
-
-        # Batched
-        A_np = np.random.randn(3, 4, 4).astype(np.float32)
-        A_mx = mx.array(A_np)
-        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-        sign_np, logabs_np = np.linalg.slogdet(A_np)
-        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
-        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
-
-        # Multi-batch
-        A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
-        A_mx = mx.array(A_np)
-        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-        sign_np, logabs_np = np.linalg.slogdet(A_np)
-        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
-        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
-
-        # Numerical stability: large matrix where det overflows
-        # 0.1 * I_100 has det = 0.1^100 which underflows in float32
-        # but slogdet should give sign=1, logabsdet = 100*log(0.1)
-        n = 100
-        A = mx.array(0.1) * mx.eye(n)
-        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
-        self.assertTrue(np.allclose(sign, 1.0))
-        self.assertTrue(np.allclose(logabsdet, n * np.log(0.1), atol=1e-3))
-
-        # Verify det = sign * exp(logabsdet) for non-singular cases
-        A_np = np.random.randn(5, 5).astype(np.float32)
-        A_mx = mx.array(A_np)
-        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-        det_mx = mx.linalg.det(A_mx, stream=mx.cpu)
-        reconstructed = float(sign_mx) * np.exp(float(logabs_mx))
-        self.assertTrue(np.allclose(float(det_mx), reconstructed, rtol=1e-4))
-
-        # float64
-        A_np = np.random.randn(4, 4).astype(np.float64)
-        A_mx = mx.array(A_np)
-        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
-        sign_np, logabs_np = np.linalg.slogdet(A_np)
-        self.assertTrue(np.allclose(sign_mx, sign_np))
-        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-10))
-
-        # Empty 0x0 matrix: sign = 1, logabsdet = 0 (empty product)
-        sign, logabsdet = mx.linalg.slogdet(mx.zeros((0, 0)), stream=mx.cpu)
-        self.assertEqual(sign.shape, ())
-        self.assertEqual(logabsdet.shape, ())
-        self.assertEqual(float(sign), 1.0)
-        self.assertEqual(float(logabsdet), 0.0)
-
-        # Batched empty matrices
-        sign, logabsdet = mx.linalg.slogdet(mx.zeros((3, 0, 0)), stream=mx.cpu)
-        sign_np, logabs_np = np.linalg.slogdet(np.zeros((3, 0, 0)))
-        self.assertTrue(np.allclose(sign, sign_np))
-        self.assertTrue(np.allclose(logabsdet, logabs_np))
-
-        # Error: non-square
-        with self.assertRaises(ValueError):
-            mx.linalg.slogdet(
-                mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu
-            )
-
-        # Error: 1D
-        with self.assertRaises(ValueError):
-            mx.linalg.slogdet(mx.array([1.0, 2.0]), stream=mx.cpu)
-
-        # Error: complex unsupported (small-matrix path)
-        with self.assertRaises(ValueError):
-            mx.linalg.slogdet(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
-
-        # Error: complex unsupported (LU path)
-        with self.assertRaises(ValueError):
-            mx.linalg.slogdet(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
-
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

python/tests/test_ops.py

@@ -92,52 +92,6 @@ class TestOps(mlx_tests.MLXTestCase):
             self.assertEqual(y.dtype, t)
             self.assertTrue(mx.array_equal(y, x))
 
-    def test_shape_overflow_error(self):
-        # Shape dimensions that don't fit in int32 should raise a clear
-        # OverflowError that names the offending value, rather than a generic
-        # "incompatible function arguments" TypeError. The overflow check
-        # lives in the mx::Shape type caster, so it applies to every op that
-        # takes a shape. See issue #2681.
-        too_big = 2**31
-
-        # Array creation ops — also exercise the scalar shape path.
-        for ctor in (mx.zeros, mx.ones):
-            with self.assertRaises(OverflowError) as cm:
-                ctor(too_big)
-            self.assertIn(str(too_big), str(cm.exception))
-            with self.assertRaises(OverflowError) as cm:
-                ctor([too_big])
-            self.assertIn(str(too_big), str(cm.exception))
-
-        with self.assertRaises(OverflowError) as cm:
-            mx.full(too_big, 0.0)
-        self.assertIn(str(too_big), str(cm.exception))
-        with self.assertRaises(OverflowError) as cm:
-            mx.full([too_big], 0.0)
-        self.assertIn(str(too_big), str(cm.exception))
-
-        # Other shape-taking ops should surface the same clean error.
-        a = mx.zeros(4)
-        with self.assertRaises(OverflowError) as cm:
-            mx.reshape(a, [too_big])
-        self.assertIn(str(too_big), str(cm.exception))
-
-        with self.assertRaises(OverflowError) as cm:
-            mx.broadcast_to(a, [too_big, 1])
-        self.assertIn(str(too_big), str(cm.exception))
-
-        # Negative overflow (< int32 min) is caught too.
-        too_negative = -(2**31) - 1
-        with self.assertRaises(OverflowError) as cm:
-            mx.zeros([too_negative])
-        self.assertIn(str(too_negative), str(cm.exception))
-
-        # Shapes that fit in int32 still go through unchanged.
-        self.assertEqual(mx.zeros(4).shape, (4,))
-        self.assertEqual(mx.zeros((2, 3)).shape, (2, 3))
-        self.assertEqual(mx.ones([2, 3]).shape, (2, 3))
-        self.assertEqual(mx.full((2, 3), 1.5).tolist(), [[1.5] * 3] * 2)
-
     def test_scalar_inputs(self):
         # Check combinations of python types
         a = mx.add(False, True)

python/tests/test_quantized.py

@@ -1046,7 +1046,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
                 y3 = scatter_unsort(y3, inv_order, indices.shape)
                 y4 = scatter_unsort(y4, inv_order, indices.shape)
 
-                tol = 1.5e-5 if (dtype == mx.float32) else 1e-3
+                tol = 1.5e-5 if (dtype == mx.float32) else 2.5e-4
 
                 self.assertLess((y1 - y2).abs().max(), tol)
                 self.assertLess((y1 - y3).abs().max(), tol)

tests/linalg_tests.cpp

@@ -637,68 +637,3 @@ TEST_CASE("test solve_triangluar") {
   expected = array({-3., 2., 3.});
   CHECK(allclose(expected, result).item<bool>());
 }
-
-TEST_CASE("test det") {
-  // 1x1 fast path
-  {
-    array a = array({5.0f}, {1, 1});
-    auto d = det(a, Device::cpu);
-    CHECK_EQ(d.item<float>(), doctest::Approx(5.0f));
-  }
-
-  // 2x2 fast path: det([[1,2],[3,4]]) = -2
-  {
-    array a = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
-    auto d = det(a, Device::cpu);
-    CHECK_EQ(d.item<float>(), doctest::Approx(-2.0f));
-  }
-
-  // 3x3 fast path: det([[1,2,3],[0,1,4],[5,6,0]]) = 1
-  {
-    array a =
-        array({1.0f, 2.0f, 3.0f, 0.0f, 1.0f, 4.0f, 5.0f, 6.0f, 0.0f}, {3, 3});
-    auto d = det(a, Device::cpu);
-    CHECK_EQ(d.item<float>(), doctest::Approx(1.0f));
-  }
-
-  // 4x4 LU path: identity matrix det = 1
-  {
-    array a = eye(4);
-    auto d = det(a, Device::cpu);
-    CHECK_EQ(d.item<float>(), doctest::Approx(1.0f));
-  }
-
-  // Non-square should throw
-  CHECK_THROWS(
-      det(array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}, {2, 3}), Device::cpu));
-
-  // 1D should throw
-  CHECK_THROWS(det(array({1.0f, 2.0f}), Device::cpu));
-}
-
-TEST_CASE("test slogdet") {
-  // 2x2: det = -2, so sign = -1, logabsdet = log(2)
-  {
-    array a = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
-    auto [s, logabs] = slogdet(a, Device::cpu);
-    CHECK_EQ(s.item<float>(), doctest::Approx(-1.0f));
-    CHECK_EQ(logabs.item<float>(), doctest::Approx(std::log(2.0f)));
-  }
-
-  // Identity: sign = 1, logabsdet = 0
-  {
-    array a = eye(4);
-    auto [s, logabs] = slogdet(a, Device::cpu);
-    CHECK_EQ(s.item<float>(), doctest::Approx(1.0f));
-    CHECK_EQ(logabs.item<float>(), doctest::Approx(0.0f));
-  }
-
-  // Singular: sign = 0, logabsdet = -inf
-  {
-    array a = array({1.0f, 2.0f, 2.0f, 4.0f}, {2, 2});
-    auto [s, logabs] = slogdet(a, Device::cpu);
-    CHECK_EQ(s.item<float>(), 0.0f);
-    CHECK(std::isinf(logabs.item<float>()));
-    CHECK(logabs.item<float>() < 0);
-  }
-}

tests/ops_tests.cpp

@@ -2737,41 +2737,11 @@ TEST_CASE("test as_strided op") {
   auto x = arange(10);
   auto y = as_strided(x, {3, 3}, {1, 1}, 0);
   auto expected = array({0, 1, 2, 1, 2, 3, 2, 3, 4}, {3, 3});
-  eval(y);
   CHECK(array_equal(y, expected).item<bool>());
-  CHECK_EQ(y.data_size(), 5);
-  CHECK_FALSE(y.flags().contiguous);
 
   y = as_strided(x, {3, 3}, {0, 3}, 0);
   expected = array({0, 3, 6, 0, 3, 6, 0, 3, 6}, {3, 3});
-  eval(y);
   CHECK(array_equal(y, expected).item<bool>());
-  CHECK_EQ(y.data_size(), 7);
-  CHECK_FALSE(y.flags().contiguous);
-
-  x = arange(24);
-  y = as_strided(x, {2, 3, 4}, {3, 1, 6}, 0);
-  expected = array(
-      {0, 6, 12, 18, 1, 7,  13, 19, 2, 8,  14, 20,
-       3, 9, 15, 21, 4, 10, 16, 22, 5, 11, 17, 23},
-      {2, 3, 4});
-  eval(y);
-  CHECK(array_equal(y, expected).item<bool>());
-  CHECK_EQ(y.data_size(), 24);
-  CHECK(y.flags().contiguous);
-  CHECK_FALSE(y.flags().row_contiguous);
-  CHECK_FALSE(y.flags().col_contiguous);
-
-  auto z = astype(y, float32);
-  CHECK(array_equal(z, astype(expected, float32)).item<bool>());
-
-  x = arange(10);
-  y = as_strided(x, {10}, {-1}, 9);
-  expected = array({9, 8, 7, 6, 5, 4, 3, 2, 1, 0}, {10});
-  eval(y);
-  CHECK(array_equal(y, expected).item<bool>());
-  CHECK_EQ(y.data_size(), 10);
-  CHECK_FALSE(y.flags().contiguous);
 
   x = reshape(x, {2, 5}); // 0 1 2 3 ...
   x = transpose(x, {1, 0}); // 0 5 1 6 2 7 ...