[CUDA] FPxINT quantized matmul for Hopper (#3160)

mlx/backend/cuda/CMakeLists.txt

@@ -56,6 +56,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm_sm90.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
@@ -116,6 +117,10 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
+# Required for generating optimized CUTLASS code.
+target_compile_options(
+  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
+
 # Suppress nvcc warnings on C++ headers.
 target_compile_options(
   mlx

mlx/backend/cuda/cutlass_utils.cuh

@@ -22,7 +22,7 @@ inline void check_cutlass_error(const char* name, cutlass::Status status) {
 }
 
 // The macro version that prints the command that failed.
-#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd))
+#define CHECK_CUTLASS_ERROR(cmd) ::mlx::core::check_cutlass_error(#cmd, (cmd))
 
 // Maps CPU types to CUTLASS types.
 template <typename T>

mlx/backend/cuda/quantized/qmm.h

@@ -0,0 +1,24 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+
+#include <optional>
+
+namespace mlx::core {
+
+void qmm_sm90(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder,
+    Stream s);
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm_sm90.cu

@@ -0,0 +1,244 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/cutlass_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmm.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+
+#include <cute/tensor.hpp>
+#include <cutlass/cutlass.h>
+#include <cutlass/epilogue/collective/collective_builder.hpp>
+#include <cutlass/gemm/collective/collective_builder.hpp>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/gemm_universal.hpp>
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+template <typename GroupSize, typename Element, typename Quant, typename F>
+void qmm_sm90(
+    const Element* A,
+    const Quant* B,
+    const Element* S,
+    const Element* Z,
+    Element* D,
+    int64_t m,
+    int64_t n,
+    int64_t k,
+    int64_t l,
+    GroupSize group_size,
+    F&& launch_kernel) {
+#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
+  using namespace cute;
+
+  constexpr int kAlignmentA = 128 / sizeof_bits<Element>::value;
+  constexpr int kAlignmentB = 128 / sizeof_bits<Quant>::value;
+  constexpr int kTileShapeK =
+      std::max(64, 128 * 8 / sizeof_bits<Element>::value);
+  static_assert(group_size % kTileShapeK == 0);
+
+  using Arch = cutlass::arch::Sm90;
+  using Accumulator = float;
+  using TileShape = Shape<_128, _16, Int<kTileShapeK>>;
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      Arch,
+      cutlass::arch::OpClassTensorOp,
+      TileShape,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      Accumulator,
+      Accumulator,
+      // ElementC:
+      void,
+      cutlass::layout::ColumnMajor,
+      kAlignmentA,
+      // ElementD:
+      Element,
+      cutlass::layout::ColumnMajor,
+      kAlignmentA,
+      cutlass::epilogue::TmaWarpSpecializedCooperative>::CollectiveOp;
+
+  // Note that A/B are swapped and transposed to use TMA epilogue.
+  using Mainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      Arch,
+      cutlass::arch::OpClassTensorOp,
+      // ElementA:
+      cute::tuple<Quant, Element, Element>,
+      cutlass::layout::RowMajor,
+      kAlignmentB,
+      // ElementB:
+      Element,
+      cutlass::layout::ColumnMajor,
+      kAlignmentA,
+      Accumulator,
+      TileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename Epilogue::SharedStorage))>,
+      cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::
+      GemmUniversal<Shape<int, int, int, int>, Mainloop, Epilogue>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+  auto dA = make_stride(k, Int<1>{}, m * k);
+  auto dB = make_stride(k, Int<1>{}, n * k);
+  auto dS = make_stride(Int<1>{}, n, n * k / group_size);
+  auto dD = make_stride(Int<1>{}, n, m * n);
+
+  Gemm gemm;
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {int(n), int(m), int(k), int(l)},
+      {B, dB, A, dA, S, dS, group_size, Z},
+      {{1.f, 0.f}, D, dD, D, dD}};
+
+  CHECK_CUTLASS_ERROR(gemm.can_implement(args));
+  CHECK_CUTLASS_ERROR(gemm.initialize(args, nullptr));
+
+  auto* kernel = &cutlass::device_kernel<GemmKernel>;
+  void* kernel_params[] = {const_cast<Gemm::Params*>(&gemm.params())};
+  launch_kernel(
+      reinterpret_cast<void*>(kernel),
+      gemm.get_grid_shape(gemm.params()),
+      GemmKernel::get_block_shape(),
+      GemmKernel::SharedStorageSize,
+      kernel_params);
+#else
+  throw std::runtime_error(
+      "[quantized_matmul] Hopper-only kernel is not available.");
+#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
+}
+
+} // namespace cutlass_gemm
+
+namespace mlx::core {
+
+inline array transpose_last_2_dims(
+    const array& x,
+    cu::CommandEncoder& encoder,
+    const Stream& s) {
+  array transposed = swapaxes_in_eval(x, -1, -2);
+  array transposed_copy = contiguous_copy_gpu(transposed, s);
+  encoder.add_temporary(transposed_copy);
+  return transposed_copy;
+}
+
+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_quant_types(int bits, const char* tag, F&& f) {
+  if (bits == 2) {
+    f.template operator()<cutlass::uint2b_t>();
+  } else if (bits == 4) {
+    f.template operator()<cutlass::uint4b_t>();
+  } else if (bits == 8) {
+    f.template operator()<uint8_t>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} {}-bit quantization is not supported.", tag, bits));
+  }
+}
+
+template <typename F>
+inline void dispatch_groups(int group_size, const char* tag, F&& f) {
+  if (group_size == 64) {
+    f(cute::Int<64>{});
+  } else if (group_size == 128) {
+    f(cute::Int<128>{});
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Group size {} is not supported.", tag, group_size));
+  }
+}
+
+void qmm_sm90(
+    const array& x_,
+    const array& w,
+    const array& scales_,
+    const std::optional<array>& biases_,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if ((mode != QuantizationMode::Affine) || !biases_) {
+    throw std::runtime_error("qmm_sm90 NYI");
+  }
+
+  const char* tag = "[quantized_matmul]";
+  int m = out.shape(-2);
+  int n = out.shape(-1);
+  int k = x_.shape(-1);
+  int l = out.size() / (m * n);
+  if (k % 64 != 0) {
+    throw std::runtime_error(fmt::format("{} K must be multiples of 64.", tag));
+  }
+  if (!w.flags().row_contiguous) {
+    throw std::runtime_error(
+        fmt::format("{} Weights must be row contiguous.", tag));
+  }
+  if (!scales_.flags().row_contiguous) {
+    throw std::runtime_error(
+        fmt::format("{} Scales must be row contiguous.", tag));
+  }
+  if (!biases_->flags().row_contiguous) {
+    throw std::runtime_error(
+        fmt::format("{} Biases must be row contiguous.", tag));
+  }
+
+  // TODO: Support column-major x.
+  array x = ensure_row_contiguous(x_, encoder, s);
+  // FIXME: Copy happens for every call.
+  array scales = transpose_last_2_dims(scales_, encoder, s);
+  array biases = transpose_last_2_dims(*biases_, encoder, s);
+
+  dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
+    dispatch_quant_types(bits, tag, [&]<typename Quant>() {
+      dispatch_groups(group_size, tag, [&](auto group_size) {
+        encoder.set_input_array(x);
+        encoder.set_input_array(w);
+        encoder.set_input_array(scales);
+        encoder.set_input_array(biases);
+        encoder.set_output_array(out);
+        cutlass_gemm::qmm_sm90(
+            gpu_ptr<Element>(x),
+            gpu_ptr<Quant>(w),
+            gpu_ptr<Element>(scales),
+            gpu_ptr<Element>(biases),
+            gpu_ptr<Element>(out),
+            m,
+            n,
+            k,
+            l,
+            group_size,
+            [&](auto* kernel,
+                dim3 num_blocks,
+                dim3 block_dims,
+                uint32_t smem_bytes,
+                void** args) {
+              encoder.add_kernel_node(
+                  kernel, num_blocks, block_dims, smem_bytes, args);
+            });
+      });
+    });
+  });
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmv.cu

