From c0fea26ed2fc3758d1c3a95af7809764039e438d Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Sun, 21 Dec 2025 06:12:41 -0800
Subject: [PATCH] Fix for non row-contig scales (#2941)

---
 mlx/backend/metal/quantized.cpp | 11 ++++++-----
 mlx/ops.cpp                     |  7 ++++++-
 python/tests/test_quantized.py  | 11 +++++++++++
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp
index 63aa2f94..723bccc6 100644
--- a/mlx/backend/metal/quantized.cpp
+++ b/mlx/backend/metal/quantized.cpp
@@ -1452,25 +1452,26 @@ void fast::Quantize::eval_gpu(
   auto& compute_encoder = d.get_command_encoder(s.index);
 
   auto w = ensure_row_contiguous(w_pre, d, s);
-  compute_encoder.set_input_array(w, 0);
   if (dequantize_) {
     auto scales = ensure_row_contiguous(inputs[1], d, s);
-    compute_encoder.set_input_array(scales, 1);
-    compute_encoder.set_output_array(out, 3);
     if (mode_ == QuantizationMode::Affine) {
       auto biases = ensure_row_contiguous(inputs[2], d, s);
       compute_encoder.set_input_array(biases, 2);
     }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_input_array(scales, 1);
+    compute_encoder.set_output_array(out, 3);
   } else {
     auto& scales = outputs[1];
     scales.set_data(allocator::malloc(scales.nbytes()));
-    compute_encoder.set_output_array(out, 1);
-    compute_encoder.set_output_array(scales, 2);
     if (mode_ == QuantizationMode::Affine) {
       auto& biases = outputs[2];
       biases.set_data(allocator::malloc(biases.nbytes()));
       compute_encoder.set_output_array(biases, 3);
     }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_output_array(scales, 2);
   }
 
   auto type_string = dequantize_ ? get_type_string(out.dtype())
diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 02d4508a..3d1c0922 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -4188,6 +4188,11 @@ std::pair<Dtype, QuantizationMode> validate_mode_with_type(
     } else {
       return {dtype, qmode};
     }
+  } else if (scales.dtype() != uint8) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Scale type must be uint8 but received type "
+        << scales.dtype() << ".";
+    throw std::invalid_argument(msg.str());
   }
   if (biases) {
     std::ostringstream msg;
@@ -6070,4 +6075,4 @@ array contiguous(
       {a});
 }
 
-} // namespace mlx::core
\ No newline at end of file
+} // namespace mlx::core
diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py
index ce63544b..87b1f2ff 100644
--- a/python/tests/test_quantized.py
+++ b/python/tests/test_quantized.py
@@ -1015,6 +1015,17 @@ class TestQuantized(mlx_tests.MLXTestCase):
 
             ds = mx.grad(gmm)(s, x, wq)
 
+    def test_quantize_strided(self):
+        N = 64
+        mode = "nvfp4"
+        w = mx.random.normal(shape=(N, N))
+        w_q, scales = mx.quantize(w, mode="nvfp4")
+
+        scales = mx.broadcast_to(mx.array(56, mx.uint8), scales.shape)
+        w_hat = mx.dequantize(w_q, scales, mode=mode)
+        expected = mx.dequantize(w_q, mx.contiguous(scales), mode=mode)
+        self.assertTrue(mx.allclose(w_hat, expected))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()