[CUDA] Add global scale support to gather_qmm

mlx/backend/cuda/device/qmm_naive.cuh

@@ -269,6 +269,10 @@ void qmm_naive_kernel(
   // For gather, use index lookup for input batch slicing.
   uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
   uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord;
+  // The global scale is per gathered expert; a non-gathered call has one.
+  if (global_scale && rhs_indices) {
+    global_scale += b_batch;
+  }
 
   // Get batch slice.
   Tensor mA = mA_mkl(_,_,a_batch); // (M,K)

mlx/backend/cuda/quantized/qmm/qmm.h

@@ -135,6 +135,7 @@ void gather_qmv(
     const array& w,
     const array& scales,
     const std::optional<array>& biases,
+    const std::optional<array>& global_scale,
     const array& lhs_indices,
     const array& rhs_indices,
     array& out,

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

@@ -222,6 +222,7 @@ __global__ void gather_qmv_kernel(
     const Q* w,
     const S* scales,
     const T* biases,
+    const float* global_scale,
     T* out,
     const uint32_t* lhs_indices,
     const uint32_t* rhs_indices,
@@ -244,8 +245,12 @@ __global__ void gather_qmv_kernel(
   x += block.group_index().y * k + m * k * x_idx;
   out += block.group_index().y * n + m * n * l;
 
+  // The global scale is per gathered expert.
+  if (global_scale) {
+    global_scale += w_idx;
+  }
   qmv_kernel_impl<elems_per_thread, group_size, has_bias, has_residue_k>(
-      x, w, scales, biases, nullptr, out, row, w_idx, n, k);
+      x, w, scales, biases, global_scale, out, row, w_idx, n, k);
 }
 
 template <
@@ -305,6 +310,7 @@ void gather_qmv(
     const Q* w,
     const S* scales,
     const T* biases,
+    const float* global_scale,
     T* out,
     const uint32_t* lhs_indices,
     const uint32_t* rhs_indices,
@@ -321,7 +327,16 @@ void gather_qmv(
       uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m), uint32_t(l)};
   dim3 block_dims{WARP_SIZE, rows_per_block};
   void* args[] = {
-      &x, &w, &scales, &biases, &out, &lhs_indices, &rhs_indices, &n, &k};
+      &x,
+      &w,
+      &scales,
+      &biases,
+      &global_scale,
+      &out,
+      &lhs_indices,
+      &rhs_indices,
+      &n,
+      &k};
 
   dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) {
     auto* kernel = &gather_qmv_kernel<
@@ -464,6 +479,7 @@ void gather_qmv(
     const array& w,
     const array& scales,
     const std::optional<array>& biases,
+    const std::optional<array>& global_scale,
     const array& lhs_indices,
     const array& rhs_indices,
     array& out,
@@ -490,6 +506,9 @@ void gather_qmv(
           if (biases) {
             encoder.set_input_array(*biases);
           }
+          if (global_scale) {
+            encoder.set_input_array(*global_scale);
+          }
           encoder.set_input_array(lhs_indices);
           encoder.set_input_array(rhs_indices);
           encoder.set_output_array(out);
@@ -499,6 +518,7 @@ void gather_qmv(
               gpu_ptr<Q>(w),
               gpu_ptr<S>(scales),
               biases ? gpu_ptr<T>(*biases) : nullptr,
+              global_scale ? gpu_ptr<float>(*global_scale) : nullptr,
               gpu_ptr<T>(out),
               gpu_ptr<uint32_t>(lhs_indices),
               gpu_ptr<uint32_t>(rhs_indices),

mlx/backend/cuda/quantized/qqmm.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
+#include "mlx/backend/common/broadcasting.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/quantized/qmm/qmm.h"
 #include "mlx/backend/cuda/quantized/qqmm_impl.h"
@@ -244,6 +245,14 @@ void GatherQQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
             ensure_contiguous(w_pre, encoder, s),
             ensure_contiguous(inputs[base_size - 1], encoder, s));
 
+  // The gather kernels take one global scale per expert.
+  if (global_scale_w) {
+    int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2);
+    array gs_e(Shape{E}, float32, nullptr, {});
+    broadcast(*global_scale_w, gs_e);
+    global_scale_w = ensure_row_contiguous(gs_e, encoder, s);
+  }
+
   // Quantize activation.
   array x = quantize_dequantize_input(
       x_pre, global_scale_x, bits_, group_size_, encoder, s);

mlx/backend/cuda/quantized/quantized.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2025 Apple Inc.
+// Copyright © 2025-2026 Apple Inc.
 
 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/device.h"
@@ -156,17 +156,16 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  if (mode_ != QuantizationMode::Affine && inputs.size() == 6) {
-    throw std::runtime_error(
-        "[GatherQMM] Global scale is only supported on the Metal backend.");
-  }
-
   array x = ensure_row_contiguous(inputs[0], encoder, s);
   const array& w = inputs[1];
   const array& scales = inputs[2];
+  // Affine gets biases at index 3, nvfp4 an optional global scale.
   std::optional<array> biases;
+  std::optional<array> global_scale;
   if (mode_ == QuantizationMode::Affine) {
     biases = inputs[3];
+  } else if (inputs.size() == 6) {
+    global_scale = ensure_row_contiguous(inputs[3], encoder, s);
   }
   array lhs_indices =
       ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s);
@@ -191,7 +190,10 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
         mode_,
         encoder.device());
   };
