Fp qmv (#2984)

mlx/backend/cpu/quantized.cpp

@@ -14,6 +14,19 @@ namespace mlx::core {
 
 namespace {
 
+array ensure_row_contiguous(
+    const array& arr,
+    cpu::CommandEncoder& encoder,
+    Stream s) {
+  if (arr.flags().row_contiguous) {
+    return arr;
+  } else {
+    auto arr_cpy = contiguous_copy_cpu(arr, s);
+    encoder.add_temporary(arr_cpy);
+    return arr_cpy;
+  }
+};
+
 const static float FP4_LUT[16] = {
     +0.0f,
     +0.5f,
@@ -922,20 +935,9 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& scales_pre = inputs[2];
 
   auto& encoder = cpu::get_command_encoder(stream());
-  auto ensure_row_contiguous = [s = stream(), &encoder](const array& arr) {
-    if (arr.flags().row_contiguous) {
-      return arr;
-    } else {
-      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
-      copy_cpu(arr, arr_cpy, CopyType::General, s);
-      encoder.add_temporary(arr_cpy);
-      return arr_cpy;
-    }
-  };
-
-  auto x = ensure_row_contiguous(x_pre);
-  auto w = ensure_row_contiguous(w_pre);
-  auto scales = ensure_row_contiguous(scales_pre);
+  auto x = ensure_row_contiguous(x_pre, encoder, stream());
+  auto w = ensure_row_contiguous(w_pre, encoder, stream());
+  auto scales = ensure_row_contiguous(scales_pre, encoder, stream());
 
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -944,7 +946,7 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   encoder.set_input_array(scales);
   encoder.set_output_array(out);
   if (mode_ == QuantizationMode::Affine) {
-    auto biases = ensure_row_contiguous(inputs[3]);
+    auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
     encoder.set_input_array(biases);
     encoder.dispatch([out = array::unsafe_weak_copy(out),
                       x = array::unsafe_weak_copy(x),
@@ -1052,6 +1054,105 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
+uint8_t to_fp8_e8m0(float x) {
+  if (!std::isfinite(x)) {
+    return 0xFF;
+  }
+  if (x < 0.0f) {
+    return 0x00;
+  }
+  float le = std::log2(x);
+  int n = int(std::round(le));
+
+  n = n < -127 ? -127 : n;
+  n = n > 127 ? 127 : n;
+  return static_cast<uint8_t>(n + 127);
+}
+
+uint8_t to_fp4_e2m1(float x) {
+  if (std::isnan(x)) {
+    return 0x7;
+  }
+
+  const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
+  x = std::abs(x);
+
+  uint8_t bits;
+  if (x > 5.0f) {
+    bits = 0x7;
+  } else if (x >= 3.5f) {
+    bits = 0x6;
+  } else if (x > 2.5f) {
+    bits = 0x5;
+  } else if (x >= 1.75f) {
+    bits = 0x4;
+  } else if (x > 1.25f) {
+    bits = 0x3;
+  } else if (x >= 0.75f) {
+    bits = 0x2;
+  } else if (x > 0.25f) {
+    bits = 0x1;
+  } else {
+    bits = 0x0;
+  }
+  return bits | sign_bit;
+}
+
+template <typename T>
+void fp_quantize_dequantize(
+    const array& w_arr,
+    array& out_arr,
+    int bits,
+    int group_size,
+    size_t w_size) {
+  auto w = w_arr.data<T>();
+  auto out = out_arr.data<T>();
+
+  size_t n_groups = w_size / group_size;
+
+  for (size_t i = 0; i < n_groups; ++i) {
+    size_t idx = i * group_size;
+    float scale = -std::numeric_limits<float>::infinity();
+    for (int j = 0; j < group_size; ++j) {
+      scale = std::max(scale, std::abs(w[idx + j]));
+    }
+    scale /= bits == 4 ? 6.0f : 448.0f;
+    if (group_size == 16) {
+      scale = dequantize_scale<float, 16>(detail::ToFP8()(scale));
+    } else {
+      scale = dequantize_scale<float, 32>(to_fp8_e8m0(scale));
+    }
+
+    for (int j = 0; j < group_size; ++j) {
+      float w_el = scale == 0 ? 0.0f : w[idx + j] / scale;
+      float output;
+      if (bits == 8) {
+        output = detail::FromFP8()(detail::ToFP8()(w_el));
+      } else {
+        output = FP4_LUT[to_fp4_e2m1(w_el)];
+      }
+      out[idx + j] = static_cast<T>(scale * output);
+    }
+  }
+}
+
+void dispatch_quantize_dequantize(
+    const array& w,
+    array& out,
+    int bits,
+    int group_size) {
+  if (w.dtype() == float16) {
+    fp_quantize_dequantize<float16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == bfloat16) {
+    fp_quantize_dequantize<bfloat16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == float32) {
+    fp_quantize_dequantize<float>(w, out, bits, group_size, w.size());
+  } else {
+    throw std::runtime_error(
+        "[quantize_dequantize] Only supports floating point inputs");
+  }
+}
+
 template <typename T, typename U>
 void quantize(
     const T* w,
@@ -1136,15 +1237,8 @@ void dispatch_quantize(
 void fast::Quantize::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  auto ensure_row_contiguous = [s = stream()](const array& arr) {
-    if (arr.flags().row_contiguous) {
-      return std::make_pair(arr, false);
-    } else {
-      return std::make_pair(contiguous_copy_cpu(arr, s), true);
-    }
-  };
-
-  auto [w, copied] = ensure_row_contiguous(inputs[0]);
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto w = ensure_row_contiguous(inputs[0], encoder, stream());
   auto& out = outputs[0];
   out.set_data(allocator::malloc(out.nbytes()));
 
@@ -1152,10 +1246,6 @@ void fast::Quantize::eval_cpu(
   auto& biases = outputs[2];
   scales.set_data(allocator::malloc(scales.nbytes()));
   biases.set_data(allocator::malloc(biases.nbytes()));
-  auto& encoder = cpu::get_command_encoder(stream());
-  if (copied) {
-    encoder.add_temporary(w);
-  }
   encoder.set_input_array(w);
   encoder.set_input_array(scales);
   encoder.set_input_array(biases);
@@ -1238,6 +1328,43 @@ void fast::ConvertFP8::eval_cpu(
 }
 
 void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
-  throw std::runtime_error("QQMatmul not implemented on CPU.");
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    bool donate_x = inputs[0].is_donatable();
+    auto x = ensure_row_contiguous(inputs[0], encoder, stream());
+    auto w = ensure_row_contiguous(inputs[1], encoder, stream());
+    auto scales = ensure_row_contiguous(inputs[2], encoder, stream());
+
+    out.set_data(allocator::malloc(out.nbytes()));
+
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    encoder.set_input_array(x);
+    encoder.set_input_array(w);
+    encoder.set_input_array(scales);
+    encoder.set_output_array(out);
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      xhat = array::unsafe_weak_copy(xhat),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_]() mutable {
+      dispatch_quantize_dequantize(x, xhat, bits_, group_size_);
+      fp_qmm_dispatch(out, xhat, w, scales, group_size_, bits_, true);
+    });
+    return;
+  } else {
+    throw std::runtime_error("[QQMatmul] NYI for the general case");
+  }
 }
+
 } // namespace mlx::core

mlx/backend/cuda/CMakeLists.txt

@@ -56,7 +56,10 @@ 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/qmv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
 
@@ -66,12 +69,12 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
   target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp)
 else()
   target_sources(
-    mlx
-    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu)
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
+                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
 endif()
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)

mlx/backend/cuda/primitives.cpp

@@ -24,20 +24,12 @@ namespace mlx::core {
     throw std::runtime_error(#func " has no CUDA implementation.");   \
   }
 
-#if CUDART_VERSION < 12080
-void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
-  throw std::runtime_error(
-      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
-}
-#endif
-
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
-NO_GPU(QuantizedMatmul)
 NO_GPU(SegmentedMM)
 NO_GPU_MULTI(SVD)
 NO_GPU(Inverse)

mlx/backend/cuda/quantized/cuda_fp4.h

@@ -81,3 +81,20 @@ struct __nv_fp4_e2m1 {
   }
   uint8_t __x{0};
 };
+
+struct __nv_fp4x4_e2m1 {
+  __device__ operator float4() {
+    float4 out;
+    auto bits = __high & 0xf;
+    out.x = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__high >> 4) & 0xf;
+    out.y = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low) & 0xf;
+    out.z = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low >> 4) & 0xf;
+    out.w = float(*(__nv_fp4_e2m1*)(&bits));
+    return out;
+  }
+  uint8_t __high{0};
+  uint8_t __low{0};
+};

mlx/backend/cuda/quantized/fp_quantize.cu

@@ -11,8 +11,6 @@
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
-#include <cuda_fp4.h>
-#include <cuda_fp8.h>
 
 namespace mlx::core {
 namespace cu {
@@ -30,6 +28,69 @@ struct Dequantize {
 
 namespace cg = cooperative_groups;
 
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void fp_quantize_dequantize(T* w, T* out, size_t size) {
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0; // reserved bits for future use
+  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 = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
+  size_t base_idx = thread_idx * group_size;
+
+  if (base_idx >= size) {
+    return;
+  }
+
+  auto w_tile = load_vector<group_size, T>(w, thread_idx);
+  float scale = 0.0f;
+
+  Tx2 amax_2x = Tx2{0.0f, 0.0f};
+
+#pragma unroll
+  for (int i = 0; i < group_size; i += 2) {
+    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
+    abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+  }
+
+  scale = static_cast<float>(
+      max(fabsf(static_cast<float>(amax_2x.x)),
+          fabsf(static_cast<float>(amax_2x.y))));
+
+  scale /= bits == 4 ? 6.0f : 448.0f;
+  // Convert to mx scale or nv scale
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto s = ScaleType(scale);
+  scale = float(s);
+  AlignedVector<T, group_size> w_hat;
+
+#pragma unroll
+  for (int i = 0; i < group_size / 4; i++) {
+    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
+    float4 dq;
+    if constexpr (bits == 8) {
+      uint32_t quantized_val =
+          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp8(quantized_val);
+    } else {
+      uint16_t quantized_val =
+          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp4(quantized_val);
+    }
+    w_hat[i * 4] = static_cast<T>(dq.x * scale);
+    w_hat[i * 4 + 1] = static_cast<T>(dq.y * scale);
+    w_hat[i * 4 + 2] = static_cast<T>(dq.z * scale);
+    w_hat[i * 4 + 3] = static_cast<T>(dq.w * scale);
+  }
+  store_vector<group_size>(out, thread_idx, w_hat);
+}
+
 template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
 __global__ void
 fp_quantize_rowwise(T* w, uint8_t* out, uint8_t* scales, size_t size) {
@@ -280,6 +341,40 @@ get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) {
 
 } // namespace cu
 
+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(w);
+  enc.set_output_array(what);
+  dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      auto kernel = cu::fp_quantize_dequantize<T, 32, 4, true, false>;
+      if (bits == 8) {
+        kernel = cu::fp_quantize_dequantize<T, 32, 8, true, false>;
+      } else if (group_size == 16) {
+        kernel = cu::fp_quantize_dequantize<T, 16, 4, false, false>;
+      }
+      bool large = w.size() > UINT_MAX;
+      auto [num_blocks, block_dims] =
+          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);
+
+      enc.add_kernel_node(
+          kernel,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<T>(w),
+          gpu_ptr<T>(what),
+          w.size());
+    }
+  });
+}
+
 void fp_quantize(
     const array& w,
     array& wq,

mlx/backend/cuda/quantized/no_qqmm_impl.cpp

@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder&,
+    int,
+    int,
+    int,
+    bool,
+    int64_t,
+    bool,
+    int64_t,
+    array&,
+    const array&,
+    const array&,
+    const array&,
+    const array&,
+    Dtype,
+    QuantizationMode,
+    float) {
+  throw std::runtime_error(
+      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
+}
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmv.cu

@@ -0,0 +1,302 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+static constexpr int rows_per_block = 8;
+
+template <typename T>
+__device__ void adjust_matrix_offsets(
+    const T*& x,
+    const uint32_t*& w,
+    const uint8_t*& scales,
+    T*& y,
+    int output_stride,
+    const int& x_batch_ndims,
+    const Shape x_shape,
+    const Strides x_strides,
+    const int& w_batch_ndims,
+    const Shape w_shape,
+    const Strides w_strides,
+    const Strides s_strides) {
+  uint32_t idx = cg::this_grid().block_index().z;
+  if (x_batch_ndims == 1) {
+    x += idx * x_strides[0];
+  } else {
+    x += elem_to_loc(idx, x_shape.data(), x_strides.data(), x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += idx * w_strides[0];
+    scales += idx * s_strides[0];
+  } else {
+    auto [w_idx, s_idx] = elem_to_loc(
+        idx, w_shape.data(), w_strides.data(), s_strides.data(), w_batch_ndims);
+    w += w_idx;
+    scales += s_idx;
+  }
+  y += idx * output_stride;
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__device__ void fp_qmv_impl(
+    const uint32_t* mat,
+    const uint8_t* scales_,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  constexpr int vals_per_item = bits == 8 ? 4 : 8;
+  constexpr int nv_per_thread = vals_per_item * n_per_thread;
+  auto g_idx = block.group_index();
+  auto t_idx = block.thread_index();
+  int row = g_idx.y * rows_per_block + t_idx.y;
+
+  vec += g_idx.x * cols;
+  out += g_idx.x * rows;
+
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto scales = (ScaleType*)(scales_);
+  auto packed_cols = cols / vals_per_item;
+
+  if (row < rows) {
+    constexpr int scales_per_step = std::max(nv_per_thread / group_size, 1);
+    constexpr int scale_step = (WARP_SIZE * nv_per_thread) / group_size;
+    constexpr int n_per_step = n_per_thread / scales_per_step;
+    // Offset scales to correct row
+    scales += row * (cols / group_size) +
+        (warp.thread_rank() * nv_per_thread) / group_size;
+    float sum = 0.0f;
+    for (int col = n_per_thread * warp.thread_rank(); col < packed_cols;
+         col += (WARP_SIZE * n_per_thread)) {
+      auto local_vec =
+          unsafe_load_vector<nv_per_thread>(vec + vals_per_item * col, 0);
+      auto local_mat =
+          unsafe_load_vector<n_per_thread>(mat + row * packed_cols + col, 0);
+#pragma unroll
+      for (int i = 0; i < scales_per_step; ++i) {
+        float2 local_sum = {0.0f, 0.0f};
+#pragma unroll
+        for (int j = 0; j < n_per_step; ++j) {
+          int k = n_per_step * i + j;
+          if constexpr (bits == 8) {
+            auto v = dequant_fp8(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.x +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.y +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+          } else {
+            auto v = dequant_fp4(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+
+            v = dequant_fp4(local_mat[k] >> 16);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k + 4]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 5]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 6]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 7]);
+          }
+        }
+        sum += (local_sum.x + local_sum.y) * float(scales[i]);
+      }
+      scales += scale_step;
+    }
+
+    sum = cg::reduce(warp, sum, cg::plus<float>{});
+    if (warp.thread_rank() == 0) {
+      out[row] = static_cast<T>(sum);
+    }
+  }
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_single(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_batched(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols,
+    int vec_batch_ndims,
+    const __grid_constant__ Shape vec_shape,
+    const __grid_constant__ Strides vec_strides,
+    int mat_batch_ndims,
+    const __grid_constant__ Shape mat_shape,
+    const __grid_constant__ Strides mat_strides,
+    const __grid_constant__ Strides scales_strides) {
+  adjust_matrix_offsets<T>(
+      vec,
+      mat,
+      scales,
+      out,
+      rows * vec_shape[vec_batch_ndims],
+      vec_batch_ndims,
+      vec_shape,
+      vec_strides,
+      mat_batch_ndims,
+      mat_shape,
+      mat_strides,
+      scales_strides);
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <typename F>
+void dispatch_1_2_4(int n, F&& f) {
+  switch (n) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+  }
+}
+
+void fp_qmv(
+    const array& mat,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(mat);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(vec);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "qmv", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      dim3 block_dims{WARP_SIZE, rows_per_block};
+      uint B = out.size() / (M * N);
+      uint blocks_y = (N + rows_per_block - 1) / rows_per_block;
+      const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
+      const T* vec_ptr = gpu_ptr<T>(vec);
+      int n = 1;
+      if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<8>(vec_ptr)) ||
+           cu::is_aligned<4>(vec_ptr))) {
+        n = 4;
+      } else if (
+          cu::is_aligned<2>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<4>(vec_ptr)) ||
+           cu::is_aligned<2>(vec_ptr))) {
+        n = 2;
+      }
+      dispatch_1_2_4(n, [&](auto n) {
+        dispatch_bool(B > 1, [&](auto batched) {
+          if (!batched()) {
+            auto kernel = fp_qmv_single<T, rows_per_block, n(), 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_single<T, rows_per_block, n(), 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_single<T, rows_per_block, n(), 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint>(M), blocks_y},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K);
+          } else {
+            auto kernel = fp_qmv_batched<T, rows_per_block, n(), 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_batched<T, rows_per_block, n(), 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_batched<T, rows_per_block, n(), 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint>(M), blocks_y, B},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K,
+                vec.ndim() - 2,
+                const_param(vec.shape()),
+                const_param(vec.strides()),
+                mat.ndim() - 2,
+                const_param(mat.shape()),
+                const_param(mat.strides()),
+                const_param(scales.strides()));
+          }
+        });
+      });
+    }
+  });
+}
+
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/qmv.h

