fix(ci): use branch_to_build input for checkout and artifact naming

ci: add macOS arm64 wheel build workflow
Builds mlx wheels for Python 3.11 and 3.12 on macos-14 (arm64) on push to main, metal-*, q-*, attn-mask-fix, fix-rope branches or via workflow_dispatch. Wheels uploaded as artifacts (30d retention). Tagged commits (v*) also publish a GitHub Release.
2026-05-10 13:41:03 +02:00 · 2026-05-10 13:13:33 +02:00 · 2026-05-09 11:11:25 +09:00 · 2026-05-09 09:02:12 +09:00 · 2026-05-08 01:29:29 -07:00
15 changed files with 1035 additions and 758 deletions
@@ -0,0 +1,94 @@
+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
@@ -19,27 +19,28 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
        "AsStrided must be used with row contiguous arrays only.");
  }

-  // 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];
+  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;
+    }
  }
+  size_t data_size = out.size() == 0 ? 0 : (h - l) + 1;
+
  auto flags = in.flags();
-  // TODO: Compute the contiguous flag in a better way cause now we are
-  //       unnecessarily strict.
-  flags.contiguous = row_contiguous || col_contiguous;
+  flags.contiguous =
+      out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size);
  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_);
 }

@@ -10,7 +10,6 @@
 #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"
@@ -316,7 +315,9 @@ void Compiled::eval_cpu(
  // Get the function
  auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
    std::ostringstream kernel;
-    kernel << get_kernel_preamble() << std::endl;
+    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 << "extern \"C\"  {" << std::endl;
    build_kernel(
        kernel,
@@ -9,4 +9,4 @@
 #include "mlx/backend/cpu/binary_ops.h"
 // clang-format on

-const char* get_kernel_preamble();
+const char* get_prebuilt_preamble();
@@ -1,6 +1,8 @@
 // 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>
@@ -86,30 +88,61 @@ 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) {
-    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
+    extra_flags += fmt::format(" /libpath:\"{}\"", lib);
  }
  return fmt::format(
      "\""
-      "cd /D \"{0}\" && "
-      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
-      "/link /out:\"{3}\" {4} 2>&1"
+      "cd /D \"{}\" && "
+      "\"{}\" /LD /EHsc /MD /Ox /nologo /std:c++17 {} \"{}\" "
+      "/link /out:\"{}\" 2>&1"
      "\"",
      dir.string(),
      info.cl_exe,
+      extra_flags,
      source_file_name,
-      shared_lib_name,
-      libpaths);
+      shared_lib_name);
 #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 \"{0}\" -o \"{1}\" 2>&1",
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1",
+      extra_flags,
      (dir / source_file_name).string(),
      (dir / shared_lib_name).string());
 #endif
@@ -7,6 +7,9 @@ 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,
@@ -15,13 +15,6 @@ $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
@@ -29,7 +22,7 @@ using namespace mlx::core::detail;
 $CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'

 $OUTPUT = @"
-const char* get_kernel_preamble() {
+const char* get_prebuilt_preamble() {
  static char preamble[] = { $CHARCODES };
  return preamble;
 }
@@ -30,12 +30,10 @@ 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_kernel_preamble() {
+const char* get_prebuilt_preamble() {
 return R"preamble(
 $INCLUDES
 $CONTENT
-using namespace mlx::core;
-using namespace mlx::core::detail;
 )preamble";
 }
 EOF
@@ -1,9 +1,8 @@
 // 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/dtype_utils.h"
+#include "mlx/backend/cuda/quantized/qmm/qmm_naive.cuh"

 // clang-format off

@@ -12,49 +11,26 @@ 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;
-};
-
-__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,
+template <bool KMajor, bool HasKResidue, bool SM80,
          typename Element, typename Quant, typename Scale,
-          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,
+          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,
    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));
@@ -64,20 +40,6 @@ __global__ void qmm_naive_kernel(
  int n_coord = int(blockIdx.y);
  int l_coord = int(blockIdx.z);

-  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)
  Tensor mB_nkl = make_tensor(make_gmem_ptr<Quant>(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L)
@@ -107,218 +69,24 @@ __global__ 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)

-  // 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)
+  // 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);

