From 9e8ef77b4280e2b5e2ddee2051c97b66230d525d Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 9 Jul 2026 11:43:26 -0700 Subject: [PATCH] up --- backends/mlx/builder/op_helpers.py | 6 +++--- examples/models/gemma4_31b/quant/quantize.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py index 85a71dc453f..98a9f779297 100644 --- a/backends/mlx/builder/op_helpers.py +++ b/backends/mlx/builder/op_helpers.py @@ -669,10 +669,10 @@ def parse_dequant_node( return None # MLX supports 2,3,4,5,6,8-bit affine quantization. to_mlx_qparams packs - # 2/4/8 via fast paths and other widths (e.g. 6) via a general contiguous - # bit-packer, so enable 6 here too. + # 2/4/8 via fast paths and other widths (e.g. 5, 6) via a general + # contiguous bit-packer, so enable 5 and 6 here too. bits = (qmax - qmin + 1).bit_length() - 1 - if bits not in [2, 4, 6, 8]: + if bits not in [2, 4, 5, 6, 8]: return None return qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim diff --git a/examples/models/gemma4_31b/quant/quantize.py b/examples/models/gemma4_31b/quant/quantize.py index f0d6d7215ea..5e2f54aa7dc 100644 --- a/examples/models/gemma4_31b/quant/quantize.py +++ b/examples/models/gemma4_31b/quant/quantize.py @@ -7,8 +7,8 @@ """Quantize weights to torchao tensor subclasses. ``quantize_weight`` quantizes a single tensor given a ``QuantConfig``, -returning an ``Int4Tensor`` (4-bit) or ``IntxUnpackedToInt8Tensor`` (6- or -8-bit). +returning an ``Int4Tensor`` (4-bit) or ``IntxUnpackedToInt8Tensor`` (5-, 6-, +or 8-bit). ``quantize_model`` walks a model's parameters, applies a ``QuantRecipe``, and returns a single state dict containing both quantized subclass tensors @@ -152,7 +152,7 @@ def _to_intx_tensor( weight: torch.Tensor, config: QuantConfig, ) -> torch.Tensor: - """Quantize to 6- or 8-bit and wrap in IntxUnpackedToInt8Tensor. + """Quantize to 5-, 6-, or 8-bit and wrap in IntxUnpackedToInt8Tensor. Quantizes in float32 for numerical precision, then constructs the subclass directly. We avoid ``from_hp`` because it quantizes in the @@ -229,10 +229,10 @@ def _to_intx_tensor( def quantize_weight(weight: torch.Tensor, config: QuantConfig) -> torch.Tensor: """Quantize ``weight`` to a torchao tensor subclass. - Returns ``Int4Tensor`` for 4-bit or ``IntxUnpackedToInt8Tensor`` for 6- - and 8-bit. + Returns ``Int4Tensor`` for 4-bit or ``IntxUnpackedToInt8Tensor`` for 5-, + 6-, and 8-bit. """ - if config.bits in (6, 8): + if config.bits in (5, 6, 8): return _to_intx_tensor(weight, config) if config.bits != 4: