mlx 1.0.0

data/mlx/mlx/backend/cuda/quantized/qqmm.cpp

@@ -0,0 +1,158 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+array pad_and_swizzle_scales(
+    const array& scale,
+    cu::CommandEncoder& encoder,
+    const Stream& s) {
+  // Compute padded dimensions for full tiles (128 rows × 4 cols)
+  auto [pad_outer, pad_inner] =
+      get_padded_scale_dims(scale.shape(-2), scale.shape(-1));
+  // cuBLAS requirements for scale factor layout:
+  // 1. Dimensions must be padded to full tiles (128 rows × 4 cols)
+  // 2. Out-of-bounds values must be filled with zeros
+  // 3. Starting addresses must be 16-byte aligned
+  //
+  // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+  // Note: cu::malloc_async already provides 256-byte alignment
+  array scale_tiled(
+      cu::malloc_async(pad_outer * pad_inner, encoder),
+      Shape{pad_outer, pad_inner},
+      scale.dtype());
+  swizzle_scales(scale, scale_tiled, encoder, s);
+
+  encoder.add_temporary(scale_tiled);
+  return scale_tiled;
+}
+
+} // namespace
+
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(
+      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
+      (inputs.size() == 2));
+  nvtx3::scoped_range r("QQMatmul::eval_gpu");
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto& device = encoder.device();
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+    bool donate_x = inputs[0].is_donatable();
+    array x = ensure_row_contiguous(inputs[0], encoder, s);
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    fp_quantize_dequantize(x, xhat, group_size_, bits_, 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);
+
+    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);
+    return;
+  }
+
+  auto cc = device.compute_capability_major() * 100 +
+      device.compute_capability_minor() * 10;
+  if (cc < 1000) {
+    throw std::runtime_error(
+        "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher.");
+  }
+  auto quantize = [&](const array& input,
+                      cu::CommandEncoder& encoder,
+                      const Stream& s) -> std::pair<array, array> {
+    auto x = ensure_contiguous(input, encoder, s);
+    auto xq_shape = x.shape();
+    xq_shape.back() = x.shape(-1) * bits_ / 32;
+
+    auto sshape = x.shape();
+    const int64_t scales_inner = x.shape(-1) / group_size_;
+    auto [pad_outer, pad_inner] =
+        get_padded_scale_dims(x.shape(-2), scales_inner);
+    sshape[x.ndim() - 2] = pad_outer;
+    sshape[x.ndim() - 1] = pad_inner;
+    sshape.back() = scales_inner;
+
+    // Allocate outputs
+    const int64_t xq_bytes = x.size() * bits_ / 8;
+    const int64_t batch = x.size() / (x.shape(-2) * x.shape(-1));
+    const int64_t scales_bytes = batch * (pad_outer * pad_inner);
+
+    array x_q(cu::malloc_async(xq_bytes, encoder), std::move(xq_shape), uint32);
+    array scales_x(
+        cu::malloc_async(scales_bytes, encoder), std::move(sshape), uint8);
+
+    fp_quantize(x, x_q, scales_x, group_size_, bits_, encoder, s);
+
+    encoder.add_temporary(x_q);
+    encoder.add_temporary(scales_x);
+    return {x_q, scales_x};
+  };
+  auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s);
+  auto [w_q, scale_w_pre] = !w_quantized ? quantize(inputs[1], encoder, s)
+                                         : std::make_pair(inputs[1], inputs[2]);
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  auto out_dtype = out.dtype();
+
+  int M = x_q.shape(-2);
+  int N = w_q.shape(-2); // always transposed
+  int K_packed = x_q.shape(-1);
+  int K = K_packed * (32 / bits_);
+
+  // Repack scales from linear to tiled layout for tensor cores
+  array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s);
+
+  bool x_transposed = false;
+  bool w_transposed = true; // always transposed
+  int64_t lda = K;
+  int64_t ldb = K;
+
+  qqmm_impl(
+      encoder,
+      M,
+      N,
+      K,
+      x_transposed,
+      lda,
+      w_transposed,
+      ldb,
+      out,
+      x_q,
+      w_q,
+      scale_x,
+      scale_w,
+      out_dtype,
+      mode_);
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp

