metal: read FP16 and BF16 affine constants directly for FP32 inputs

mlx/backend/metal/kernels/quantized.h

@@ -847,11 +847,16 @@ METAL_FUNC void qmv_quad_impl(
   }
 }
 
-template <typename T, int group_size, int bits, bool partial_rows = false>
+template <
+    typename T,
+    int group_size,
+    int bits,
+    bool partial_rows = false,
+    typename S = T>
 METAL_FUNC void qmv_fast_impl(
     const device uint32_t* w,
-    const device T* scales,
-    const device T* biases,
+    const device S* scales,
+    const device S* biases,
     const device T* x,
     device T* y,
     const constant int& in_vec_size,
@@ -911,8 +916,8 @@ METAL_FUNC void qmv_fast_impl(
         src = min(row, last_row);
       }
       auto wl = (const device uint8_t*)(ws + src * in_vec_size_w);
-      const device T* sl = scales + src * in_vec_size_g;
-      const device T* bl = biases + src * in_vec_size_g;
+      const device S* sl = scales + src * in_vec_size_g;
+      const device S* bl = biases + src * in_vec_size_g;
 
       U s = sl[0];
       U b = bl[0];
@@ -942,8 +947,8 @@ METAL_FUNC void qmv_fast_impl(
           src = min(row, last_row);
         }
         auto wl = (const device uint8_t*)(ws + src * in_vec_size_w);
-        const device T* sl = scales + src * in_vec_size_g;
-        const device T* bl = biases + src * in_vec_size_g;
+        const device S* sl = scales + src * in_vec_size_g;
+        const device S* bl = biases + src * in_vec_size_g;
 
         U s = sl[0];
         U b = bl[0];
@@ -1739,6 +1744,31 @@ template <typename T, int group_size, int bits, int D, bool batched>
       quad_lid);
 }
 
+template <typename S, int group_size, int bits, bool partial_rows>
+[[kernel]] void affine_qmv_fast_mixed(
+    const device uint32_t* w [[buffer(0)]],
+    const device S* scales [[buffer(1)]],
+    const device S* biases [[buffer(2)]],
+    const device float* x [[buffer(3)]],
+    device float* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  qmv_fast_impl<float, group_size, bits, partial_rows, S>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
 template <
     typename T,
     int group_size,

mlx/backend/metal/kernels/quantized.metal

@@ -226,4 +226,30 @@
 instantiate_quantized_all_groups(float)
 instantiate_quantized_all_groups(float16_t)
 instantiate_quantized_all_groups(bfloat16_t)
-// clang-format on
+
+#define instantiate_qmv_mixed(type, group_size, bits, partial_rows) \
+  instantiate_kernel( \
+      "affine_qmv_fast_mixed_" #type "_gs_" #group_size "_b_" #bits "_pr_" #partial_rows, \
+      affine_qmv_fast_mixed, type, group_size, bits, partial_rows)
+
+#define instantiate_qmv_mixed_rows(type, group_size, bits) \
+  instantiate_qmv_mixed(type, group_size, bits, 0)        \
+  instantiate_qmv_mixed(type, group_size, bits, 1)
+
+#define instantiate_qmv_mixed_bits(type, group_size) \
+  instantiate_qmv_mixed_rows(type, group_size, 1)    \
+  instantiate_qmv_mixed_rows(type, group_size, 2)    \
+  instantiate_qmv_mixed_rows(type, group_size, 3)    \
+  instantiate_qmv_mixed_rows(type, group_size, 4)    \
+  instantiate_qmv_mixed_rows(type, group_size, 5)    \
+  instantiate_qmv_mixed_rows(type, group_size, 6)    \
+  instantiate_qmv_mixed_rows(type, group_size, 8)
+
+#define instantiate_qmv_mixed_groups(type) \
+  instantiate_qmv_mixed_bits(type, 32)     \
+  instantiate_qmv_mixed_bits(type, 64)     \
+  instantiate_qmv_mixed_bits(type, 128)
+
+instantiate_qmv_mixed_groups(float16_t)
+instantiate_qmv_mixed_groups(bfloat16_t)
+    // clang-format on

mlx/backend/metal/quantized.cpp

@@ -2142,6 +2142,53 @@ void dispatch_qmv(
       mode);
 }
 
