diff --git a/mlx/backend/cuda/device/qmm_naive.cuh b/mlx/backend/cuda/device/qmm_naive.cuh index e2ad855b9f..5eea030c50 100644 --- a/mlx/backend/cuda/device/qmm_naive.cuh +++ b/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) diff --git a/mlx/backend/cuda/quantized/qmm/qmm.h b/mlx/backend/cuda/quantized/qmm/qmm.h index 64cfecfd5a..6991baad12 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/mlx/backend/cuda/quantized/qmm/qmm.h @@ -135,6 +135,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, diff --git a/mlx/backend/cuda/quantized/qmm/qmv.cu b/mlx/backend/cuda/quantized/qmm/qmv.cu index 4ec6b95a7a..63c8478442 100644 --- a/mlx/backend/cuda/quantized/qmm/qmv.cu +++ b/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( - 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& biases, + const std::optional& 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(w), gpu_ptr(scales), biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, gpu_ptr(out), gpu_ptr(lhs_indices), gpu_ptr(rhs_indices), diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index 196bd9c05e..5a9fba33d5 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/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& 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); diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 03e6b4f028..5c3e255805 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/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& 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 biases; + std::optional 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& 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& 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& inputs, array& out) { w, scales, biases, + global_scale, lhs_indices, rhs_indices, out, diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 7ec5369820..95010eb854 100644 --- a/mlx/ops.cpp +++ b/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) { diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 678cd984cd..a36b63b3d6 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4865,7 +4865,7 @@ void init_ops(nb::module_& m) { :ref:`table of quantization 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``. diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 74e6980869..7a5547411e 100644 --- a/python/tests/test_quantized.py +++ b/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"