qqmm (#2789)

Co-authored-by: root <root@bolt-t9a77vmteu-94s9t6ymth.bolt-pods.turi-bolt.svc.cluster.local>
Co-authored-by: root <root@bolt-5azkyvd8ga-kgfzk84y6m.bolt-pods.turi-bolt.svc.cluster.local>
Co-authored-by: root <root@bolt-y4nktpaecv-ssnx24rdha.bolt-pods.turi-bolt.svc.cluster.local>

python/src/ops.cpp

@@ -5433,4 +5433,54 @@ void init_ops(nb::module_& m) {
         Returns:
             array or Sequence[array]: The outputs which depend on dependencies.
       )pbdoc");
+  m.def(
+      "qqmm",
+      &mx::qqmm,
+      nb::arg(), // x
+      nb::arg(), // w_q
+      "scales"_a = nb::none(), // scales w
+      "group_size"_a = nb::none(),
+      "bits"_a = nb::none(),
+      "mode"_a = "nvfp4",
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def qqmm(x: array, w: array, scales: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Perform a matrix multiplication using a possibly quantized weight matrix
+      ``w`` and a non-quantized input ``x``. The input ``x`` is quantized on the
+      fly. The weight matrix ``w`` is used as-is if it is already quantized;
+      otherwise, it is quantized on the fly.
+
+      If ``w`` is quantized, ``scales`` must be provided, and ``group_size``,
+      ``bits``, and ``mode`` must match the parameters that were used to quantize
+      ``w``.
+
+      Notes:
+        If ``w`` is expected to receive gradients, it must be provided in
+        non-quantized form.
+
+        If ``x`` and `w`` are not quantized, their data types must be ``float32``, 
+        ``float16``, or ``bfloat16``.
+        If ``w`` is quantized, it must be packed in unsigned integers.
+        
+      Args:
+        x (array): Input array.
+        w (array): Weight matrix. If quantized, it is packed in unsigned integers.
+        scales (array, optional): The scales to use per ``group_size`` elements of
+          ``w`` if ``w`` is quantized. Default: ``None``.
+        group_size (int, optional): Number of elements in ``x`` and ``w`` that
+          share a scale. See supported values and defaults in the
+          :ref:`table of quantization modes <quantize-modes>`. Default: ``None``.
+        bits (int, optional): Number of bits used to represent each element of
+          ``x`` and ``w``. See supported values and defaults in the
+          :ref:`table of quantization modes <quantize-modes>`. Default: ``None``.
+        mode (str, optional): The quantization mode. Default: ``"nvfp4"``.
+          Supported modes are ``nvfp4`` and ``mxfp8``. See the
+          :ref:`table of quantization modes <quantize-modes>` for details.
+
+      Returns:
+        array: The result of the multiplication of quantized ``x`` with quantized ``w``.
+        needed).
+  )pbdoc");
 }