@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+namespace mlx::core {
+
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha) {
+  // Invoke CublasQQMM
+  std::string qmode = quantization_mode_to_string(mode);
+
+  // Currently only supports non-batched QQMM operations
+  // that covers all use cases for training, we will just collapse (batch,
+  // seq_len) into (tokens)
+  CublasQQMM qqmm(
+      encoder.device(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      1, // batch_count
+      0, // a_batch_stride
+      0, // b_batch_stride
+      out_dtype,
+      qmode);
+
+  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/qqmm_impl.h

@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha = 1.0f);
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu

@@ -0,0 +1,227 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cg = cooperative_groups;
+
+constexpr int TILE_ROWS = 128;
+constexpr int TILE_COLS = 4;
+constexpr int TILES_PER_LANE = 1;
+constexpr int LANES_PER_BLOCK = 32;
+
+// To pass scales to tensor cores, they need to be repacked into a tiled layout
+// https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+// Tiled layout for scale factors is very well described in CUTLASS
+// documentation:
+// https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md#scale-factor-layouts
+// Conceptually, it should be like this:
+// q_w = mx.zeros(shape=(M, N)) <-- zeros just for an example
+// s.shape = (M, N // 16) -- packed in row contigous order, group_size = 16
+// cbg_cnt = N // 16 // 4
+// rb_cnt = M // 128
+// tmp = x.reshape(rb_cnt, 4, 32, cbg_cnt, 4)
+// repacked_scales = tmp.transpose(0, 3, 2, 1, 4)
+// example: indecis of intial tile 128 x 4 of scales (packed in row major tensor
+// (M, K // 16), where M = 128, K = 64): array([[0, 1, 2, 3],
+//       [4, 5, 6, 7],
+//       [8, 9, 10, 11],
+//       ...,
+//       [500, 501, 502, 503],
+//       [504, 505, 506, 507],
+//       [508, 509, 510, 511]]
+// packed scales within tile 128 x 4:
+// array([[[[[0, 1, 2, 3], <-- s_0,0..s_0,3 scales
+//          [128, 129, 130, 131], <-- s_32,0..s_32,3 scales
+//          [256, 257, 258, 259], <-- s_64,0..s_64,3 scales
+//          [384, 385, 386, 387]], <-- s_96,0..s_96,3 scales
+//         [[4, 5, 6, 7], <-- s_1,0..s_1,3 scales
+//          [132, 133, 134, 135], ...
+//          [260, 261, 262, 263],
+//          [388, 389, 390, 391]],
+//         [[124, 125, 126, 127],
+//          [252, 253, 254, 255],
+//          [380, 381, 382, 383],
+//          [508, 509, 510, 511]]]]],
+
+inline std::tuple<dim3, dim3> get_swizzle_launch_args(
+    size_t M_swizzled,
+    size_t K_swizzled) {
+  constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+  constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32
+
+  const int num_tiles_k = K_swizzled / TILE_COLS;
+  const int num_tiles_m = M_swizzled / TILE_ROWS;
+
+  dim3 grid;
+  grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block);
+  grid.y = num_tiles_m;
+  grid.z = 1;
+  // Block is always (32, 32) = 1024 threads
+  dim3 block(LANES_PER_BLOCK, warps_per_block, 1);
+
+  return std::make_tuple(grid, block);
+}
+
+namespace cu {
+
+__global__ void swizzle_scales(
+    const uint8_t* scales_linear,
+    uint8_t* scales_swizzled,
+    const size_t M,
+    const size_t K,
+    const size_t M_swizzled,
+    const size_t K_swizzled) {
+  constexpr int tile_size = TILE_ROWS * TILE_COLS;
+  constexpr int num_tile_rows_per_thread = 4;
+  constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+
+  constexpr int tile_stride = tile_size / 16; // 32 int4s per tile
+
+  // Each thread loads 16 scales from 4 rows (stride 32) and packs them into
+  // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0,
+  // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc.
+  // The store is strided within a warp (stride 32 int4s), so we first
+  // write to shared memory, then do a coalesced store from shared to global
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+  auto linear_tid = tidy * block_size.x + tidx;
+
+  const int bid_x = block_idx.x;
+  const int bid_y = block_idx.y;
+
+  const int K_int = K_swizzled / 4;
+
+  const size_t output_offset = static_cast<size_t>(bid_y) * TILE_ROWS * K_int +
+      static_cast<size_t>(bid_x) * max_tiles_per_block * tile_size / 4;
+  int* output_block = reinterpret_cast<int*>(scales_swizzled) + output_offset;
+
+  const int grid_dim_x = cg::this_grid().dim_blocks().x;
+  const int grid_dim_y = cg::this_grid().dim_blocks().y;
+
+  int remaining = K_int - bid_x * max_tiles_per_block;
+  int tiles_in_block = min(remaining, max_tiles_per_block);
+  bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block;
+
+  __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride];
+
+  // Initialize to zero for padding
+  int thread_tile_rows[num_tile_rows_per_thread] = {0};
+
+  if (valid_tile) {
+    const size_t col_base =
+        static_cast<size_t>(bid_x) * max_tiles_per_block * TILE_COLS +
+        tidx * TILE_COLS;
+
+    const bool aligned_k = (K % 4 == 0);
+
+    if (aligned_k) {
+      // fast path: K is aligned, use vectorized loads with stride K/4
+      const int K_stride = K / 4;
+      const size_t block_offset =
+          static_cast<size_t>(bid_y) * TILE_ROWS * K_stride +
+          static_cast<size_t>(bid_x) * max_tiles_per_block;
+      const int* input_block =
+          reinterpret_cast<const int*>(scales_linear) + block_offset;
+// load
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        const int thread_offset =
+            (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE;
+        if (row < M && col_base + TILE_COLS <= K) {
+          thread_tile_rows[i] = __ldg(input_block + thread_offset);
+        } else if (row < M) {
+// partial tile at K boundary: load byte-by-byte
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row * K + col_base + c];
+            }
+          }
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        if (row < M) {
+          const size_t row_start = row * K;
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row_start + col_base + c];
+            }
+          }
+        }
+      }
+    }
+    // store to shared with XOR swizzle to avoid bank conflicts
+    int base_idx = tidx * tile_stride + tidy;
+    int xor_bits = (tidy >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    strided_scales_thread[swizzled_idx] =
+        *reinterpret_cast<int4*>(thread_tile_rows);
+  }
+
+  cg::thread_block block = cg::this_thread_block();
+  cg::sync(block);
+
+  const int total_int4s = tiles_in_block * tile_stride;
+#pragma unroll
+  for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) {
+    int tile_idx = i / tile_stride;
+    int row_idx = i % tile_stride;
+    int base_idx = tile_idx * tile_stride + row_idx;
+    int xor_bits = (row_idx >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    reinterpret_cast<int4*>(output_block)[i] =
+        strided_scales_thread[swizzled_idx];
+  }
+}
+} // namespace cu
+
+void swizzle_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(scales);
+  enc.set_output_array(scales_tiled);
+  // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
+  // are not multiples of tile sizes
+  size_t input_rows = scales.shape(-2);
+  size_t input_cols = scales.shape(-1);
+
+  size_t output_rows = scales_tiled.shape(-2);
+  size_t output_cols = scales_tiled.shape(-1);
+
+  auto [num_blocks, block_dims] =
+      get_swizzle_launch_args(output_rows, output_cols);
+  enc.add_kernel_node(
+      cu::swizzle_scales,
+      num_blocks,
+      block_dims,
+      0,
+      gpu_ptr<uint8_t>(scales),
+      gpu_ptr<uint8_t>(scales_tiled),
+      input_rows,
+      input_cols,
+      output_rows,
+      output_cols);
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/qqmm_utils.h