-  // 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));
-    }
-  }
+  qmm_naive_mainloop<KMajor, HasKResidue, SM80>(
+      cta_tiler,
+      gA,
+      gB,
+      gS,
+      gZ,
+      gC,
+      mma,
+      m_max_coord, n_max_coord, k_residue,
+      thread_idx);
 }

-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,
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80,
          typename Element, typename Quant, typename Scale>
 void qmm_naive(
    const Element* A,
@@ -333,14 +101,12 @@ void qmm_naive(
    auto group_size,
    auto&& launch_kernel) {
  // Define shapes (dynamic).
-  auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L)
+  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)

-  // Define TN strides (mixed).
+  // Define layouts (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.
@@ -349,45 +115,41 @@ void qmm_naive(
    get<2>(stride(S_layout)) = 0;
  }

-  // 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 CTA tile size (static).
+  auto cta_tiler = make_cta_tiler<TileM, SM80>(group_size);

  // Define MMA.
-  TiledMMA mma = make_tiled_mma<TileM, SM80, Element>();
+  auto mma = make_tiled_mma<SM80, Element>(cta_tiler);
  auto num_threads = size(mma);

-  // 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);
+  // 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)>);

  auto* kernel = &qmm_naive_kernel<
-      HasKResidue, decltype(prob_shape), decltype(cta_tiler),
+      KMajor, HasKResidue, SM80,
      Element, Quant, Scale,
-      decltype(dA), decltype(sA_layout), decltype(copy_a),
-      decltype(dB), decltype(sB_layout), decltype(copy_b),
-      decltype(dC), decltype(S_layout), decltype(mma)>;
-
-  // Set L1 to be SMEM only.
-  size_t smem_bytes = sizeof(SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>);
+      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(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l);
-  dim3 block_dims(num_threads);
+  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};
  void* args[] = {
-      &prob_shape, &cta_tiler,
-      &A, &dA, &sA_layout, &copy_a,
-      &B, &dB, &sB_layout, &copy_b,
-      &C, &dC,
+      &shape_MNKL,
+      &cta_tiler,
+      &A, &dA,
+      &B, &dB,
      &S, &Z, &S_layout,
      &lhs_indices, &rhs_indices,
+      &C, &dC,
      &mma};
  launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
 }
@@ -398,69 +160,6 @@ 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,
@@ -518,7 +217,7 @@ void qmm_naive_impl(
              [&](auto* kernel,
                  dim3 num_blocks,
                  dim3 block_dims,
-                  uint32_t smem_bytes,
+                  size_t smem_bytes,
                  void** args) {
                encoder.add_kernel_node_raw(
                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
@@ -0,0 +1,381 @@
+// 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
@@ -1,8 +1,7 @@
 // Copyright © 2026 Apple Inc.

-#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
 #include "mlx/backend/cuda/quantized/qmm/qmm.h"
-#include "mlx/dtype_utils.h"
+#include "mlx/backend/cuda/quantized/qmm/qmm_sm80.cuh"

 // clang-format off

@@ -11,38 +10,24 @@ 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;
-};
-
-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(
+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(
    ProblemShape shape_MNKL, CtaTiler cta_tiler,
-    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 Element* A, StrideA dA,
+    const Quant* B, StrideB dB,
+    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(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));
@@ -81,201 +66,23 @@ __global__ 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)

-  // 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)
-
-  // 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.
+  // Compute tile residues for predication.
  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);
+  qmm_sm80_mainloop(
+      cta_tiler,
+      gA,
+      gB,
+      gS,
+      gZ,
+      gC,
+      mma,
+      m_max_coord,
+      thread_idx);
 }

-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>
+template <int TileM,
+          typename Element, typename Quant, typename Scale>
 void qmm_sm80(
    const Element* A,
    const Quant*   B,
@@ -286,20 +93,16 @@ void qmm_sm80(
    Element* C,
    int m, int n, int k, int l,
    bool broadcast_b,
-    GroupSize group_size,
+    auto group_size,
    auto&& launch_kernel) {
  // Define shapes (dynamic).
-  auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L)
+  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)