-  bool can_use_qmm_sm80 = supports(supports_qmm_sm80);
+  // qmm_sm80 does not apply global scales yet; route such calls to the
+  // naive kernel until it does.
+  bool can_use_qmm_sm80 =
+      !global_scale.has_value() && supports(supports_qmm_sm80);
   bool can_use_qmm_naive = supports(supports_qmm_naive);
   bool can_use_qmv = supports(supports_qmv);
 
@@ -217,7 +219,7 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
         w,
         scales,
         biases,
-        std::nullopt,
+        global_scale,
         lhs_indices,
         rhs_indices,
         out,
@@ -234,6 +236,7 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
         w,
         scales,
         biases,
+        global_scale,
         lhs_indices,
         rhs_indices,
         out,

mlx/ops.cpp

@@ -5575,9 +5575,9 @@ array gather_qmm(
           << ".";
       throw std::invalid_argument(msg.str());
     }
-    if (to_stream(s).device != Device::gpu || !metal::is_available()) {
+    if (to_stream(s).device != Device::gpu) {
       throw std::invalid_argument(
-          "[gather_qmm] Global scale is only supported on the Metal backend.");
+          "[gather_qmm] Global scale is only supported on the GPU.");
     }
   }
   if (qmode == QuantizationMode::Affine) {

python/src/ops.cpp

@@ -4865,7 +4865,7 @@ void init_ops(nb::module_& m) {
               :ref:`table of quantization modes <quantize-modes>`. Default: ``None``.
             mode (str, optional): The quantization mode. Default: ``"affine"``.
             global_scale (array, optional): The per-input float32 scale used for
-              ``nvfp4`` quantization of ``w``. Only supported on Metal.
+              ``nvfp4`` quantization of ``w``. Only supported on the GPU.
               Default: ``None``.
             sorted_indices (bool, optional): May allow a faster implementation
               if the passed indices are sorted. Default: ``False``.

python/tests/test_quantized.py

@@ -1601,7 +1601,6 @@ def test_gather_qqmm(self):
                 self.assertEqual(y_q.shape, y_hat.shape)
                 self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
-    @unittest.skipIf(mx.cuda.is_available(), "Not implemented for CUDA")
     def test_gather_qqmm_global_scale_matrix_paths(self):
         if mx.default_device() == mx.cpu:
             self.skipTest("Not implemented for CPU")
@@ -2090,7 +2089,7 @@ def gmm(s, x, wq):
                 ds = mx.grad(gmm)(s, x, wq)
 
     @unittest.skipIf(
-        not mx.metal.is_available(), "Global scale is only supported on Metal backend"
+        not mx.is_available(mx.gpu), "Global scale is only supported on the GPU"
     )
     def test_gather_qmm_global_scale(self):
         mx.random.seed(0)
@@ -2164,6 +2163,29 @@ def quantize_experts(w):
                 wrong = mx.gather_qmm(x, wq, s, global_scale=rotated, **kwargs)
                 self.assertGreater(rel_err(wrong, expected), 0.5)
 
+        # Repeated and out-of-order experts with explicit lhs indices on the
+        # M=1 vector path; each output must pick its own expert's scale. The
+        # long K is the decode shape that long-K vector kernels take.
+        E = 6
+        for K in (K, 2048):
+            with self.subTest(K=K):
+                w = (mx.random.normal((E, N, K)) / K**0.5).astype(mx.bfloat16)
+                wq, s, gs, w_hat = quantize_experts(w)
+                indices = mx.array([3, 0, 5, 3, 1, 0], mx.uint32)
+                x = (mx.random.normal((2, 1, K)) / K**0.5).astype(mx.bfloat16)
+                lhs = mx.array([0, 1, 0, 1, 0, 1], mx.uint32)
+                expected = x[lhs] @ w_hat[indices].swapaxes(-1, -2)
+                out = mx.gather_qmm(
+                    x,
+                    wq,
+                    s,
+                    lhs_indices=lhs,
+                    rhs_indices=indices,
+                    mode="nvfp4",
+                    global_scale=gs,
+                )
+                self.assertLess(rel_err(out, expected), 3e-2)
+
     def test_quantize_strided(self):
         N = 64
         mode = "nvfp4"