@@ -4,6 +4,7 @@
 #include "mlx/backend/cuda/kernel_utils.cuh"
 #include "mlx/backend/cuda/quantized/qmv.h"
 #include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/dtype_utils.h"
 
 #include <cooperative_groups.h>
@@ -214,16 +215,23 @@ void dispatch_1_2_4(int n, F&& f) {
 }
 
 void fp_qmv(
-    const array& mat,
-    const array& scales,
-    const array& vec,
+    const array& x,
+    const array& w,
+    const array& scales_,
     array& out,
     int bits,
     int group_size,
     int M,
     int N,
     int K,
-    CommandEncoder& encoder) {
+    CommandEncoder& encoder,
+    Stream s) {
+  // Make sure the last two dims of x and w, s, b are contiguous. This should
+  // be relaxed for x.
+  array vec = ensure_row_contiguous_matrix(x, encoder, s);
+  array mat = ensure_row_contiguous_matrix(w, encoder, s);
+  array scales = ensure_row_contiguous_matrix(scales_, encoder, s);
+
   encoder.set_input_array(mat);
   encoder.set_input_array(scales);
   encoder.set_input_array(vec);

mlx/backend/cuda/quantized/qmv.h

@@ -7,15 +7,16 @@
 namespace mlx::core::cu {
 
 void fp_qmv(
+    const array& x,
     const array& w,
     const array& scales,
-    const array& vec,
     array& out,
     int bits,
     int group_size,
     int M,
     int N,
     int K,
-    CommandEncoder& encoder);
+    CommandEncoder& encoder,
+    Stream s);
 
 } // namespace mlx::core::cu