+void qmv_mixed(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s) {
+  std::string kname;
+  auto scale_type = get_type_string(scales.dtype());
+  bool partial_rows = N % 8 != 0;
+  concatenate(
+      kname,
+      "affine_qmv_fast_mixed_",
+      scale_type,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_pr_",
+      partial_rows);
+  auto kernel = get_quantized_kernel_wrapped(
+      d,
+      kname,
+      "qmv_fast_mixed",
+      "affine",
+      scale_type,
+      group_size,
+      bits,
+      partial_rows);
+  auto& encoder = metal::get_command_encoder(s);
+  encoder.set_compute_pipeline_state(kernel);
+  encoder.set_input_array(w, 0);
+  encoder.set_input_array(scales, 1);
+  encoder.set_input_array(biases, 2);
+  encoder.set_input_array(x, 3);
+  encoder.set_output_array(out, 4);
+  encoder.set_bytes(K, 5);
+  encoder.set_bytes(N, 6);
+  encoder.dispatch_threadgroups(
+      MTL::Size(1, (N + 7) / 8, 1), MTL::Size(32, 2, 1));
+}
+
 void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
@@ -2167,6 +2214,23 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
   auto mode = quantization_mode_to_string(mode_);
 
+  if (mode_ == QuantizationMode::Affine && scales.dtype() != x.dtype()) {
+    if (transpose_ && non_batched && M == 1 &&
+        K % qmv_fast_k_alignment(bits_) == 0) {
+      qmv_mixed(x, w, scales, *biases, out, group_size_, bits_, N, K, d, s);
+      return;
+    }
+    // Other kernels require constants with the activation dtype.
+    auto widen = [&](const array& a) {
+      array copy(a.shape(), x.dtype(), nullptr, {});
+      copy_gpu(a, copy, CopyType::General, s);
+      metal::get_command_encoder(s).add_temporary(copy);
+      return copy;
+    };
+    scales = widen(scales);
+    biases = widen(*biases);
+  }
+
   // For few rows use qmv_mma when N/K is large enough.
   if (transpose_ && non_batched && (M >= 6) && (M <= 16) &&
       use_qmv_mma(x, w, scales, mode_, M, N, K, d)) {

mlx/ops.cpp

@@ -4966,7 +4966,14 @@ array quantized_matmul(
     throw std::invalid_argument(msg.str());
   }
   std::vector<array> inputs;
-  if (qmode == QuantizationMode::Affine) {
+  bool narrow_constants = qmode == QuantizationMode::Affine &&
+      x.dtype() == float32 &&
+      (scales.dtype() == float16 || scales.dtype() == bfloat16) &&
+      biases->dtype() == scales.dtype() && to_stream(s).device == Device::gpu &&
+      metal::is_available();
+  if (narrow_constants) {
+    inputs = {x, w, scales, *biases};
+  } else if (qmode == QuantizationMode::Affine) {
     inputs = {
         astype(x, dtype, s),
         w,

mlx/primitives.cpp

@@ -3673,6 +3673,7 @@ std::vector<array> QuantizedMatmul::vjp(
         wq = unflatten(wq, -1, {-1, group_size_}, stream());
         vjps.push_back(sum(multiply(*dsb, wq, stream()), -1, false, stream()));
       }
+      vjps.back() = astype(vjps.back(), primals[arg].dtype(), stream());
     }
   }
   return vjps;

python/tests/test_export_import.py

@@ -12,6 +12,29 @@
 
 
 class TestExportImport(mlx_tests.MLXTestCase):
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_quantized_matmul_narrow_constants(self):
+        for dtype in [mx.float16, mx.bfloat16]:
+            for m in [1, 3]:
+                x = mx.random.normal((m, 1024))
+                wq, scales, biases = mx.quantize(
+                    mx.random.normal((17, 1024)).astype(dtype), group_size=64, bits=4
+                )
+
+                def fun(x, scales, biases):
+                    return mx.quantized_matmul(x, wq, scales, biases, True, 64, 4)
+
+                path = os.path.join(self.test_dir, "qmm-narrow.mlxfn")
+                mx.export_function(path, fun, x, scales, biases)
+                imported = mx.import_function(path)
+                fresh = mx.random.normal(x.shape)
+                (actual,) = imported(fresh, scales, biases)
+                expected = fun(
+                    fresh, scales.astype(mx.float32), biases.astype(mx.float32)
+                )
+                self.assertEqual(actual.dtype, mx.float32)
+                self.assertTrue(mx.array_equal(actual, expected))
+
     @classmethod
     def setUpClass(cls):
         cls.test_dir_fid = tempfile.TemporaryDirectory()

python/tests/test_quantized.py

@@ -683,6 +683,120 @@ def test_qmm_small_m_block(self):
                 )
                 self.assertLess((y_ref - y).abs().max(), 1e-3)
 
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_qmm_narrow_affine_constants(self):
+        for dtype, bits, gs, n in product(
+            [mx.float16, mx.bfloat16], [1, 2, 3, 4, 5, 6, 8], [32, 64, 128], [1, 16, 17]
+        ):
+            k = 1536 if bits in [3, 6] else 2048
+            with self.subTest(dtype=dtype, bits=bits, gs=gs, n=n):
+                key_x, key_w = mx.random.split(mx.random.key(7))
+                x = mx.random.normal((1, k), key=key_x)
+                w = mx.random.normal((n, k), key=key_w).astype(dtype)
+                wq, scales, biases = mx.quantize(w, group_size=gs, bits=bits)
+                actual = mx.quantized_matmul(
+                    x, wq, scales, biases, group_size=gs, bits=bits, stream=mx.gpu
+                )
+                expected = mx.quantized_matmul(
+                    x,
+                    wq,
+                    scales.astype(mx.float32),
+                    biases.astype(mx.float32),
+                    group_size=gs,
+                    bits=bits,
+                    stream=mx.gpu,
+                )
+                self.assertEqual(actual.dtype, mx.float32)
+                self.assertTrue(mx.array_equal(actual, expected))
+
+    def test_qmm_narrow_affine_fallbacks(self):
+        for dtype, m, k, transpose in product(
+            [mx.float16, mx.bfloat16], [1, 2, 33], [128, 768], [False, True]
+        ):
+            with self.subTest(dtype=dtype, m=m, k=k, transpose=transpose):
+                x = mx.random.normal((m, k))
+                shape = (128, k) if transpose else (k, 128)
+                wq, scales, biases = mx.quantize(
+                    mx.random.normal(shape).astype(dtype), group_size=32, bits=4
+                )
+                for device in [mx.cpu, mx.default_device()]:
+                    actual = mx.quantized_matmul(
+                        x, wq, scales, biases, transpose, 32, 4, stream=device
+                    )
+                    expected = mx.quantized_matmul(
+                        x,
+                        wq,
+                        scales.astype(mx.float32),
+                        biases.astype(mx.float32),
+                        transpose,
+                        32,
+                        4,
+                        stream=device,
+                    )
+                    self.assertTrue(mx.array_equal(actual, expected))
+
+    def test_qmm_narrow_affine_layouts(self):
+        for dtype in [mx.float16, mx.bfloat16]:
+            wq, scales, biases = mx.quantize(
+                mx.random.normal((2, 34, 1024)).astype(dtype), group_size=64, bits=2
+            )
+            wq, scales, biases = wq[:, ::2], scales[:, ::2], biases[:, ::2]
+            x = mx.random.normal((2, 1, 2048))[..., ::2]
+            for inputs in [
+                (x, wq, scales, biases),
+                (x[0, 0], wq[0], scales[0], biases[0]),
+            ]:
+                a, w, s, b = inputs
+                actual = mx.quantized_matmul(a, w, s, b, True, 64, 2)
+                expected = mx.quantized_matmul(
+                    a, w, s.astype(mx.float32), b.astype(mx.float32), True, 64, 2
+                )
+                self.assertTrue(mx.array_equal(actual, expected))
+
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_qmm_narrow_affine_transforms(self):
+        for dtype in [mx.float16, mx.bfloat16]:
+            x = mx.random.normal((1, 1024))
+            wq, scales, biases = mx.quantize(
+                mx.random.normal((9, 1024)).astype(dtype), group_size=64, bits=4
+            )
+
+            def narrow(x, scales, biases):
+                return mx.quantized_matmul(x, wq, scales, biases, True, 64, 4)
+
+            def widened(x, scales, biases):
+                return narrow(x, scales.astype(mx.float32), biases.astype(mx.float32))
+
+            actual = mx.grad(lambda x, s, b: narrow(x, s, b).sum(), (0, 1, 2))(
+                x, scales, biases
+            )
+            expected = mx.grad(lambda x, s, b: widened(x, s, b).sum(), (0, 1, 2))(
+                x, scales, biases
+            )
+            for a, e in zip(actual, expected):
+                self.assertEqual(a.dtype, e.dtype)
+                self.assertTrue(mx.array_equal(a, e))
+
+            compiled = mx.compile(narrow)
+            x_new = mx.random.normal(x.shape, key=mx.random.key(29))
+            s_new = scales * mx.array(0.75, dtype=dtype)
+            b_new = biases + mx.array(0.03125, dtype=dtype)
+            self.assertTrue(
+                mx.array_equal(
+                    compiled(x_new, s_new, b_new), widened(x_new, s_new, b_new)
+                )
+            )
+            compiled_grad = mx.compile(
+                mx.grad(lambda x, s, b: narrow(x, s, b).sum(), (0, 1, 2))
+            )
+            actual = compiled_grad(x_new, s_new, b_new)
+            expected = mx.grad(lambda x, s, b: widened(x, s, b).sum(), (0, 1, 2))(
+                x_new, s_new, b_new
+            )
+            for a, e in zip(actual, expected):
+                self.assertEqual(a.dtype, e.dtype)
+                self.assertTrue(mx.array_equal(a, e))
+
     @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
     def test_qmm_splitk_precision(self):
         # Split-K sums the partitions in float32 and rounds the result to the