L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tensor scale nvfp4 (#3022)

Browse Source
This commit is contained in:
Anastasiia Filippova
2026-02-18 11:19:26 +01:00
committed by GitHub
parent 360639c2df
commit 06305022ab
19 changed files with 616 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/ops.cpp
+16 -4
View File
@@ -4304,10 +4304,11 @@ void init_ops(nb::module_& m) {
"group_size"_a = nb::none(),
"bits"_a = nb::none(),
"mode"_a = "affine",
"global_scale"_a = nb::none(),
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def quantize(w: array, /, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, stream: Union[None, Stream, Device] = None) -> tuple[array, array, array]"),
"def quantize(w: array, /, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, global_scale: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> tuple[array, array, array]"),
R"pbdoc(
Quantize the array ``w``.
@@ -4332,6 +4333,8 @@ void init_ops(nb::module_& m) {
``w`` in the quantized array. See supported values and defaults in the
:ref:`table of quantization modes <quantize-modes>`. Default: ``None``.
mode (str, optional): The quantization mode. Default: ``"affine"``.
global_scale (array, optional): The per-input float32 scale used for
``"nvfp4"`` quantization if provided. Default: ``None``.
Returns:
tuple: A tuple with either two or three elements containing:
@@ -4401,11 +4404,12 @@ void init_ops(nb::module_& m) {
"group_size"_a = nb::none(),
"bits"_a = nb::none(),
"mode"_a = "affine",
"global_scale"_a = nb::none(),
"dtype"_a = nb::none(),
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
"def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', global_scale: Optional[array] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Dequantize the matrix ``w`` using quantization parameters.
@@ -4420,6 +4424,8 @@ void init_ops(nb::module_& m) {
bits (int, optional): The number of bits occupied by each element of
``w`` in the quantized array. See supported values and defaults in the
:ref:`table of quantization modes <quantize-modes>`. Default: ``None``.
global_scale (array, optional): The per-input float32 scale used for
``"nvfp4"`` quantization if provided. Default: ``None``.
dtype (Dtype, optional): The data type of the dequantized output. If
``None`` the return type is inferred from the scales and biases
when possible and otherwise defaults to ``bfloat16``.
@@ -5511,10 +5517,12 @@ void init_ops(nb::module_& m) {
"group_size"_a = nb::none(),
"bits"_a = nb::none(),
"mode"_a = "nvfp4",
"global_scale_x"_a = nb::none(),
"global_scale_w"_a = nb::none(),
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"),
"def qqmm(x: array, w: array, scales: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, 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
@@ -5532,6 +5540,7 @@ void init_ops(nb::module_& m) {
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.
``global_scale_x`` and ``global_scale_w`` are only used for ``nvfp4`` quantization.
Args:
x (array): Input array.
@@ -5547,7 +5556,10 @@ void init_ops(nb::module_& m) {
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.
global_scale (array, optional): The per-input float32 scale used for x
with ``"nvfp4"`` quantization. Default: ``None``.
global_scale_w (array, optional): The per-input float32 scale used for w
with ``"nvfp4"`` quantization. Default: ``None``.
Returns:
array: The result of the multiplication of quantized ``x`` with quantized ``w``.
needed).
python/tests/test_quantized.py
+13
View File
@@ -160,6 +160,19 @@ class TestQuantized(mlx_tests.MLXTestCase):
w_hat = mx.dequantize(w_q, scales, mode="nvfp4")
self.assertTrue(mx.all(w_hat == 0))
# Test nvfp4 quantize/dequantize with tensor-scale global_scale
# currently supported only on cpu and cuda
if not mx.metal.is_available():
global_scale = w.abs().max().astype(mx.float32)
else:
global_scale = None
w_q, scales = mx.quantize(w, mode="nvfp4", global_scale=global_scale)
w_hat = mx.dequantize(
w_q, scales, group_size=16, bits=4, mode="nvfp4", global_scale=global_scale
)
self.assertTrue(mx.allclose(w, w_hat, rtol=1e-5, atol=1e-5))
def test_qqmv(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)