@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core::cu {
+
+void fp_qmv(
+    const array& w,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder);
+
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/qqmm.cpp

@@ -1,10 +1,11 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/quantized/cublas_qqmm.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/gpu/copy.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/primitives.h"
 
 #include <nvtx3/nvtx3.hpp>
@@ -13,16 +14,6 @@ namespace mlx::core {
 
 namespace {
 
-inline array
-ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) {
-  if (x.flags().row_contiguous || x.flags().col_contiguous) {
-    return x;
-  }
-  array x_copy = contiguous_copy_gpu(x, s);
-  enc.add_temporary(x_copy);
-  return x_copy;
-}
-
 array pad_and_swizzle_scales(
     const array& scale,
     cu::CommandEncoder& encoder,
@@ -47,64 +38,54 @@ array pad_and_swizzle_scales(
   return scale_tiled;
 }
 
-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) {
-  // 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
 
 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;
+  }
+  std::cout << "RUNNING FULL?" << std::endl;
+
   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.");
   }
-  assert(
-      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
-      (inputs.size() == 2));
-
   auto quantize = [&](const array& input,
                       cu::CommandEncoder& encoder,
                       const Stream& s) -> std::pair<array, array> {
@@ -136,9 +117,8 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     return {x_q, scales_x};
   };
   auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s);