mlx/backend/cuda/quantized/qqmm.cpp

@@ -106,16 +106,15 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     fp_quantize_dequantize(
         x, xhat, group_size_, bits_, global_scale, encoder, s);
 
-    // Make sure the last two dims of w and s are contiguous
-    array w = ensure_row_contiguous_matrix(inputs[1], encoder, s);
-    array scales = ensure_row_contiguous_matrix(inputs[2], encoder, s);
+    const array& w = inputs[1];
+    const array& scales = inputs[2];
 
     bool non_batched = w.ndim() == 2;
     int K = x.shape(-1);
     int M = non_batched ? x.size() / K : x.shape(-2);
     int N = out.shape(-1);
 
-    fp_qmv(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder);
+    fp_qmv(xhat, w, scales, out, bits_, group_size_, M, N, K, encoder, s);
     return;
   }
 

mlx/backend/cuda/quantized/quantized.cpp

@@ -2,6 +2,7 @@
 
 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/qmm.h"
 #include "mlx/backend/cuda/quantized/qmv.h"
 #include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/fast_primitives.h"
@@ -14,34 +15,35 @@ namespace mlx::core {
 void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   nvtx3::scoped_range r("QuantizedMatmul::eval_gpu");
   auto& s = stream();
-  auto& d = cu::device(s.device);
-  auto& enc = d.get_command_encoder(s);
+  auto& encoder = cu::get_command_encoder(s);
 
-  out.set_data(cu::malloc_async(out.nbytes(), enc));
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
 
-  // Make sure the last two dims of x and w, s, b are contiguous. This should
-  // be relaxed for x.
-  array x = ensure_row_contiguous_matrix(inputs[0], enc, s);
-  array w = ensure_row_contiguous_matrix(inputs[1], enc, s);
-  array scales = ensure_row_contiguous_matrix(inputs[2], enc, s);
-  std::optional<array> biases = std::nullopt;
-  if (inputs.size() == 4) {
-    biases = ensure_row_contiguous_matrix(inputs[3], enc, s);
+  const array& x = inputs[0];
+  const array& w = inputs[1];
+  const array& scales = inputs[2];
+  std::optional<array> biases;
+  if (inputs.size() > 3) {
+    biases = inputs[3];
   }
 
-  bool non_batched = w.ndim() == 2 && x.flags().row_contiguous;
+  bool non_batched = w.ndim() == 2;
   int K = x.shape(-1);
-  int M = non_batched ? x.size() / K : x.shape(-2);
   int N = out.shape(-1);
+  int vec_batch = non_batched ? x.size() / K : x.shape(-2);
 
-  if (M > 8 || !transpose_ || mode_ == QuantizationMode::Affine) {
-    throw std::runtime_error("QMM NYI");
-  }
-
-  if (transpose_) {
-    fp_qmv(w, scales, x, out, bits_, group_size_, M, N, K, enc);
+  if (transpose_ && vec_batch <= 8 && mode_ != QuantizationMode::Affine) {
+    assert(!biases);
+    fp_qmv(x, w, scales, out, bits_, group_size_, vec_batch, N, K, encoder, s);
     return;
   }
+
+  if (transpose_ && encoder.device().compute_capability_major() == 9) {
+    qmm_sm90(x, w, scales, biases, out, bits_, group_size_, mode_, encoder, s);
+    return;
+  }
+
+  throw std::runtime_error("QMM NYI");
 }
 
 void fast::Quantize::eval_gpu(