Implied-bias affine quantized_matmul / dequantize (0-d biases factor)

mlx/backend/common/quantized.h

@@ -1,5 +1,11 @@
 // Copyright © 2026 Apple Inc.
 
+#pragma once
+
+#include <optional>
+
+#include "mlx/array.h"
+
 namespace mlx::core {
 
 inline constexpr short get_pack_factor(int bits, int wsize = 8) {
@@ -11,4 +17,14 @@ inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
   return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
 }
 
+// Implied bias (affine mode): a 0-d `biases` array holds a single factor f
+// standing for the per-group bias `scales * T(f)` in the scales' dtype T.
+inline bool is_implied_bias(const array& biases) {
+  return biases.ndim() == 0;
+}
+
+inline bool is_implied_bias(const std::optional<array>& biases) {
+  return biases.has_value() && is_implied_bias(*biases);
+}
+
 } // namespace mlx::core

mlx/backend/cpu/quantized.cpp

@@ -950,6 +950,66 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
     }
     auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
     encoder.set_input_array(biases);
+    if (is_implied_bias(biases)) {
+      // Implied bias: `biases` is a single factor f, exact in the scales'
+      // dtype T by contract. The CPU kernels index a per-group bias array, so
+      // materialize f * scales into a temporary: the float product of two T
+      // values is exact, so rounding it to T matches `scales * T(f)`.
+      array full_biases(scales.shape(), scales.dtype(), nullptr, {});
+      full_biases.set_data(allocator::malloc(full_biases.nbytes()));
+      encoder.add_temporary(full_biases);
+      encoder.set_output_array(full_biases);
+      encoder.dispatch([out = array::unsafe_weak_copy(out),
+                        x = array::unsafe_weak_copy(x),
+                        w = array::unsafe_weak_copy(w),
+                        scales = array::unsafe_weak_copy(scales),
+                        factor = array::unsafe_weak_copy(biases),
+                        full_biases = array::unsafe_weak_copy(full_biases),
+                        group_size_ = group_size_,
+                        bits_ = bits_,
+                        transpose_ = transpose_]() mutable {
+        float f;
+        switch (factor.dtype()) {
+          case float32:
+            f = factor.data<float>()[0];
+            break;
+          case float16:
+            f = static_cast<float>(factor.data<float16_t>()[0]);
+            break;
+          case bfloat16:
+            f = static_cast<float>(factor.data<bfloat16_t>()[0]);
+            break;
+          default:
+            throw std::runtime_error(
+                "[QuantizedMatmul::eval_cpu] Implied-bias factor must be a "
+                "real floating type.");
+        }
+        auto fill = [&](auto* dst, const auto* src) {
+          using T = std::remove_pointer_t<decltype(dst)>;
+          for (size_t i = 0; i < scales.size(); ++i) {
+            dst[i] = static_cast<T>(f * static_cast<float>(src[i]));
+          }
+        };
+        switch (scales.dtype()) {
+          case float32:
+            fill(full_biases.data<float>(), scales.data<float>());
+            break;
+          case float16:
+            fill(full_biases.data<float16_t>(), scales.data<float16_t>());
+            break;
+          case bfloat16:
+            fill(full_biases.data<bfloat16_t>(), scales.data<bfloat16_t>());
+            break;
+          default:
+            throw std::runtime_error(
+                "[QuantizedMatmul::eval_cpu] Only real floating scales are "
+                "supported with an implied bias.");
+        }
+        _qmm_dispatch(
+            out, x, w, scales, full_biases, group_size_, bits_, transpose_);
+      });
+      return;
+    }
     encoder.dispatch([out = array::unsafe_weak_copy(out),
                       x = array::unsafe_weak_copy(x),
                       w = array::unsafe_weak_copy(w),