-  auto [w_q, scale_w_pre] = (inputs[1].dtype() != uint32)
-      ? quantize(inputs[1], encoder, s)
-      : std::make_pair(inputs[1], inputs[2]);
+  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));
 

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

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

mlx/backend/cuda/quantized/quantized.cpp

@@ -2,60 +2,48 @@
 
 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/copy.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 {
 
-namespace {
+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);
 
-inline array ensure_row_contiguous(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (!x.flags().row_contiguous) {
-    array x_copy = contiguous_copy_gpu(x, s);
-    enc.add_temporary(x_copy);
-    return x_copy;
-  } else {
-    return x;
+  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;
   }
 }
 
-inline array
-ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) {
-  if (x.flags().row_contiguous || x.flags().col_contiguous) {
-    return x;
-  }
-  array x_copy = contiguous_copy_gpu(x, s);
-  enc.add_temporary(x_copy);
-  return x_copy;
-}
-
-inline array ensure_row_contiguous_matrix(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (x.ndim() < 2) {
-    if (x.strides()[0] == 1) {
-      return x;
-    }
-  } else {
-    auto stride_0 = x.strides()[x.ndim() - 2];
-    auto stride_1 = x.strides()[x.ndim() - 1];
-    if (stride_0 == x.shape(-1) && stride_1 == 1) {
-      return x;
-    }
-  }
-  array x_copy = contiguous_copy_gpu(x, s);
-  enc.add_temporary(x_copy);
-  return x_copy;
-}
-
-} // namespace
-
 void fast::Quantize::eval_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {

mlx/backend/cuda/quantized/quantized.h

@@ -42,4 +42,12 @@ void fp_dequantize(
     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

mlx/backend/cuda/quantized/quantized_utils.cuh

@@ -1,9 +1,22 @@
 // Copyright © 2025 Apple Inc.
 
+#include <cuda_fp4.h>
+#include <cuda_fp8.h>
+
 namespace mlx::core {
 
 namespace cu {
 
+inline __device__ float4 dequant_fp8(uint32_t bits) {
+  auto out = *(__nv_fp8x4_e4m3*)(&bits);
+  return out.operator float4();
+}
+
+inline __device__ float4 dequant_fp4(uint16_t bits) {
+  auto out = *(__nv_fp4x4_e2m1*)(&bits);
+  return out.operator float4();
+}
+
 template <int bits, int wsize = 8>
 inline constexpr __device__ short get_pack_factor() {
   return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);

mlx/backend/cuda/quantized/quantized_utils.h

@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/copy.h"
+
+namespace mlx::core {
+inline array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    enc.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline array ensure_row_contiguous_matrix(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (x.ndim() < 2) {
+    if (x.strides()[0] == 1) {
+      return x;
+    }
+  } else {
+    auto stride_0 = x.strides()[x.ndim() - 2];
+    auto stride_1 = x.strides()[x.ndim() - 1];
+    if (stride_0 == x.shape(-1) && stride_1 == 1) {
+      return x;
+    }
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+inline array
+ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) {
+  if (x.flags().row_contiguous || x.flags().col_contiguous) {
+    return x;
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+} // namespace mlx::core

mlx/backend/metal/kernels/fp_quantized.h

@@ -1766,7 +1766,6 @@ template <typename T, const int group_size, const int bits>
   uint8_t q_scale = s.bits;
   scale = float(s);
 
-  // Write out the scales and biases
   size_t gindex = index / group_size;
   if (index % group_size == 0) {
     scales[gindex] = q_scale;
@@ -1786,7 +1785,7 @@ template <typename T, const int group_size, const int bits>
 template <typename T, const int group_size, const int bits>
 [[kernel]] void fp_dequantize(
     const device uint8_t* w [[buffer(0)]],
-    const device T* scales [[buffer(1)]],
+    const device uint8_t* scales [[buffer(1)]],
     device T* out [[buffer(3)]],
     uint2 index [[thread_position_in_grid]],
     uint2 grid_dim [[threads_per_grid]]) {
@@ -1814,3 +1813,32 @@ template <typename T, const int group_size, const int bits>
     out[i] = static_cast<T>(scale * Dequantize<bits>{}(d));
   }
 }
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void fp_quantize_dequantize(
+    const device T* w [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    uint2 tidx [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr bool use_mx_scale = group_size == 32;
+  size_t index = tidx.x + grid_dim.x * size_t(tidx.y);
+
+  float scale;
+  float w_thread = w[index];
+  if (use_mx_scale) {
+    scale = simd_max(abs(w_thread));
+  } else {
+    float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0);
+    float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0);
+    scale = tidx.x < 16 ? w_max_l : w_max_r;
+  }
+  scale /= bits == 4 ? 6.0f : 448.0f;
+
+  using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;
+  auto s = ScaleType(scale);
+  scale = float(s);
+
+  uint8_t output = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
+
+  out[index] = static_cast<T>(scale * Dequantize<bits>{}(output));
+}

mlx/backend/metal/kernels/fp_quantized.metal

@@ -114,6 +114,12 @@
   instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits)
 
 #define instantiate_quantize_dequantize(type, mode, group_size, bits) \
+  instantiate_kernel( \
+    #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits, \
+    fp_quantize_dequantize, \
+    type, \
+    group_size,  \
+    bits) \
   instantiate_kernel( \
     #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits, \
     fp_quantize, \

mlx/backend/metal/primitives.cpp

@@ -239,8 +239,4 @@ void LUF::eval_gpu(
   throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI.");
 }
 
-void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& output) {
-  throw std::runtime_error("[QQMatmul::eval_gpu] Metal QQMatmul NYI.");
-}
-
 } // namespace mlx::core

mlx/backend/metal/quantized.cpp

@@ -1267,6 +1267,30 @@ void gather_qmm_rhs(
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
+void dispatch_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  // It is a qmv with a small inner dimension so route to qmv_quad kernel
+  if ((K == 128 || K == 64) && is_power_of_2(bits)) {
+    qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode);
+    return;
+  }
+
+  // Run of the mill qmv
+  qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode);
+}
+
 void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
@@ -1310,16 +1334,10 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
-  // It is a qmv with a small inner dimension so route to qmv_quad kernel
-  if (transpose_ && (K == 128 || K == 64) && is_power_of_2(bits_)) {
-    qmv_quad(
-        x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
-    return;
-  }
-
   // Run of the mill qmv
   if (transpose_) {
-    qmv(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
+    dispatch_qmv(
+        x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
     return;
   }
 
@@ -1443,6 +1461,99 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
       mode);
 }
 
+void quantize_dequantize(
+    const array& in,
+    array& out,
+    std::string mode,
+    int group_size,
+    int bits,
+    metal::Device& d,
+    const Stream& s) {
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  auto w = ensure_row_contiguous(in, d, s);
+  compute_encoder.set_input_array(w, 0);
+  compute_encoder.set_output_array(out, 1);
+  auto type_string = get_type_string(in.dtype());
+  std::string kname;
+  concatenate(
+      kname,
+      mode + "_quantize_dequantize_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits);
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "quantize_dequantize", mode, type_string, group_size, bits);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  constexpr int uint8_per_uint32 = 4;
+  constexpr int simd_size = 32;
+  int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits;
+  int per_thread = std::max(group_size / simd_size, 1);
+  size_t nthreads = w.size() / per_thread;
+
+  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  auto group_dims = MTL::Size(thread_group_size, 1, 1);
+  bool use_2d = nthreads > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() /= per_thread;
+  MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides())
+                               : MTL::Size(nthreads, 1, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto mode = quantization_mode_to_string(mode_);
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    out.set_data(allocator::malloc(out.nbytes()));
+
+    bool donate_x = inputs[0].is_donatable();
+    array x = ensure_row_contiguous(inputs[0], d, s);
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
+    quantize_dequantize(x, xhat, mode, group_size_, bits_, d, s);
+
+    // Make sure the last two dims of w and s are contiguous
+    array w = ensure_row_contiguous_matrix(inputs[1], d, s);
+    array scales = ensure_row_contiguous_matrix(inputs[2], d, 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);
+    dispatch_qmv(
+        xhat,
+        w,
+        scales,
+        std::nullopt,
+        out,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  } else {
+    throw std::runtime_error("[QQMatmul] NYI for the general case");
+  }
+}
+
 void fast::Quantize::eval_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {

mlx/ops.cpp

@@ -4345,10 +4345,6 @@ array qqmm(
     const std::string& mode /* = "nvfp4" */,
     StreamOrDevice s /* = {} */) {
   auto stream = to_stream(s);
-  if (stream.device != Device::gpu || !cu::is_available()) {
-    throw std::invalid_argument(
-        "[qqmm] Only supported on GPU with the CUDA backend.");
-  }
   auto qmode = string_to_quantization_mode(mode, "qqmm");
   // cuBLAS block scaled matmul only supports nvfp4 and mxfp8
   if (qmode != QuantizationMode::Nvfp4 && qmode != QuantizationMode::Mxfp8) {

python/tests/cuda_skip.py

@@ -47,7 +47,6 @@ cuda_skip = {
     "TestQuantized.test_qmm_shapes",
     "TestQuantized.test_qmm_vjp",
     "TestQuantized.test_qmv",
-    "TestQuantized.test_fp_qmv",
     "TestQuantized.test_fp_qvm",
     "TestQuantized.test_qvm",
     "TestQuantized.test_qvm_splitk",

python/tests/test_quantized.py

@@ -160,6 +160,38 @@ class TestQuantized(mlx_tests.MLXTestCase):
         w_hat = mx.dequantize(w_q, scales, mode="nvfp4")
         self.assertTrue(mx.all(w_hat == 0))
 
+    def test_qqmv(self):
+        key = mx.random.key(0)
+        k1, k2 = mx.random.split(key)
+        tests = product(
+            [256, 512, 67],  # M
+            [64, 256],  # N
+        )
+        modes = ["nvfp4", "mxfp8"]
+        for M, N in tests:
+            for mode in modes:
+                with self.subTest(shape=(M, N), mode=mode):
+                    x_shape = (1, N)
+                    w_shape = (M, N)
+
+                    x = mx.random.normal(shape=x_shape, key=k1)
+                    x_hat = mx.dequantize(
+                        *mx.quantize(x, mode=mode), mode=mode, dtype=mx.float32
+                    )
+
+                    w = mx.random.normal(shape=w_shape, key=k2)
+                    w_q, scales = mx.quantize(w, mode=mode)
+                    w_hat = mx.dequantize(w_q, scales, mode=mode, dtype=mx.float32)
+                    y_q = mx.qqmm(
+                        x,
+                        w_q,
+                        scales,
+                        mode=mode,
+                    )
+                    y_hat = x_hat @ mx.swapaxes(w_hat, -1, -2)
+                    self.assertEqual(y_q.shape, y_hat.shape)
+                    self.assertLess((y_q - y_hat).abs().max(), 1e-3)
+
     def test_qmm(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)
@@ -338,6 +370,28 @@ class TestQuantized(mlx_tests.MLXTestCase):
                     self.assertEqual(y_q.shape, y_hat.shape)
                     self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
+        # Test multiple of 16 but not 32
+        M = 128
+        N = 48
+        mode = "nvfp4"
+        with self.subTest(shape=(B, M, N), mode=mode):
+            x_shape = (1, N)
+            w_shape = (M, N)
+            x = mx.random.normal(shape=x_shape, key=k1)
+            w = mx.random.normal(shape=w_shape, key=k2)
+            w_q, scales = mx.quantize(w, mode=mode)
+            w_hat = mx.dequantize(w_q, scales, mode=mode)
+            y_q = mx.quantized_matmul(
+                x,
+                w_q,
+                scales,
+                transpose=True,
+                mode=mode,
+            )
+            y_hat = x @ mx.swapaxes(w_hat, -1, -2)
+            self.assertEqual(y_q.shape, y_hat.shape)
+            self.assertLess((y_q - y_hat).abs().max(), 1e-3)
+
     def test_qvm(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)