-  // Define TN strides (mixed).
+  // Define layouts (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)
-
-  // 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));
+  auto S_layout = make_scales_layout(n, k, l, group_size);

  // Handle broadcasting.
  if (broadcast_b) {
@@ -308,70 +111,41 @@ void qmm_sm80(
  }

  // Define CTA tile sizes (static).
-  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)
+  auto cta_tiler = make_cta_tiler<TileM>(group_size);

  // Define MMA.
  TiledMMA mma = make_tiled_mma<TileM, Element>();
  auto num_threads = size(mma);

-  // 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.
+  // Shared memory size.
+  auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler);
  size_t smem_bytes = sizeof(SharedStorage<Element, Quant,
                                           decltype(sA_layout),
                                           decltype(sB_layout),
                                           decltype(sC_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);
+  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);
+
+  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};
  void* args[] = {
-      &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,
+      &shape_MNKL, &cta_tiler,
+      &A, &dA,
+      &B, &dB,
+      &S, &Z, &S_layout,
      &lhs_indices, &rhs_indices,
+      &C, &dC,
      &mma};
  launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
 }
@@ -382,59 +156,6 @@ 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,
@@ -492,7 +213,7 @@ void qmm_sm80_impl(
              [&](auto* kernel,
                  dim3 num_blocks,
                  dim3 block_dims,
-                  uint32_t smem_bytes,
+                  size_t smem_bytes,
                  void** args) {
                encoder.add_kernel_node_raw(
                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
@@ -0,0 +1,346 @@
+// 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
@@ -133,9 +133,8 @@ void RoPE::eval_gpu(
  if (single) {
    compute_encoder.set_bytes(out_strides, 1, 4);
    uint32_t dim0 = dims_ / 2;
-    uint32_t dim1 = B * N;
-    group_dims = get_block_dims(dim0, dim1, 1);
-    grid_dims = MTL::Size(dim0, dim1, 1);
+    group_dims = get_block_dims(dim0, N, 1);
+    grid_dims = MTL::Size(dim0, N, 1);
  } else {
    compute_encoder.set_bytes(strides, 3, 4);
    compute_encoder.set_bytes(out_strides, 3, 5);
@@ -365,28 +365,6 @@ class TestFast(mlx_tests.MLXTestCase):
        rx = rope_orig(x, dims, traditional, base, scale, offset)
        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)

-    def test_rope_single_batch(self):
-        base = 10000.0
-        scale = 1.0
-        offset = 5
-
-        for traditional in [True, False]:
-            for B in [2, 4, 8]:
-                for n_head in [1, 4, 7]:
-                    for dims in [64, 128]:
-                        x = mx.random.uniform(shape=(B, n_head, 1, dims))
-                        mx.eval(x)
-                        rx_fast = mx.fast.rope(
-                            x,
-                            dims,
-                            traditional=traditional,
-                            base=base,
-                            scale=scale,
-                            offset=offset,
-                        )
-                        rx = rope_orig(x, dims, traditional, base, scale, offset)
-                        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
-
    def test_rope_with_large_offset(self):
        x = mx.random.normal(shape=(1, 1, 1024, 32))
        rx_fp32 = mx.fast.rope(
@@ -2737,11 +2737,41 @@ 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 ...
Author	SHA1	Message	Date
electron-rare	73f33e92cd	fix(ci): use branch_to_build input for checkout and artifact naming	2026-05-10 13:41:03 +02:00
electron-rare	765ebd7ef2	ci: add macOS arm64 wheel build workflow Builds mlx wheels for Python 3.11 and 3.12 on macos-14 (arm64) on push to main, metal-, q-, attn-mask-fix, fix-rope branches or via workflow_dispatch. Wheels uploaded as artifacts (30d retention). Tagged commits (v*) also publish a GitHub Release.	2026-05-10 13:13:33 +02:00
Cheng	84961223c0	[CUDA] Separate main loop into a function in qmm (#3443 )	2026-05-09 11:11:25 +09:00
Cheng	662115c1f0	Do not use prebuilt cpu compile preamble when headers are installed (#3463 )	2026-05-09 09:02:12 +09:00
Valeriy Sofin	a1c0b6f9ac	Compute contiguity from the actual occupied data (#3475 )	2026-05-08 01:29:29 -07:00