Add optional int8 path for quantized matmul with affine weights

docs/src/usage/environment_variables.rst

@@ -26,6 +26,13 @@ General
    on supported hardware. The default is ``1``. Set it to ``0`` to keep these
    operations in full ``float32`` precision. See :doc:`precision`.
 
+.. envvar:: MLX_QMM_INT8
+
+   Run affine quantized matmuls on 8-bit integer weights and activations on
+   supported hardware. The default is ``0``. Set it to ``1``, ``4`` or ``2``
+   to enable it, from most to least precise. Other values are ignored. See
+   :doc:`precision`.
+
 Distributed
 -----------
 

docs/src/usage/precision.rst

@@ -31,3 +31,26 @@ To disable Winograd, set :envvar:`MLX_CONV_WINOGRAD` to ``0``:
 .. code-block:: shell
 
   MLX_CONV_WINOGRAD=0 python my_script.py
+
+8-bit Quantized Matmul
+----------------------
+
+Quantized matmuls dequantize the weights to the type of the input and
+multiply in floating point. On hardware with dedicated
+matrix-multiplication units MLX can instead round the weights and the input
+to 8 bits and multiply those, which is faster and less precise. This is off
+by default.
+
+For affine quantized weights set :envvar:`MLX_QMM_INT8`. The weights and the
+input are rounded to 8-bit integers with one scale per row. With ``1`` the
+scale of the input is refined for every block of 64 elements, with ``4`` for
+every four blocks, and ``2`` uses the row scale alone. Other values keep the
+default path:
+
+.. code-block:: shell
+
+  MLX_QMM_INT8=1 python my_script.py
+
+Expect a relative difference of one to a few percent from the floating point
+result. Only products of 64 rows or more with transposed weights take this
+path. The variable is read each time a quantized matmul is built.