@@ -0,0 +1,30 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+// Compute padded dimensions for tiled layout
+// Tiles are 128 rows × 4 columns, must allocate full tiles
+inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
+  constexpr int rows_per_tile = 128;
+  constexpr int cols_per_tile = 4;
+
+  int padded_rows =
+      ((num_rows + rows_per_tile - 1) / rows_per_tile) * rows_per_tile;
+  int padded_cols =
+      ((num_cols + cols_per_tile - 1) / cols_per_tile) * cols_per_tile;
+
+  return {padded_rows, padded_cols};
+}
+
+void swizzle_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/quantized.cpp

@@ -0,0 +1,85 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+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);
+
+  out.set_data(cu::malloc_async(out.nbytes(), enc));
+
+  // 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);
+  }
+
+  bool non_batched = w.ndim() == 2 && x.flags().row_contiguous;
+  int K = x.shape(-1);
+  int M = non_batched ? x.size() / K : x.shape(-2);
+  int N = out.shape(-1);
+
+  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);
+    return;
+  }
+}
+
+void fast::Quantize::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("Quantize::eval_gpu");
+  auto& s = stream();
+  auto& d = cu::device(s.device);
+  auto& enc = d.get_command_encoder(s);
+
+  if (dequantize_) {
+    auto wq = ensure_row_contiguous(inputs[0], enc, s);
+    auto scales = ensure_row_contiguous(inputs[1], enc, s);
+    auto& w = outputs[0];
+
+    w.set_data(cu::malloc_async(w.nbytes(), enc));
+
+    if (mode_ == QuantizationMode::Affine) {
+      auto biases = ensure_row_contiguous(inputs[2], enc, s);
+      affine_dequantize(wq, scales, biases, w, group_size_, bits_, enc, s);
+    } else {
+      fp_dequantize(wq, scales, w, group_size_, bits_, enc, s);
+    }
+  } else {
+    auto w = ensure_contiguous(inputs[0], enc, s);
+    auto& wq = outputs[0];
+    auto& scales = outputs[1];
+
+    wq.set_data(cu::malloc_async(wq.nbytes(), enc));
+    scales.set_data(cu::malloc_async(scales.nbytes(), enc));
+    if (mode_ == QuantizationMode::Affine) {
+      auto& biases = outputs[2];
+      biases.set_data(cu::malloc_async(biases.nbytes(), enc));
+      affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s);
+    } else {
+      fp_quantize(w, wq, scales, group_size_, bits_, enc, s);
+    }
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/quantized/quantized.h

@@ -0,0 +1,53 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+void affine_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    array& biases,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void affine_dequantize(
+    const array& wq,
+    const array& scales,
+    const array& biases,
+    array& w,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_dequantize(
+    const array& wq,
+    const array& scales,
+    array& w,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core