From 32d61bc2fafda1fdc44cbe76f52e31a55f86a842 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 09:27:12 -0700 Subject: [PATCH 1/9] Add IQ quantization codecs and backend Signed-off-by: Hung-Yueh Chiang --- .pre-commit-config.yaml | 2 + LICENSE | 1 + modelopt/torch/quantization/__init__.py | 6 + modelopt/torch/quantization/ggml/__init__.py | 25 ++ modelopt/torch/quantization/ggml/backend.py | 35 +++ modelopt/torch/quantization/ggml/common.py | 59 ++++ modelopt/torch/quantization/ggml/iq1_s.py | 304 +++++++++++++++++++ modelopt/torch/quantization/ggml/iq2_xs.py | 302 ++++++++++++++++++ 8 files changed, 734 insertions(+) create mode 100644 modelopt/torch/quantization/ggml/__init__.py create mode 100644 modelopt/torch/quantization/ggml/backend.py create mode 100644 modelopt/torch/quantization/ggml/common.py create mode 100644 modelopt/torch/quantization/ggml/iq1_s.py create mode 100644 modelopt/torch/quantization/ggml/iq2_xs.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2f4fdd595e3..0202b9b2fec 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,6 +103,8 @@ repos: exclude: > (?x)^( modelopt/torch/quantization/utils/calib_utils.py| + modelopt/torch/quantization/ggml/iq1_s.py| + modelopt/torch/quantization/ggml/iq2_xs.py| modelopt/onnx/quantization/operators.py| modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| diff --git a/LICENSE b/LICENSE index c58bddda878..57c1104ae67 100644 --- a/LICENSE +++ b/LICENSE @@ -250,6 +250,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2023 DeepSeek Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors + Copyright (c) 2023-2026 The ggml authors Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py index 87dbf30bb57..5d069369ec2 100644 --- a/modelopt/torch/quantization/__init__.py +++ b/modelopt/torch/quantization/__init__.py @@ -15,6 +15,8 @@ """Quantization package.""" +from importlib import import_module as _import_module + # Initialize mode and plugins from . import mode, plugins, utils @@ -25,3 +27,7 @@ from .model_quant import * from .nn.modules.quant_module import QuantModuleRegistry from .utils import update_quant_cfg_with_kv_cache_quant + +# Loading this before the core imports above creates a cycle through quantization.qtensor. +ggml = _import_module(".ggml", __name__) +del _import_module diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py new file mode 100644 index 00000000000..0fbb63ee405 --- /dev/null +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GGML-compatible block quantization formats.""" + +# Importing the backend installs its TensorQuantizer dispatch entry. +from . import backend as _backend +from .iq1_s import * +from .iq1_s import __all__ as _iq1_s_all +from .iq2_xs import * +from .iq2_xs import __all__ as _iq2_xs_all + +__all__ = [*_iq1_s_all, *_iq2_xs_all] # noqa: PLE0604 diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py new file mode 100644 index 00000000000..95547659bb5 --- /dev/null +++ b/modelopt/torch/quantization/ggml/backend.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TensorQuantizer backend dispatch for GGML-compatible IQ formats.""" + +import torch + +from ..nn.modules.tensor_quantizer import register_quant_backend +from .iq1_s import iq1_s_fake_quant +from .iq2_xs import iq2_xs_fake_quant + + +def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """Dispatch an IQ quantizer to its format-specific implementation.""" + num_bits = getattr(quantizer, "num_bits", None) + if num_bits == "iq1_s": + return iq1_s_fake_quant(inputs, quantizer) + if num_bits == "iq2_xs": + return iq2_xs_fake_quant(inputs, quantizer) + raise ValueError("The ggml backend requires num_bits='iq1_s' or 'iq2_xs'") + + +register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py new file mode 100644 index 00000000000..ff977813ac1 --- /dev/null +++ b/modelopt/torch/quantization/ggml/common.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared validation for GGML-compatible block quantizers.""" + +import math + +import torch + +GGML_BLOCK_SIZE = 256 + + +def validate_weight(weight: torch.Tensor, format_name: str) -> None: + """Validate a weight accepted by the current GGML block encoders.""" + if weight.numel() == 0: + raise ValueError(f"{format_name} requires a non-empty weight") + if weight.dim() == 0 or weight.shape[-1] % GGML_BLOCK_SIZE: + raise ValueError( + f"{format_name} requires the last weight dimension to be divisible by " + f"{GGML_BLOCK_SIZE}, got shape {tuple(weight.shape)}" + ) + if not weight.is_floating_point(): + raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError(f"{format_name} requires finite weight values") + + +def validate_packed_weights( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + block_bytes: int, + format_name: str, +) -> tuple[int, ...]: + """Validate a packed payload and return its logical shape.""" + if packed_weights.dtype != torch.uint8 or packed_weights.shape[-1] != block_bytes: + raise ValueError( + f"packed_weights must be uint8 with last dimension {block_bytes}, " + f"got {packed_weights.dtype} {tuple(packed_weights.shape)}" + ) + shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist()) + if not shape or shape[-1] % GGML_BLOCK_SIZE: + raise ValueError(f"invalid {format_name} logical weight shape: {shape}") + expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes + if packed_weights.numel() != expected_payload_values: + raise ValueError("packed_weights size does not match weight_shape") + return shape diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py new file mode 100644 index 00000000000..b8aae0d3df3 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -0,0 +1,304 @@ +# This file includes the IQ1_S codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""IQ1_S fake quantization and GGML-compatible block packing. + +The encoder follows the canonical ``search_impl="auto"`` search. Every 256 +logical values become one 50-byte ``block_iq1_s`` payload: + +* bytes 0..1: little-endian FP16 super-block scale ``d`` +* bytes 2..33: low eight bits of 32 codebook indices +* bytes 34..49: eight little-endian uint16 metadata words + +Each metadata word describes four consecutive eight-value vectors. Bits 0..11 +hold the three high index bits, bits 12..14 select one of eight local scales, +and bit 15 selects the shared -0.125 rather than +0.125 delta. The canonical +2048 x 8 ternary grid below comes from llama.cpp ``ggml-common.h`` revision +9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +import zlib +from functools import cache + +import torch + +from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight + +__all__ = [ + "IQ1_S_BLOCK_BYTES", + "IQ1_S_BLOCK_SIZE", + "IQ1_S_EFFECTIVE_BITS", + "dequantize_iq1_s", + "iq1_s_fake_quant", + "iq1_s_grid", + "quantize_iq1_s", +] + +IQ1_S_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ1_S_BLOCK_BYTES = 50 +IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE +_IQ1_S_DELTA = 0.125 +_IQ1_S_NATIVE_MAX = 16.875 + +# zlib-compressed little-endian bytes of the canonical uint64_t table. The +# decoded int8 values are -1, 0, and 1. +_IQ1_S_GRID_ZLIB_B64 = ( + "eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv" + "c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY" + "4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX" + "bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8" + "Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W" + "YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH" + "9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL" + "2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E" + "UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg" + "MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0" + "jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+" + "LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs" + "fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2" + "a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g" + "n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq" + "ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ" + "q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/" + "1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2" + "kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM" + "T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8" + "0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT" + "oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy" + "i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44" + "rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg" + "O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO" + "k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI" + "Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8" + "73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn" + "gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6" + "hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+" + "EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA" + "0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH" + "OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD" + "FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn" + "h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6" + "DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB" + "MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI" + "4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh" + "eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ" + "zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep" + "V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5" + "/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C" + "qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T" + "rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD" + "hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr" + "x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO" + "umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0" + "Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA" + "MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ" + "jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8=" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64)) + + +def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ1_S ternary grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + raw = torch.tensor(list(_grid_bytes()), dtype=torch.uint8).view(torch.int8) + _GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to( + device=resolved_device, dtype=torch.float32 + ) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + xnorm = vectors.square().sum(dim=-1) + xsum = vectors.sum(dim=-1) + + amax = x.abs().amax(dim=1) + d = ((amax / _IQ1_S_NATIVE_MAX) * 0.61).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + best_error = torch.full((block_count, 32, 16), torch.inf, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + grid_norm = grid.square().sum(dim=-1) + grid_sum = grid.sum(dim=-1) + + # Tile the 2048-entry codebook to bound temporary memory. A strict update + # retains the lowest codebook index when two candidates have equal error. + for entry_start in range(0, 2048, 128): + grid_tile = grid[entry_start : entry_start + 128] + dot = torch.matmul(vectors, grid_tile.T) + tile_norm = grid_norm[entry_start : entry_start + 128].reshape(1, 1, -1) + tile_sum = grid_sum[entry_start : entry_start + 128].reshape(1, 1, -1) + + for shift in range(2): + delta = -_IQ1_S_DELTA if shift else _IQ1_S_DELTA + shifted_dot = dot + delta * xsum.unsqueeze(-1) + shifted_norm = tile_norm + 2 * delta * tile_sum + 8 * delta * delta + for local in range(8): + choice = shift * 8 + local + scale = d_float.reshape(-1, 1, 1) * (2 * local + 1) + error = ( + xnorm.unsqueeze(-1) - 2 * scale * shifted_dot + scale.square() * shifted_norm + ).clamp_min_(0) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, choice] + best_error[:, :, choice] = torch.where( + replace, tile_error, best_error[:, :, choice] + ) + best_entry[:, :, choice] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, choice] + ) + + group_error = best_error.reshape(block_count, 8, 4, 16).sum(dim=2) + selected_choice = group_error.argmin(dim=-1) + vector_choice = selected_choice.repeat_interleave(4, dim=1) + selected_entry = best_entry.gather(2, vector_choice.unsqueeze(-1)).squeeze(-1) + selected_local = selected_choice & 0x7 + selected_shift = selected_choice >> 3 + + high = (selected_entry >> 8).reshape(block_count, 8, 4) + qh = ( + high[:, :, 0] + | (high[:, :, 1] << 3) + | (high[:, :, 2] << 6) + | (high[:, :, 3] << 9) + | (selected_local << 12) + | (selected_shift << 15) + ) + + packed = torch.empty((block_count, IQ1_S_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:34] = (selected_entry & 0xFF).to(torch.uint8) + packed[:, 34:50:2] = (qh & 0xFF).to(torch.uint8) + packed[:, 35:50:2] = (qh >> 8).to(torch.uint8) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_iq1_s( + weight: torch.Tensor, *, block_chunk_size: int = 64 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ1_S blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 50]`` + and ``[weight.ndim]``. Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ1_S") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE) + grid = iq1_s_grid(weight.device) + if weight.is_cuda: + from .. import extensions + + get_extension = getattr(extensions, "get_cuda_ext_iq1_s", None) + extension = get_extension() if get_extension is not None else None + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ1_S_BLOCK_SIZE, + IQ1_S_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ1_S_BLOCK_SIZE, + IQ1_S_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq1_s( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ1_S payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ1_S_BLOCK_BYTES, format_name="IQ1_S" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = blocks[:, 2:34].to(torch.int64).reshape(-1, 8, 4) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) + high = (qh.unsqueeze(-1) >> shifts) & 0x7 + entries = low | (high << 8) + + local = (qh >> 12) & 0x7 + delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA) + values = iq1_s_grid(blocks.device)[entries] + delta.unsqueeze(-1).unsqueeze(-1) + scales = d.unsqueeze(-1) * (2 * local + 1).float() + decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + return decoded.reshape(shape).to(dtype) + + +def iq1_s_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ1_S backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq1_s": + raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ1_S search_impl='auto' is currently supported") + packed, shape = quantize_iq1_s(inputs) + reconstructed = dequantize_iq1_s(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py new file mode 100644 index 00000000000..73b26f61c64 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -0,0 +1,302 @@ +# This file includes the IQ2_XS codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""IQ2_XS fake quantization and GGML-compatible block packing. + +The encoder follows the canonical search_impl="auto" search. Every 256 +logical values become one 74-byte block_iq2_xs payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign) +* bytes 66..73: 16 four-bit local scales, two per byte + +The canonical 512 x 8 magnitude grid below comes from llama.cpp +ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +from functools import cache + +import torch + +from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight + +__all__ = [ + "IQ2_XS_BLOCK_BYTES", + "IQ2_XS_BLOCK_SIZE", + "IQ2_XS_EFFECTIVE_BITS", + "dequantize_iq2_xs", + "iq2_xs_fake_quant", + "iq2_xs_grid", + "quantize_iq2_xs", +] + +IQ2_XS_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ2_XS_BLOCK_BYTES = 74 +IQ2_XS_EFFECTIVE_BITS = IQ2_XS_BLOCK_BYTES * 8 / IQ2_XS_BLOCK_SIZE + +# Compact byte representation of the canonical [512, 8] grid. Values are only +# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids +# adding a pickle-backed torch.save artifact to the wheel. +_IQ2_XS_GRID_B64 = ( + "CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr" + "CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI" + "GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI" + "GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr" + "CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr" + "CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI" + "CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI" + "KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ" + "CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI" + "CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI" + "CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ" + "GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI" + "CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI" + "KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ" + "CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI" + "GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI" + "GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI" + "GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ" + "CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI" + "CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI" + "KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI" + "GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr" + "KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr" + "KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI" + "CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI" + "GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ" + "GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr" + "CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ" + "CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI" + "CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr" + "CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI" + "CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI" + "CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ" + "GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI" + "CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr" + "GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ" + "GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI" + "CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ" + "KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI" + "KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr" + "CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI" + "KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI" + "GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr" + "GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr" + "CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI" + "KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ" + "GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI" + "CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ" + "GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr" + "KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ" + "CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw==" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return base64.b64decode(_IQ2_XS_GRID_B64) + + +def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_XS magnitude grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + values = torch.tensor(list(_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + amax = x.abs().amax(dim=1) + rms = x.square().mean(dim=1).sqrt() + peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) + anchor_ratio = (1.0 - 0.035 * peak_to_rms).clamp(0.65, 0.92) + d = ((amax / 166.625) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + # Search the codebook in tiles to cap temporary memory. Strict comparison + # preserves the lowest grid index on equal error, matching the CUDA key. + for entry_start in range(0, 512, 64): + grid_tile = grid[entry_start : entry_start + 64] + products = magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8) + dot = products.sum(dim=-1) + dot = torch.where(odd_parity.unsqueeze(-1), dot - 2.0 * products.amin(dim=-1), dot) + tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1) + + for local in range(16): + scale = d_float.reshape(-1, 1, 1) * ((2 * local + 1) / 8.0) + error = ( + xnorm.unsqueeze(-1) - 2.0 * scale * dot + scale.square() * tile_qnorm + ).clamp_min_(0) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, local] + best_error[:, :, local] = torch.where(replace, tile_error, best_error[:, :, local]) + best_entry[:, :, local] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, local] + ) + + group_error = best_error.reshape(block_count, 16, 2, 16).sum(dim=2) + selected_local = group_error.argmin(dim=-1) + vector_local = selected_local.repeat_interleave(2, dim=1) + selected_entry = best_entry.gather(2, vector_local.unsqueeze(-1)).squeeze(-1) + + selected_grid = grid[selected_entry] + weakest_index = (magnitudes * selected_grid).argmin(dim=-1) + flip = torch.nn.functional.one_hot(weakest_index, num_classes=8).bool() + encoded_negative = negative ^ (flip & odd_parity.unsqueeze(-1)) + sign_bits = torch.arange(8, dtype=torch.int64, device=x.device) + sign_mask = (encoded_negative.to(torch.int64) << sign_bits).sum(dim=-1) + + codes = selected_entry | ((sign_mask & 0x7F) << 9) + packed = torch.empty((block_count, IQ2_XS_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:66:2] = (codes & 0xFF).to(torch.uint8) + packed[:, 3:66:2] = (codes >> 8).to(torch.uint8) + packed[:, 66:] = (selected_local[:, 0::2] | (selected_local[:, 1::2] << 4)).to(torch.uint8) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_iq2_xs( + weight: torch.Tensor, *, block_chunk_size: int = 64 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ2_XS blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 74]`` + and ``[weight.ndim]``. Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ2_XS") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE) + grid = iq2_xs_grid(weight.device) + if weight.is_cuda: + from .. import extensions + + get_extension = getattr(extensions, "get_cuda_ext_iq2_xs", None) + extension = get_extension() if get_extension is not None else None + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XS_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XS_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_xs( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ2_XS payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_XS_BLOCK_BYTES, format_name="IQ2_XS" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + entries = codes & 0x1FF + sign_index = codes >> 9 + + parity = torch.zeros_like(sign_index) + for bit in range(7): + parity ^= (sign_index >> bit) & 1 + sign_mask = sign_index | (parity << 7) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + + scale_bytes = blocks[:, 66:].to(torch.int64) + local = torch.empty((blocks.shape[0], 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 + values = iq2_xs_grid(blocks.device)[entries] * signs + decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) + return decoded.reshape(shape).to(dtype) + + +def iq2_xs_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ2_XS backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_xs": + raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ2_XS search_impl='auto' is currently supported") + packed, shape = quantize_iq2_xs(inputs) + reconstructed = dequantize_iq2_xs(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() From 887d357174b2c889af573a2f58351821347ecada Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 14:51:44 -0700 Subject: [PATCH 2/9] Test IQ quantization codecs Signed-off-by: Hung-Yueh Chiang --- .../gpu/torch/quantization/test_iq1_s_cuda.py | 68 +++++++++++ .../torch/quantization/test_iq2_xs_cuda.py | 75 ++++++++++++ tests/unit/torch/quantization/test_iq1_s.py | 107 ++++++++++++++++++ tests/unit/torch/quantization/test_iq2_xs.py | 94 +++++++++++++++ 4 files changed, 344 insertions(+) create mode 100644 tests/gpu/torch/quantization/test_iq1_s_cuda.py create mode 100644 tests/gpu/torch/quantization/test_iq2_xs_cuda.py create mode 100644 tests/unit/torch/quantization/test_iq1_s.py create mode 100644 tests/unit/torch/quantization/test_iq2_xs.py diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py new file mode 100644 index 00000000000..69c6a1d1460 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +import modelopt.torch.quantization.extensions as extensions +from modelopt.torch.quantization.extensions import get_cuda_ext_iq1_s +from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, iq1_s_grid, quantize_iq1_s + + +def _extension(): + extension = get_cuda_ext_iq1_s(raise_if_failed=True) + assert extension is not None + return extension + + +def test_iq1_s_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) + packed_again = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) + dispatched, shape = quantize_iq1_s(weight) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (8, 1, 50) + assert torch.equal(packed, packed_again) + assert torch.equal(packed, dispatched) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) + shape = torch.tensor(weight.shape, device="cuda") + + assert not packed.any() + assert torch.equal(dequantize_iq1_s(packed, shape), weight) + + +def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch): + monkeypatch.setattr(extensions, "get_cuda_ext_iq1_s", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + reconstructed = dequantize_iq1_s(packed, shape) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + + assert packed.shape == (2, 1, 50) + assert normalized_mse < 0.25 diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py new file mode 100644 index 00000000000..b8f07f9d6f6 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +import modelopt.torch.quantization.extensions as extensions +from modelopt.torch.quantization.extensions import get_cuda_ext_iq2_xs +from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, iq2_xs_grid, quantize_iq2_xs + + +def _extension(): + extension = get_cuda_ext_iq2_xs(raise_if_failed=True) + assert extension is not None + return extension + + +def test_iq2_xs_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) + packed_again = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) + dispatched, shape = quantize_iq2_xs(weight) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (8, 2, 74) + assert torch.equal(packed, packed_again) + assert torch.equal(packed, dispatched) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + +def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(1, 1, 74) + shape = torch.tensor(weight.shape, device="cuda") + + assert not packed.any() + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) + + +def test_iq2_xs_cuda_underflowed_scale_has_canonical_zero_encoding(): + weight = torch.full((1, 256), -1e-6, device="cuda", dtype=torch.bfloat16) + packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(1, 1, 74) + + assert not packed.any() + + +def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch): + monkeypatch.setattr(extensions, "get_cuda_ext_iq2_xs", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + reconstructed = dequantize_iq2_xs(packed, shape) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + + assert packed.shape == (2, 1, 74) + assert normalized_mse < 0.1 diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py new file mode 100644 index 00000000000..a623be608a4 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.ggml.iq1_s import ( + IQ1_S_BLOCK_BYTES, + dequantize_iq1_s, + iq1_s_fake_quant, + iq1_s_grid, + quantize_iq1_s, +) + + +def test_iq1_s_canonical_grid(): + grid = iq1_s_grid() + + assert grid.shape == (2048, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0} + assert grid[0].tolist() == [-1.0] * 8 + + +def test_iq1_s_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + + assert packed.shape == (2, 1, IQ1_S_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq1_s(packed, shape), weight) + + +def test_iq1_s_dequantizes_ggml_metadata_bit_fields(): + packed = torch.zeros((1, 1, 50), dtype=torch.uint8) + d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8) + packed[0, 0, :2] = d + entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64) + packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8) + qh = ( + ((entries[0] >> 8) & 7) + | (((entries[1] >> 8) & 7) << 3) + | (((entries[2] >> 8) & 7) << 6) + | (((entries[3] >> 8) & 7) << 9) + | (3 << 12) + | (1 << 15) + ) + packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8) + packed[0, 0, 35] = (qh >> 8).to(torch.uint8) + + decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32) + expected = (iq1_s_grid()[entries] - 0.125) * 14.0 + + assert torch.equal(decoded[0, :32].reshape(4, 8), expected) + + +def test_iq1_s_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight, block_chunk_size=1) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (2, 1, 50) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + blocks = packed.reshape(-1, 50) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + assert torch.all(((qh >> 12) & 0x7) < 8) + assert torch.all((qh & 0xFFF) < 0x1000) + + +def test_iq1_s_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="last weight dimension"): + quantize_iq1_s(torch.ones(2, 257)) + + +def test_iq1_s_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq1_s" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq1_s_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py new file mode 100644 index 00000000000..773e36e93c4 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.ggml.iq2_xs import ( + IQ2_XS_BLOCK_BYTES, + dequantize_iq2_xs, + iq2_xs_fake_quant, + iq2_xs_grid, + quantize_iq2_xs, +) + + +def test_iq2_xs_canonical_grid(): + grid = iq2_xs_grid() + + assert grid.shape == (512, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {8.0, 25.0, 43.0} + assert grid[0].tolist() == [8.0] * 8 + assert grid[-1].tolist() == [43.0] * 8 + + +def test_iq2_xs_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + + assert packed.shape == (2, 1, IQ2_XS_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) + + +def test_iq2_xs_underflowed_scale_has_canonical_zero_encoding(): + weight = torch.full((1, 256), -1e-6, dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + + assert not packed.any() + assert torch.equal(dequantize_iq2_xs(packed, shape), torch.zeros_like(weight)) + + +def test_iq2_xs_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight, block_chunk_size=2) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (2, 2, 74) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + blocks = packed.reshape(-1, 74) + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + assert torch.all((codes & 0x1FF) < 512) + assert torch.all((codes >> 9) < 128) + + +def test_iq2_xs_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="last weight dimension"): + quantize_iq2_xs(torch.ones(2, 257)) + + +def test_iq2_xs_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq2_xs" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq2_xs_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) From d40192f9ec419758b30ce9af403bf6fe30225279 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 16:29:31 -0700 Subject: [PATCH 3/9] Fix IQ codec review findings Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/quantization/__init__.py | 3 +- modelopt/torch/quantization/ggml/common.py | 4 +- modelopt/torch/quantization/ggml/iq1_s.py | 32 ++++++++------- modelopt/torch/quantization/ggml/iq2_xs.py | 40 ++++++++++++------- .../gpu/torch/quantization/test_iq1_s_cuda.py | 5 ++- .../torch/quantization/test_iq2_xs_cuda.py | 5 ++- tests/unit/torch/quantization/test_iq1_s.py | 16 +++++++- tests/unit/torch/quantization/test_iq2_xs.py | 14 ++++++- 8 files changed, 79 insertions(+), 40 deletions(-) diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py index 5d069369ec2..a597401a7bf 100644 --- a/modelopt/torch/quantization/__init__.py +++ b/modelopt/torch/quantization/__init__.py @@ -28,6 +28,7 @@ from .nn.modules.quant_module import QuantModuleRegistry from .utils import update_quant_cfg_with_kv_cache_quant -# Loading this before the core imports above creates a cycle through quantization.qtensor. +# Imported last to register the backend without cycling through quantization.qtensor. +# A dynamic import prevents isort from hoisting it into the import block above. ggml = _import_module(".ggml", __name__) del _import_module diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index ff977813ac1..e5d6b9c16a8 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -23,7 +23,7 @@ def validate_weight(weight: torch.Tensor, format_name: str) -> None: - """Validate a weight accepted by the current GGML block encoders.""" + """Validate weight metadata accepted by the current GGML block encoders.""" if weight.numel() == 0: raise ValueError(f"{format_name} requires a non-empty weight") if weight.dim() == 0 or weight.shape[-1] % GGML_BLOCK_SIZE: @@ -33,8 +33,6 @@ def validate_weight(weight: torch.Tensor, format_name: str) -> None: ) if not weight.is_floating_point(): raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") - if not torch.isfinite(weight).all(): - raise ValueError(f"{format_name} requires finite weight values") def validate_packed_weights( diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index b8aae0d3df3..942884fc61f 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -40,8 +40,10 @@ """IQ1_S fake quantization and GGML-compatible block packing. -The encoder follows the canonical ``search_impl="auto"`` search. Every 256 -logical values become one 50-byte ``block_iq1_s`` payload: +The encoder performs a single-pass squared-error grid search at a fixed, +empirically anchored super-block scale. It does not iteratively refine the +scale or apply importance weights. Every 256 logical values become one 50-byte +``block_iq1_s`` payload: * bytes 0..1: little-endian FP16 super-block scale ``d`` * bytes 2..33: low eight bits of 32 codebook indices @@ -60,6 +62,7 @@ import torch +from ..extensions import get_cuda_ext_iq1_s from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight __all__ = [ @@ -77,6 +80,9 @@ IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE _IQ1_S_DELTA = 0.125 _IQ1_S_NATIVE_MAX = 16.875 +_IQ1_S_SCALE_ANCHOR = 0.61 +# At 1024 blocks, each largest IQ1_S search temporary is about 16 MiB in FP32. +_DEFAULT_BLOCK_CHUNK_SIZE = 1024 # zlib-compressed little-endian bytes of the canonical uint64_t table. The # decoded int8 values are -1, 0, and 1. @@ -161,10 +167,12 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: xsum = vectors.sum(dim=-1) amax = x.abs().amax(dim=1) - d = ((amax / _IQ1_S_NATIVE_MAX) * 0.61).clamp(max=65504.0).to(torch.float16) + # The fixed-scale search favors a compressed super-block scale. This + # empirical anchor initializes d below the full-range value. + d = ((amax / _IQ1_S_NATIVE_MAX) * _IQ1_S_SCALE_ANCHOR).clamp(max=65504.0).to(torch.float16) d_float = d.float() - best_error = torch.full((block_count, 32, 16), torch.inf, device=x.device) + best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device) best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) grid_norm = grid.square().sum(dim=-1) grid_sum = grid.sum(dim=-1) @@ -223,25 +231,23 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: @torch.no_grad() def quantize_iq1_s( - weight: torch.Tensor, *, block_chunk_size: int = 64 + weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE ) -> tuple[torch.Tensor, torch.Tensor]: """Pack a floating-point weight into GGML-compatible IQ1_S blocks. Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 50]`` - and ``[weight.ndim]``. Both tensors remain on the weight's device. + and ``[weight.ndim]``. The packed payload remains on the weight's device; + the logical-shape metadata is kept on CPU. """ validate_weight(weight, "IQ1_S") if block_chunk_size <= 0: raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") - logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE) grid = iq1_s_grid(weight.device) if weight.is_cuda: - from .. import extensions - - get_extension = getattr(extensions, "get_cuda_ext_iq1_s", None) - extension = get_extension() if get_extension is not None else None + extension = get_cuda_ext_iq1_s() if extension is not None: packed = extension.pack(blocks, grid) packed_shape = ( @@ -295,10 +301,6 @@ def iq1_s_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: """IQ1_S backend for TensorQuantizer, with pass-through backward.""" if getattr(quantizer, "num_bits", None) != "iq1_s": raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'") - extra_args = getattr(quantizer, "backend_extra_args", None) or {} - search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) - if search_impl != "auto": - raise NotImplementedError("Only IQ1_S search_impl='auto' is currently supported") packed, shape = quantize_iq1_s(inputs) reconstructed = dequantize_iq1_s(packed, shape, dtype=inputs.dtype) return inputs + (reconstructed - inputs).detach() diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 73b26f61c64..83fdd5c293b 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -40,8 +40,10 @@ """IQ2_XS fake quantization and GGML-compatible block packing. -The encoder follows the canonical search_impl="auto" search. Every 256 -logical values become one 74-byte block_iq2_xs payload: +The encoder performs a single-pass squared-error grid search at a fixed, +empirically anchored super-block scale. It does not iteratively refine the +scale or apply importance weights. Every 256 logical values become one 74-byte +block_iq2_xs payload: * bytes 0..1: little-endian FP16 super-block scale d * bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign) @@ -49,6 +51,8 @@ The canonical 512 x 8 magnitude grid below comes from llama.cpp ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +The matching dequantization formula is in ggml-quants.c at the same revision: +https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-quants.c#L2516-L2538 """ import base64 @@ -56,6 +60,7 @@ import torch +from ..extensions import get_cuda_ext_iq2_xs from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight __all__ = [ @@ -71,6 +76,12 @@ IQ2_XS_BLOCK_SIZE = GGML_BLOCK_SIZE IQ2_XS_BLOCK_BYTES = 74 IQ2_XS_EFFECTIVE_BITS = IQ2_XS_BLOCK_BYTES * 8 / IQ2_XS_BLOCK_SIZE +_IQ2_XS_NATIVE_MAX = 43 * 31 / 8 +_IQ2_XS_SCALE_ANCHOR_MIN = 0.65 +_IQ2_XS_SCALE_ANCHOR_MAX = 0.92 +_IQ2_XS_PEAK_TO_RMS_TAPER = 0.035 +# At 256 blocks, the largest IQ2_XS search temporary is about 64 MiB in FP32. +_DEFAULT_BLOCK_CHUNK_SIZE = 256 # Compact byte representation of the canonical [512, 8] grid. Values are only # 8, 25, and 43. Keeping this as checkpoint-independent package data avoids @@ -162,8 +173,12 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: amax = x.abs().amax(dim=1) rms = x.square().mean(dim=1).sqrt() peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) - anchor_ratio = (1.0 - 0.035 * peak_to_rms).clamp(0.65, 0.92) - d = ((amax / 166.625) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + # The fixed-scale search favors a compressed super-block scale. This + # empirical predictor tapers the anchor for outlier-heavy blocks. + anchor_ratio = (1.0 - _IQ2_XS_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( + _IQ2_XS_SCALE_ANCHOR_MIN, _IQ2_XS_SCALE_ANCHOR_MAX + ) + d = ((amax / _IQ2_XS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) d_float = d.float() xnorm = vectors.square().sum(dim=-1) @@ -214,25 +229,23 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: @torch.no_grad() def quantize_iq2_xs( - weight: torch.Tensor, *, block_chunk_size: int = 64 + weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE ) -> tuple[torch.Tensor, torch.Tensor]: """Pack a floating-point weight into GGML-compatible IQ2_XS blocks. Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 74]`` - and ``[weight.ndim]``. Both tensors remain on the weight's device. + and ``[weight.ndim]``. The packed payload remains on the weight's device; + the logical-shape metadata is kept on CPU. """ validate_weight(weight, "IQ2_XS") if block_chunk_size <= 0: raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") - logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE) grid = iq2_xs_grid(weight.device) if weight.is_cuda: - from .. import extensions - - get_extension = getattr(extensions, "get_cuda_ext_iq2_xs", None) - extension = get_extension() if get_extension is not None else None + extension = get_cuda_ext_iq2_xs() if extension is not None: packed = extension.pack(blocks, grid) packed_shape = ( @@ -283,6 +296,7 @@ def dequantize_iq2_xs( local = torch.empty((blocks.shape[0], 16), dtype=torch.int64, device=blocks.device) local[:, 0::2] = scale_bytes & 0x0F local[:, 1::2] = scale_bytes >> 4 + # Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8. scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 values = iq2_xs_grid(blocks.device)[entries] * signs decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) @@ -293,10 +307,6 @@ def iq2_xs_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: """IQ2_XS backend for TensorQuantizer, with pass-through backward.""" if getattr(quantizer, "num_bits", None) != "iq2_xs": raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'") - extra_args = getattr(quantizer, "backend_extra_args", None) or {} - search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) - if search_impl != "auto": - raise NotImplementedError("Only IQ2_XS search_impl='auto' is currently supported") packed, shape = quantize_iq2_xs(inputs) reconstructed = dequantize_iq2_xs(packed, shape, dtype=inputs.dtype) return inputs + (reconstructed - inputs).detach() diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py index 69c6a1d1460..bad32a3c718 100644 --- a/tests/gpu/torch/quantization/test_iq1_s_cuda.py +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -15,7 +15,7 @@ import torch -import modelopt.torch.quantization.extensions as extensions +import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module from modelopt.torch.quantization.extensions import get_cuda_ext_iq1_s from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, iq1_s_grid, quantize_iq1_s @@ -38,6 +38,7 @@ def test_iq1_s_cuda_pack_is_deterministic_and_decodable(): assert packed.shape == (8, 1, 50) assert torch.equal(packed, packed_again) assert torch.equal(packed, dispatched) + assert shape.device.type == "cpu" normalized_mse = ( reconstructed.float() - weight.float() ).square().mean() / weight.float().square().mean() @@ -54,7 +55,7 @@ def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch): - monkeypatch.setattr(extensions, "get_cuda_ext_iq1_s", lambda: None) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py index b8f07f9d6f6..5775c52bbd7 100644 --- a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -15,7 +15,7 @@ import torch -import modelopt.torch.quantization.extensions as extensions +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module from modelopt.torch.quantization.extensions import get_cuda_ext_iq2_xs from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, iq2_xs_grid, quantize_iq2_xs @@ -38,6 +38,7 @@ def test_iq2_xs_cuda_pack_is_deterministic_and_decodable(): assert packed.shape == (8, 2, 74) assert torch.equal(packed, packed_again) assert torch.equal(packed, dispatched) + assert shape.device.type == "cpu" normalized_mse = ( reconstructed.float() - weight.float() ).square().mean() / weight.float().square().mean() @@ -61,7 +62,7 @@ def test_iq2_xs_cuda_underflowed_scale_has_canonical_zero_encoding(): def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch): - monkeypatch.setattr(extensions, "get_cuda_ext_iq2_xs", lambda: None) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py index a623be608a4..cf332b48e52 100644 --- a/tests/unit/torch/quantization/test_iq1_s.py +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -90,6 +90,21 @@ def test_iq1_s_round_trip_and_payload_fields(): assert torch.all((qh & 0xFFF) < 0x1000) +def test_iq1_s_search_is_independent_of_default_dtype(): + generator = torch.Generator().manual_seed(0) + weight = torch.randn((8, 256), generator=generator, dtype=torch.float32) + expected, _ = quantize_iq1_s(weight) + + default_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + actual, _ = quantize_iq1_s(weight) + finally: + torch.set_default_dtype(default_dtype) + + assert torch.equal(actual, expected) + + def test_iq1_s_requires_complete_last_dimension_blocks(): with pytest.raises(ValueError, match="last weight dimension"): quantize_iq1_s(torch.ones(2, 257)) @@ -98,7 +113,6 @@ def test_iq1_s_requires_complete_last_dimension_blocks(): def test_iq1_s_fake_quant_has_pass_through_gradient(): class Quantizer: num_bits = "iq1_s" - backend_extra_args = {"search_impl": "auto"} weight = torch.randn(1, 256, requires_grad=True) output = iq1_s_fake_quant(weight, Quantizer()) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py index 773e36e93c4..408db1e457c 100644 --- a/tests/unit/torch/quantization/test_iq2_xs.py +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -77,6 +77,19 @@ def test_iq2_xs_round_trip_and_payload_fields(): assert torch.all((codes >> 9) < 128) +def test_iq2_xs_dequantizes_pinned_scale_factor(): + packed = torch.zeros((1, 1, 74), dtype=torch.uint8) + packed[0, 0, :2] = torch.tensor([1.0], dtype=torch.float16).view(torch.uint8) + packed[0, 0, 2:66:2] = 0xFF + packed[0, 0, 3:66:2] = 0x01 + packed[0, 0, 66:] = 0xFF + + decoded = dequantize_iq2_xs(packed, torch.tensor([1, 256]), dtype=torch.float32) + + # Entry 511 contains eight 43s and local code 15 gives (2 * 15 + 1) / 8. + assert torch.equal(decoded, torch.full((1, 256), 43 * 31 / 8, dtype=torch.float32)) + + def test_iq2_xs_requires_complete_last_dimension_blocks(): with pytest.raises(ValueError, match="last weight dimension"): quantize_iq2_xs(torch.ones(2, 257)) @@ -85,7 +98,6 @@ def test_iq2_xs_requires_complete_last_dimension_blocks(): def test_iq2_xs_fake_quant_has_pass_through_gradient(): class Quantizer: num_bits = "iq2_xs" - backend_extra_args = {"search_impl": "auto"} weight = torch.randn(1, 256, requires_grad=True) output = iq2_xs_fake_quant(weight, Quantizer()) From 9af3c13c2aa6ad2d3d533de80316a9f93b4da178 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 16:49:23 -0700 Subject: [PATCH 4/9] Test IQ backend routing and CUDA parity Signed-off-by: Hung-Yueh Chiang --- .../gpu/torch/quantization/test_iq1_s_cuda.py | 7 +-- .../torch/quantization/test_iq2_xs_cuda.py | 7 +-- .../torch/quantization/test_ggml_backend.py | 53 +++++++++++++++++++ 3 files changed, 61 insertions(+), 6 deletions(-) create mode 100644 tests/unit/torch/quantization/test_ggml_backend.py diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py index bad32a3c718..7b27e81998a 100644 --- a/tests/gpu/torch/quantization/test_iq1_s_cuda.py +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -26,18 +26,19 @@ def _extension(): return extension -def test_iq1_s_cuda_pack_is_deterministic_and_decodable(): +def test_iq1_s_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16) packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) packed_again = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) - dispatched, shape = quantize_iq1_s(weight) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) + reference, shape = quantize_iq1_s(weight) reconstructed = dequantize_iq1_s(packed, shape) assert packed.shape == (8, 1, 50) assert torch.equal(packed, packed_again) - assert torch.equal(packed, dispatched) + assert torch.equal(packed, reference) assert shape.device.type == "cpu" normalized_mse = ( reconstructed.float() - weight.float() diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py index 5775c52bbd7..cb429c8a9f2 100644 --- a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -26,18 +26,19 @@ def _extension(): return extension -def test_iq2_xs_cuda_pack_is_deterministic_and_decodable(): +def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) packed_again = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) - dispatched, shape = quantize_iq2_xs(weight) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) + reference, shape = quantize_iq2_xs(weight) reconstructed = dequantize_iq2_xs(packed, shape) assert packed.shape == (8, 2, 74) assert torch.equal(packed, packed_again) - assert torch.equal(packed, dispatched) + assert torch.equal(packed, reference) assert shape.device.type == "cpu" normalized_mse = ( reconstructed.float() - weight.float() diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py new file mode 100644 index 00000000000..9dd80e7791c --- /dev/null +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.ggml.backend import ggml_fake_quant + + +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_ggml_backend_via_quantize(num_bits): + torch.manual_seed(1234) + model = torch.nn.Linear(256, 2, bias=False) + inputs = torch.randn(2, 256) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": num_bits, "backend": "ggml"}, + "enable": True, + }, + ], + "algorithm": "max", + } + + mtq.quantize(model, config, forward_loop=lambda module: module(inputs)) + output = model(inputs) + + assert model.weight_quantizer.backend == "ggml" + assert model.weight_quantizer.num_bits == num_bits + assert output.shape == (2, 2) + assert torch.isfinite(output).all() + + +def test_ggml_backend_rejects_unknown_format(): + with pytest.raises(ValueError, match="requires num_bits"): + ggml_fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits="unknown")) From e8d937081d8cd01cf8e44d43915df443b79deb17 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 18:59:19 -0700 Subject: [PATCH 5/9] Address IQ codec review findings Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/quantization/__init__.py | 1 + modelopt/torch/quantization/ggml/backend.py | 10 ++- modelopt/torch/quantization/ggml/common.py | 79 +++++++++++++++++- modelopt/torch/quantization/ggml/iq1_s.py | 36 +++++++-- modelopt/torch/quantization/ggml/iq2_xs.py | 80 +++++++++++++------ .../gpu/torch/quantization/test_iq1_s_cuda.py | 11 +++ .../torch/quantization/test_iq2_xs_cuda.py | 25 +++++- .../torch/quantization/test_ggml_backend.py | 66 +++++++++++++++ tests/unit/torch/quantization/test_iq1_s.py | 41 ++++++++++ tests/unit/torch/quantization/test_iq2_xs.py | 41 ++++++++++ 10 files changed, 350 insertions(+), 40 deletions(-) diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py index a597401a7bf..29f42e8f72b 100644 --- a/modelopt/torch/quantization/__init__.py +++ b/modelopt/torch/quantization/__init__.py @@ -31,4 +31,5 @@ # Imported last to register the backend without cycling through quantization.qtensor. # A dynamic import prevents isort from hoisting it into the import block above. ggml = _import_module(".ggml", __name__) +globals().update({name: getattr(ggml, name) for name in ggml.__all__}) del _import_module diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py index 95547659bb5..97da3332398 100644 --- a/modelopt/torch/quantization/ggml/backend.py +++ b/modelopt/torch/quantization/ggml/backend.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TensorQuantizer backend dispatch for GGML-compatible IQ formats.""" +"""TensorQuantizer backend dispatch for GGML-compatible weight-only IQ formats.""" import torch @@ -25,10 +25,14 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: """Dispatch an IQ quantizer to its format-specific implementation.""" num_bits = getattr(quantizer, "num_bits", None) + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + unknown_args = set(extra_args) - {"block_chunk_size"} + if unknown_args: + raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}") if num_bits == "iq1_s": - return iq1_s_fake_quant(inputs, quantizer) + return iq1_s_fake_quant(inputs, quantizer, **extra_args) if num_bits == "iq2_xs": - return iq2_xs_fake_quant(inputs, quantizer) + return iq2_xs_fake_quant(inputs, quantizer, **extra_args) raise ValueError("The ggml backend requires num_bits='iq1_s' or 'iq2_xs'") diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index e5d6b9c16a8..175e1ce1419 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -16,12 +16,80 @@ """Shared validation for GGML-compatible block quantizers.""" import math +import weakref +from collections.abc import Callable +from dataclasses import dataclass import torch GGML_BLOCK_SIZE = 256 +@dataclass +class _PackedWeightCache: + input_ref: weakref.ReferenceType + input_key: tuple[object, ...] + format_name: str + block_chunk_size: int + packed_weights: torch.Tensor + weight_shape: torch.Tensor + + +def _input_cache_key(inputs: torch.Tensor) -> tuple[object, ...] | None: + try: + version = inputs._version + except RuntimeError: + # Inference tensors can omit version counters, so changes cannot be detected safely. + return None + return ( + inputs.data_ptr(), + tuple(inputs.shape), + tuple(inputs.stride()), + inputs.dtype, + inputs.device, + version, + ) + + +def fake_quantize_with_cache( + inputs: torch.Tensor, + quantizer, + *, + format_name: str, + block_chunk_size: int, + quantize: Callable[..., tuple[torch.Tensor, torch.Tensor]], + dequantize: Callable[..., torch.Tensor], +) -> torch.Tensor: + """Fake-quantize a weight while caching its compact packed representation.""" + input_key = _input_cache_key(inputs) + cache = getattr(quantizer, "_quantizer_cache", None) + if ( + isinstance(cache, _PackedWeightCache) + and input_key is not None + and cache.input_ref() is inputs + and cache.input_key == input_key + and cache.format_name == format_name + and cache.block_chunk_size == block_chunk_size + ): + packed_weights, weight_shape = cache.packed_weights, cache.weight_shape + else: + packed_weights, weight_shape = quantize(inputs, block_chunk_size=block_chunk_size) + if input_key is not None: + quantizer._quantizer_cache = _PackedWeightCache( + input_ref=weakref.ref(inputs), + input_key=input_key, + format_name=format_name, + block_chunk_size=block_chunk_size, + packed_weights=packed_weights, + weight_shape=weight_shape, + ) + else: + quantizer._quantizer_cache = None + + reconstructed = dequantize(packed_weights, weight_shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() + + def validate_weight(weight: torch.Tensor, format_name: str) -> None: """Validate weight metadata accepted by the current GGML block encoders.""" if weight.numel() == 0: @@ -43,13 +111,20 @@ def validate_packed_weights( format_name: str, ) -> tuple[int, ...]: """Validate a packed payload and return its logical shape.""" - if packed_weights.dtype != torch.uint8 or packed_weights.shape[-1] != block_bytes: + if ( + packed_weights.dim() == 0 + or packed_weights.dtype != torch.uint8 + or packed_weights.shape[-1] != block_bytes + ): raise ValueError( f"packed_weights must be uint8 with last dimension {block_bytes}, " f"got {packed_weights.dtype} {tuple(packed_weights.shape)}" ) + integral_dtypes = {torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64} + if weight_shape.dim() != 1 or weight_shape.dtype not in integral_dtypes: + raise ValueError("weight_shape must be a one-dimensional integral tensor") shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist()) - if not shape or shape[-1] % GGML_BLOCK_SIZE: + if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % GGML_BLOCK_SIZE: raise ValueError(f"invalid {format_name} logical weight shape: {shape}") expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes if packed_weights.numel() != expected_payload_values: diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 942884fc61f..04eda2d032f 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -63,7 +63,12 @@ import torch from ..extensions import get_cuda_ext_iq1_s -from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight +from .common import ( + GGML_BLOCK_SIZE, + fake_quantize_with_cache, + validate_packed_weights, + validate_weight, +) __all__ = [ "IQ1_S_BLOCK_BYTES", @@ -150,6 +155,8 @@ def _grid_bytes() -> bytes: def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: """Return the canonical IQ1_S ternary grid as float32.""" resolved_device = torch.device(device or "cpu") + if resolved_device.type == "cuda" and resolved_device.index is None: + resolved_device = torch.device("cuda", torch.cuda.current_device()) if resolved_device not in _GRID_CACHE: raw = torch.tensor(list(_grid_bytes()), dtype=torch.uint8).view(torch.int8) _GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to( @@ -160,7 +167,7 @@ def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: """Encode a moderate-size batch of flattened 256-value blocks.""" - x = blocks.float() + x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) block_count = x.shape[0] vectors = x.reshape(block_count, 32, 8) xnorm = vectors.square().sum(dim=-1) @@ -237,9 +244,12 @@ def quantize_iq1_s( Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 50]`` and ``[weight.ndim]``. The packed payload remains on the weight's device; - the logical-shape metadata is kept on CPU. + the logical-shape metadata is kept on CPU. Non-finite input elements are + treated as zero during packing. """ validate_weight(weight, "IQ1_S") + if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int): + raise TypeError("block_chunk_size must be an integer") if block_chunk_size <= 0: raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") @@ -297,10 +307,20 @@ def dequantize_iq1_s( return decoded.reshape(shape).to(dtype) -def iq1_s_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: - """IQ1_S backend for TensorQuantizer, with pass-through backward.""" +def iq1_s_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, +) -> torch.Tensor: + """IQ1_S weight backend for TensorQuantizer, with pass-through backward.""" if getattr(quantizer, "num_bits", None) != "iq1_s": raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'") - packed, shape = quantize_iq1_s(inputs) - reconstructed = dequantize_iq1_s(packed, shape, dtype=inputs.dtype) - return inputs + (reconstructed - inputs).detach() + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq1_s", + block_chunk_size=block_chunk_size, + quantize=quantize_iq1_s, + dequantize=dequantize_iq1_s, + ) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 83fdd5c293b..e78760ec1fc 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -61,7 +61,12 @@ import torch from ..extensions import get_cuda_ext_iq2_xs -from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight +from .common import ( + GGML_BLOCK_SIZE, + fake_quantize_with_cache, + validate_packed_weights, + validate_weight, +) __all__ = [ "IQ2_XS_BLOCK_BYTES", @@ -80,8 +85,9 @@ _IQ2_XS_SCALE_ANCHOR_MIN = 0.65 _IQ2_XS_SCALE_ANCHOR_MAX = 0.92 _IQ2_XS_PEAK_TO_RMS_TAPER = 0.035 -# At 256 blocks, the largest IQ2_XS search temporary is about 64 MiB in FP32. +# At 256 blocks, the largest IQ2_XS search temporary is about 16 MiB in FP32. _DEFAULT_BLOCK_CHUNK_SIZE = 256 +_SCALE_BLOCK_CHUNK_SIZE = 4096 # Compact byte representation of the canonical [512, 8] grid. Values are only # 8, 25, and 43. Keeping this as checkpoint-independent package data avoids @@ -155,21 +161,17 @@ def _grid_bytes() -> bytes: def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: """Return the canonical IQ2_XS magnitude grid as float32.""" resolved_device = torch.device(device or "cpu") + if resolved_device.type == "cuda" and resolved_device.index is None: + resolved_device = torch.device("cuda", torch.cuda.current_device()) if resolved_device not in _GRID_CACHE: values = torch.tensor(list(_grid_bytes()), dtype=torch.float32) _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) return _GRID_CACHE[resolved_device] -def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: - """Encode a moderate-size batch of flattened 256-value blocks.""" - x = blocks.float() - block_count = x.shape[0] - vectors = x.reshape(block_count, 32, 8) - magnitudes = vectors.abs() - negative = vectors < 0 - odd_parity = negative.sum(dim=-1).remainder(2).bool() - +def _predict_iq2_xs_scales(blocks: torch.Tensor) -> torch.Tensor: + """Predict one FP16 super-block scale for each flattened block.""" + x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) amax = x.abs().amax(dim=1) rms = x.square().mean(dim=1).sqrt() peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) @@ -178,7 +180,21 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: anchor_ratio = (1.0 - _IQ2_XS_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( _IQ2_XS_SCALE_ANCHOR_MIN, _IQ2_XS_SCALE_ANCHOR_MAX ) - d = ((amax / _IQ2_XS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + return ((amax / _IQ2_XS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + + +def _encode_blocks( + blocks: torch.Tensor, grid: torch.Tensor, scales: torch.Tensor | None = None +) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + d = _predict_iq2_xs_scales(x) if scales is None else scales d_float = d.float() xnorm = vectors.square().sum(dim=-1) @@ -235,9 +251,12 @@ def quantize_iq2_xs( Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 74]`` and ``[weight.ndim]``. The packed payload remains on the weight's device; - the logical-shape metadata is kept on CPU. + the logical-shape metadata is kept on CPU. Non-finite input elements are + treated as zero during packing. """ validate_weight(weight, "IQ2_XS") + if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int): + raise TypeError("block_chunk_size must be an integer") if block_chunk_size <= 0: raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") @@ -247,7 +266,11 @@ def quantize_iq2_xs( if weight.is_cuda: extension = get_cuda_ext_iq2_xs() if extension is not None: - packed = extension.pack(blocks, grid) + scale_chunks = [ + _predict_iq2_xs_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) + for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) + ] + packed = extension.pack(blocks, grid, torch.cat(scale_chunks)) packed_shape = ( *weight.shape[:-1], weight.shape[-1] // IQ2_XS_BLOCK_SIZE, @@ -255,10 +278,11 @@ def quantize_iq2_xs( ) return packed.reshape(packed_shape), logical_shape - chunks = [ - _encode_blocks(blocks[start : start + block_chunk_size], grid) - for start in range(0, blocks.shape[0], block_chunk_size) - ] + chunks = [] + for start in range(0, blocks.shape[0], block_chunk_size): + block_chunk = blocks[start : start + block_chunk_size] + scales = _predict_iq2_xs_scales(block_chunk) + chunks.append(_encode_blocks(block_chunk, grid, scales)) packed_shape = ( *weight.shape[:-1], weight.shape[-1] // IQ2_XS_BLOCK_SIZE, @@ -303,10 +327,20 @@ def dequantize_iq2_xs( return decoded.reshape(shape).to(dtype) -def iq2_xs_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: - """IQ2_XS backend for TensorQuantizer, with pass-through backward.""" +def iq2_xs_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, +) -> torch.Tensor: + """IQ2_XS weight backend for TensorQuantizer, with pass-through backward.""" if getattr(quantizer, "num_bits", None) != "iq2_xs": raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'") - packed, shape = quantize_iq2_xs(inputs) - reconstructed = dequantize_iq2_xs(packed, shape, dtype=inputs.dtype) - return inputs + (reconstructed - inputs).detach() + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq2_xs", + block_chunk_size=block_chunk_size, + quantize=quantize_iq2_xs, + dequantize=dequantize_iq2_xs, + ) diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py index 7b27e81998a..2716bd51e27 100644 --- a/tests/gpu/torch/quantization/test_iq1_s_cuda.py +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -55,6 +55,17 @@ def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): assert torch.equal(dequantize_iq1_s(packed, shape), weight) +def test_iq1_s_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch): + weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") + + packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) + reference, _ = quantize_iq1_s(weight) + + assert torch.equal(packed, reference) + + def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch): monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py index cb429c8a9f2..b6450bc1071 100644 --- a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -26,12 +26,18 @@ def _extension(): return extension +def _pack(weight): + blocks = weight.contiguous().reshape(-1, 256) + scales = iq2_xs_module._predict_iq2_xs_scales(blocks) + return _extension().pack(weight, iq2_xs_grid("cuda"), scales) + + def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) - packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) - packed_again = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(8, 2, 74) + packed = _pack(weight).reshape(8, 2, 74) + packed_again = _pack(weight).reshape(8, 2, 74) monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) reference, shape = quantize_iq2_xs(weight) reconstructed = dequantize_iq2_xs(packed, shape) @@ -48,7 +54,7 @@ def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout(): weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) - packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(1, 1, 74) + packed = _pack(weight).reshape(1, 1, 74) shape = torch.tensor(weight.shape, device="cuda") assert not packed.any() @@ -57,11 +63,22 @@ def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout(): def test_iq2_xs_cuda_underflowed_scale_has_canonical_zero_encoding(): weight = torch.full((1, 256), -1e-6, device="cuda", dtype=torch.bfloat16) - packed = _extension().pack(weight, iq2_xs_grid("cuda")).reshape(1, 1, 74) + packed = _pack(weight).reshape(1, 1, 74) assert not packed.any() +def test_iq2_xs_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch): + weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") + + packed = _pack(weight).reshape(1, 1, 74) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) + reference, _ = quantize_iq2_xs(weight) + + assert torch.equal(packed, reference) + + def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch): monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 9dd80e7791c..6aade467791 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -19,6 +19,9 @@ import torch import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.ggml.backend as backend_module +import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module from modelopt.torch.quantization.ggml.backend import ggml_fake_quant @@ -51,3 +54,66 @@ def test_ggml_backend_via_quantize(num_bits): def test_ggml_backend_rejects_unknown_format(): with pytest.raises(ValueError, match="requires num_bits"): ggml_fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits="unknown")) + + +def test_ggml_codecs_are_exported_from_quantization_package(): + assert mtq.quantize_iq1_s is iq1_s_module.quantize_iq1_s + assert mtq.quantize_iq2_xs is iq2_xs_module.quantize_iq2_xs + + +def test_ggml_backend_forwards_block_chunk_size(monkeypatch): + received = {} + + def fake_quant(inputs, _quantizer, *, block_chunk_size): + received["block_chunk_size"] = block_chunk_size + return inputs + + monkeypatch.setattr(backend_module, "iq1_s_fake_quant", fake_quant) + inputs = torch.ones(1, 256) + quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"block_chunk_size": 17}) + + assert ggml_fake_quant(inputs, quantizer) is inputs + assert received == {"block_chunk_size": 17} + + +def test_ggml_backend_rejects_unknown_extra_arg(): + quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"unknown": 1}) + + with pytest.raises(ValueError, match="Unsupported ggml backend_extra_args"): + ggml_fake_quant(torch.ones(1, 256), quantizer) + + +@pytest.mark.parametrize( + ("num_bits", "module", "fake_quant_name", "quantize_name"), + [ + ("iq1_s", iq1_s_module, "iq1_s_fake_quant", "quantize_iq1_s"), + ("iq2_xs", iq2_xs_module, "iq2_xs_fake_quant", "quantize_iq2_xs"), + ], +) +def test_ggml_backend_caches_packed_weight_and_invalidates_on_change( + monkeypatch, num_bits, module, fake_quant_name, quantize_name +): + weight = torch.randn(1, 256) + quantizer = SimpleNamespace(num_bits=num_bits, _quantizer_cache=None) + original_quantize = getattr(module, quantize_name) + call_count = 0 + + def counted_quantize(*args, **kwargs): + nonlocal call_count + call_count += 1 + return original_quantize(*args, **kwargs) + + monkeypatch.setattr(module, quantize_name, counted_quantize) + fake_quant = getattr(module, fake_quant_name) + + fake_quant(weight, quantizer, block_chunk_size=1) + fake_quant(weight, quantizer, block_chunk_size=1) + assert call_count == 1 + + fake_quant(weight, quantizer, block_chunk_size=2) + assert call_count == 2 + + with torch.no_grad(): + weight.add_(0.01) + fake_quant(weight, quantizer, block_chunk_size=2) + assert call_count == 3 diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py index cf332b48e52..2d12934b653 100644 --- a/tests/unit/torch/quantization/test_iq1_s.py +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -16,6 +16,7 @@ import pytest import torch +import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module from modelopt.torch.quantization.ggml.iq1_s import ( IQ1_S_BLOCK_BYTES, dequantize_iq1_s, @@ -34,6 +35,15 @@ def test_iq1_s_canonical_grid(): assert grid[0].tolist() == [-1.0] * 8 +def test_iq1_s_grid_normalizes_unindexed_cuda_device(monkeypatch): + cached = torch.empty(0) + indexed_device = torch.device("cuda", 7) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 7) + monkeypatch.setitem(iq1_s_module._GRID_CACHE, indexed_device, cached) + + assert iq1_s_grid("cuda") is cached + + def test_iq1_s_zero_block_has_canonical_zero_encoding(): weight = torch.zeros((2, 256), dtype=torch.bfloat16) @@ -110,6 +120,37 @@ def test_iq1_s_requires_complete_last_dimension_blocks(): quantize_iq1_s(torch.ones(2, 257)) +def test_iq1_s_treats_nonfinite_values_as_zero(): + weight = torch.randn(1, 256) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf]) + + packed, _ = quantize_iq1_s(weight) + expected, _ = quantize_iq1_s(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0)) + + assert torch.equal(packed, expected) + + +@pytest.mark.parametrize( + "weight_shape", + [ + torch.tensor(256), + torch.tensor([[1, 256]]), + torch.tensor([1.0, 256.0]), + torch.tensor([0, 256]), + ], +) +def test_iq1_s_rejects_invalid_shape_metadata(weight_shape): + packed = torch.zeros((1, 1, 50), dtype=torch.uint8) + + with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"): + dequantize_iq1_s(packed, weight_shape) + + +def test_iq1_s_rejects_scalar_packed_payload(): + with pytest.raises(ValueError, match="packed_weights"): + dequantize_iq1_s(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256])) + + def test_iq1_s_fake_quant_has_pass_through_gradient(): class Quantizer: num_bits = "iq1_s" diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py index 408db1e457c..5b16542737a 100644 --- a/tests/unit/torch/quantization/test_iq2_xs.py +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -16,6 +16,7 @@ import pytest import torch +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module from modelopt.torch.quantization.ggml.iq2_xs import ( IQ2_XS_BLOCK_BYTES, dequantize_iq2_xs, @@ -35,6 +36,15 @@ def test_iq2_xs_canonical_grid(): assert grid[-1].tolist() == [43.0] * 8 +def test_iq2_xs_grid_normalizes_unindexed_cuda_device(monkeypatch): + cached = torch.empty(0) + indexed_device = torch.device("cuda", 7) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 7) + monkeypatch.setitem(iq2_xs_module._GRID_CACHE, indexed_device, cached) + + assert iq2_xs_grid("cuda") is cached + + def test_iq2_xs_zero_block_has_canonical_zero_encoding(): weight = torch.zeros((2, 256), dtype=torch.bfloat16) @@ -95,6 +105,37 @@ def test_iq2_xs_requires_complete_last_dimension_blocks(): quantize_iq2_xs(torch.ones(2, 257)) +def test_iq2_xs_treats_nonfinite_values_as_zero(): + weight = torch.randn(1, 256) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf]) + + packed, _ = quantize_iq2_xs(weight) + expected, _ = quantize_iq2_xs(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0)) + + assert torch.equal(packed, expected) + + +@pytest.mark.parametrize( + "weight_shape", + [ + torch.tensor(256), + torch.tensor([[1, 256]]), + torch.tensor([1.0, 256.0]), + torch.tensor([0, 256]), + ], +) +def test_iq2_xs_rejects_invalid_shape_metadata(weight_shape): + packed = torch.zeros((1, 1, 74), dtype=torch.uint8) + + with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"): + dequantize_iq2_xs(packed, weight_shape) + + +def test_iq2_xs_rejects_scalar_packed_payload(): + with pytest.raises(ValueError, match="packed_weights"): + dequantize_iq2_xs(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256])) + + def test_iq2_xs_fake_quant_has_pass_through_gradient(): class Quantizer: num_bits = "iq2_xs" From 5b465df27da2cb67c20fc9772ccfcb2d6070d19a Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Thu, 17 Sep 2026 11:53:02 -0700 Subject: [PATCH 6/9] [OMNIML-5899] Bound IQ fake-quant decode memory Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/quantization/ggml/common.py | 15 ++++- modelopt/torch/quantization/ggml/iq1_s.py | 39 +++++++------ modelopt/torch/quantization/ggml/iq2_xs.py | 55 +++++++++++-------- .../torch/quantization/test_ggml_backend.py | 2 + tests/unit/torch/quantization/test_iq1_s.py | 4 +- tests/unit/torch/quantization/test_iq2_xs.py | 4 +- 6 files changed, 76 insertions(+), 43 deletions(-) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 175e1ce1419..88689e6e087 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -86,7 +86,12 @@ def fake_quantize_with_cache( else: quantizer._quantizer_cache = None - reconstructed = dequantize(packed_weights, weight_shape, dtype=inputs.dtype) + reconstructed = dequantize( + packed_weights, + weight_shape, + dtype=inputs.dtype, + block_chunk_size=block_chunk_size, + ) return inputs + (reconstructed - inputs).detach() @@ -103,6 +108,14 @@ def validate_weight(weight: torch.Tensor, format_name: str) -> None: raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") +def validate_block_chunk_size(block_chunk_size: int) -> None: + """Validate the common encoder and decoder block-chunk limit.""" + if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int): + raise TypeError("block_chunk_size must be an integer") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + def validate_packed_weights( packed_weights: torch.Tensor, weight_shape: torch.Tensor, diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 04eda2d032f..117660f704f 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -66,6 +66,7 @@ from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, + validate_block_chunk_size, validate_packed_weights, validate_weight, ) @@ -248,10 +249,7 @@ def quantize_iq1_s( treated as zero during packing. """ validate_weight(weight, "IQ1_S") - if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int): - raise TypeError("block_chunk_size must be an integer") - if block_chunk_size <= 0: - raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + validate_block_chunk_size(block_chunk_size) logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE) @@ -285,26 +283,35 @@ def dequantize_iq1_s( weight_shape: torch.Tensor, *, dtype: torch.dtype = torch.bfloat16, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, ) -> torch.Tensor: """Decode GGML-compatible IQ1_S payload bytes.""" shape = validate_packed_weights( packed_weights, weight_shape, block_bytes=IQ1_S_BLOCK_BYTES, format_name="IQ1_S" ) + validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES) - d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() - low = blocks[:, 2:34].to(torch.int64).reshape(-1, 8, 4) - qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) - high = (qh.unsqueeze(-1) >> shifts) & 0x7 - entries = low | (high << 8) - - local = (qh >> 12) & 0x7 - delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA) - values = iq1_s_grid(blocks.device)[entries] + delta.unsqueeze(-1).unsqueeze(-1) - scales = d.unsqueeze(-1) * (2 * local + 1).float() - decoded = values * scales.unsqueeze(-1).unsqueeze(-1) - return decoded.reshape(shape).to(dtype) + grid = iq1_s_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ1_S_BLOCK_SIZE), dtype=dtype, device=blocks.device) + for start in range(0, blocks.shape[0], block_chunk_size): + stop = min(start + block_chunk_size, blocks.shape[0]) + block_chunk = blocks[start:stop] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = block_chunk[:, 2:34].to(torch.int64).reshape(-1, 8, 4) + qh = block_chunk[:, 34:50:2].to(torch.int64) | ( + block_chunk[:, 35:50:2].to(torch.int64) << 8 + ) + high = (qh.unsqueeze(-1) >> shifts) & 0x7 + entries = low | (high << 8) + local = (qh >> 12) & 0x7 + delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA) + values = grid[entries] + delta.unsqueeze(-1).unsqueeze(-1) + scales = d.unsqueeze(-1) * (2 * local + 1).float() + chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_S_BLOCK_SIZE) + return decoded.reshape(shape) def iq1_s_fake_quant( diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index e78760ec1fc..6321eb8b0b3 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -64,6 +64,7 @@ from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, + validate_block_chunk_size, validate_packed_weights, validate_weight, ) @@ -255,10 +256,7 @@ def quantize_iq2_xs( treated as zero during packing. """ validate_weight(weight, "IQ2_XS") - if isinstance(block_chunk_size, bool) or not isinstance(block_chunk_size, int): - raise TypeError("block_chunk_size must be an integer") - if block_chunk_size <= 0: - raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + validate_block_chunk_size(block_chunk_size) logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE) @@ -297,34 +295,43 @@ def dequantize_iq2_xs( weight_shape: torch.Tensor, *, dtype: torch.dtype = torch.bfloat16, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, ) -> torch.Tensor: """Decode GGML-compatible IQ2_XS payload bytes.""" shape = validate_packed_weights( packed_weights, weight_shape, block_bytes=IQ2_XS_BLOCK_BYTES, format_name="IQ2_XS" ) + validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES) - d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() - codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) - entries = codes & 0x1FF - sign_index = codes >> 9 - - parity = torch.zeros_like(sign_index) - for bit in range(7): - parity ^= (sign_index >> bit) & 1 - sign_mask = sign_index | (parity << 7) bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) - signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() - - scale_bytes = blocks[:, 66:].to(torch.int64) - local = torch.empty((blocks.shape[0], 16), dtype=torch.int64, device=blocks.device) - local[:, 0::2] = scale_bytes & 0x0F - local[:, 1::2] = scale_bytes >> 4 - # Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8. - scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 - values = iq2_xs_grid(blocks.device)[entries] * signs - decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) - return decoded.reshape(shape).to(dtype) + grid = iq2_xs_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ2_XS_BLOCK_SIZE), dtype=dtype, device=blocks.device) + for start in range(0, blocks.shape[0], block_chunk_size): + stop = min(start + block_chunk_size, blocks.shape[0]) + block_chunk = blocks[start:stop] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + codes = block_chunk[:, 2:66:2].to(torch.int64) | ( + block_chunk[:, 3:66:2].to(torch.int64) << 8 + ) + entries = codes & 0x1FF + sign_index = codes >> 9 + parity = torch.zeros_like(sign_index) + for bit in range(7): + parity ^= (sign_index >> bit) & 1 + sign_mask = sign_index | (parity << 7) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + + scale_bytes = block_chunk[:, 66:].to(torch.int64) + local = torch.empty((block_chunk.shape[0], 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + # Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8. + scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 + values = grid[entries] * signs + chunk_decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XS_BLOCK_SIZE) + return decoded.reshape(shape) def iq2_xs_fake_quant( diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 6aade467791..869e602e4bb 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -30,6 +30,7 @@ def test_ggml_backend_via_quantize(num_bits): torch.manual_seed(1234) model = torch.nn.Linear(256, 2, bias=False) inputs = torch.randn(2, 256) + unquantized_output = model(inputs).detach() config = { "quant_cfg": [ {"quantizer_name": "*", "enable": False}, @@ -49,6 +50,7 @@ def test_ggml_backend_via_quantize(num_bits): assert model.weight_quantizer.num_bits == num_bits assert output.shape == (2, 2) assert torch.isfinite(output).all() + assert not torch.equal(output, unquantized_output) def test_ggml_backend_rejects_unknown_format(): diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py index 2d12934b653..42c32942abc 100644 --- a/tests/unit/torch/quantization/test_iq1_s.py +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -84,11 +84,13 @@ def test_iq1_s_round_trip_and_payload_fields(): weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16) packed, shape = quantize_iq1_s(weight, block_chunk_size=1) - reconstructed = dequantize_iq1_s(packed, shape) + reconstructed = dequantize_iq1_s(packed, shape, block_chunk_size=1) + default_reconstructed = dequantize_iq1_s(packed, shape) assert packed.shape == (2, 1, 50) assert reconstructed.shape == weight.shape assert reconstructed.dtype == torch.bfloat16 + assert torch.equal(reconstructed, default_reconstructed) normalized_mse = ( reconstructed.float() - weight.float() ).square().mean() / weight.float().square().mean() diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py index 5b16542737a..644bda6a906 100644 --- a/tests/unit/torch/quantization/test_iq2_xs.py +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -71,11 +71,13 @@ def test_iq2_xs_round_trip_and_payload_fields(): weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) packed, shape = quantize_iq2_xs(weight, block_chunk_size=2) - reconstructed = dequantize_iq2_xs(packed, shape) + reconstructed = dequantize_iq2_xs(packed, shape, block_chunk_size=1) + default_reconstructed = dequantize_iq2_xs(packed, shape) assert packed.shape == (2, 2, 74) assert reconstructed.shape == weight.shape assert reconstructed.dtype == torch.bfloat16 + assert torch.equal(reconstructed, default_reconstructed) normalized_mse = ( reconstructed.float() - weight.float() ).square().mean() / weight.float().square().mean() From e4f6270d0f19aa8bc9ad7196f54041afaa81ab00 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 17 Sep 2026 21:01:07 +0000 Subject: [PATCH 7/9] Adopt the consolidated GGML CUDA extension #2462 merged the two IQ packing extensions into one. get_cuda_ext_iq1_s and get_cuda_ext_iq2_xs are gone, replaced by get_cuda_ext_ggml, and the packer each exposed as `pack` is now `iq1_s_pack` / `iq2_xs_pack` on the shared module. Update both codecs and their CUDA tests accordingly. The monkeypatched fallback tests are unaffected in substance: each codec still imports the getter into its own module namespace, so patching it out isolates to one format. Verified on an RTX PRO 6000 Blackwell: 36 unit tests and 33 GPU tests, covering both codec parity suites and the extension-boundary suite. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/iq1_s.py | 6 +++--- modelopt/torch/quantization/ggml/iq2_xs.py | 6 +++--- .../gpu/torch/quantization/test_iq1_s_cuda.py | 18 +++++++++--------- .../gpu/torch/quantization/test_iq2_xs_cuda.py | 12 ++++++------ 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 117660f704f..55718087981 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -62,7 +62,7 @@ import torch -from ..extensions import get_cuda_ext_iq1_s +from ..extensions import get_cuda_ext_ggml from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, @@ -255,9 +255,9 @@ def quantize_iq1_s( blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE) grid = iq1_s_grid(weight.device) if weight.is_cuda: - extension = get_cuda_ext_iq1_s() + extension = get_cuda_ext_ggml() if extension is not None: - packed = extension.pack(blocks, grid) + packed = extension.iq1_s_pack(blocks, grid) packed_shape = ( *weight.shape[:-1], weight.shape[-1] // IQ1_S_BLOCK_SIZE, diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 6321eb8b0b3..265746c47be 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -60,7 +60,7 @@ import torch -from ..extensions import get_cuda_ext_iq2_xs +from ..extensions import get_cuda_ext_ggml from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, @@ -262,13 +262,13 @@ def quantize_iq2_xs( blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE) grid = iq2_xs_grid(weight.device) if weight.is_cuda: - extension = get_cuda_ext_iq2_xs() + extension = get_cuda_ext_ggml() if extension is not None: scale_chunks = [ _predict_iq2_xs_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) ] - packed = extension.pack(blocks, grid, torch.cat(scale_chunks)) + packed = extension.iq2_xs_pack(blocks, grid, torch.cat(scale_chunks)) packed_shape = ( *weight.shape[:-1], weight.shape[-1] // IQ2_XS_BLOCK_SIZE, diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py index 2716bd51e27..9788b34ffb1 100644 --- a/tests/gpu/torch/quantization/test_iq1_s_cuda.py +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -16,12 +16,12 @@ import torch import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module -from modelopt.torch.quantization.extensions import get_cuda_ext_iq1_s +from modelopt.torch.quantization.extensions import get_cuda_ext_ggml from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, iq1_s_grid, quantize_iq1_s def _extension(): - extension = get_cuda_ext_iq1_s(raise_if_failed=True) + extension = get_cuda_ext_ggml(raise_if_failed=True) assert extension is not None return extension @@ -30,9 +30,9 @@ def test_iq1_s_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16) - packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) - packed_again = _extension().pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) - monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) + packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) + packed_again = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None) reference, shape = quantize_iq1_s(weight) reconstructed = dequantize_iq1_s(packed, shape) @@ -48,7 +48,7 @@ def test_iq1_s_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) - packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) + packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) shape = torch.tensor(weight.shape, device="cuda") assert not packed.any() @@ -59,15 +59,15 @@ def test_iq1_s_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch): weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16) weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") - packed = _extension().pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) - monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) + packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None) reference, _ = quantize_iq1_s(weight) assert torch.equal(packed, reference) def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch): - monkeypatch.setattr(iq1_s_module, "get_cuda_ext_iq1_s", lambda: None) + monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py index b6450bc1071..aeeb40171d9 100644 --- a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -16,12 +16,12 @@ import torch import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module -from modelopt.torch.quantization.extensions import get_cuda_ext_iq2_xs +from modelopt.torch.quantization.extensions import get_cuda_ext_ggml from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, iq2_xs_grid, quantize_iq2_xs def _extension(): - extension = get_cuda_ext_iq2_xs(raise_if_failed=True) + extension = get_cuda_ext_ggml(raise_if_failed=True) assert extension is not None return extension @@ -29,7 +29,7 @@ def _extension(): def _pack(weight): blocks = weight.contiguous().reshape(-1, 256) scales = iq2_xs_module._predict_iq2_xs_scales(blocks) - return _extension().pack(weight, iq2_xs_grid("cuda"), scales) + return _extension().iq2_xs_pack(weight, iq2_xs_grid("cuda"), scales) def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): @@ -38,7 +38,7 @@ def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): packed = _pack(weight).reshape(8, 2, 74) packed_again = _pack(weight).reshape(8, 2, 74) - monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None) reference, shape = quantize_iq2_xs(weight) reconstructed = dequantize_iq2_xs(packed, shape) @@ -73,14 +73,14 @@ def test_iq2_xs_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch): weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") packed = _pack(weight).reshape(1, 1, 74) - monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None) reference, _ = quantize_iq2_xs(weight) assert torch.equal(packed, reference) def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch): - monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_iq2_xs", lambda: None) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None) generator = torch.Generator(device="cuda").manual_seed(1234) weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) From 16712b69ce70d7578f9243e12cd89710c1877ef7 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 17 Sep 2026 21:18:52 +0000 Subject: [PATCH 8/9] Move the IQ codebooks into their own module The two base64 blobs dominated the codec files: about 110 of the 333 lines in iq1_s.py and 116 of 353 in iq2_xs.py were opaque data a reviewer has to scroll past to reach the encoders. Move both tables, and the base64/zlib decoding that belongs with them, into ggml/codebooks.py. The codecs keep their public iq1_s_grid/iq2_xs_grid accessors, which are real logic -- device resolution, caching, and tensor shaping -- and now call iq1_s_grid_bytes/iq2_xs_grid_bytes for the raw data. This also bounds the third-party surface. The tables are the only GGML material reproduced in this package, so the MIT header travels with them and codebooks.py becomes the single file carrying it: iq1_s.py and iq2_xs.py go back to a plain Apache-2.0 header and come off the insert-license exclude list, which now names codebooks.py alone. Their module docstrings already cited the pinned ggml-common.h revision; those sentences now point at the new module instead of claiming the grid sits "below". Verified the move is lossless: both decoded grids are byte-identical to before, sha256 4ca82266881c8a77 for the [2048, 8] ternary table and 989f82d20f8b93e2 for the [512, 8] magnitude table. 36 unit tests and 33 GPU tests pass on an RTX PRO 6000 Blackwell. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .pre-commit-config.yaml | 3 +- modelopt/torch/quantization/ggml/codebooks.py | 178 ++++++++++++++++++ modelopt/torch/quantization/ggml/iq1_s.py | 96 +--------- modelopt/torch/quantization/ggml/iq2_xs.py | 102 +--------- 4 files changed, 189 insertions(+), 190 deletions(-) create mode 100644 modelopt/torch/quantization/ggml/codebooks.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0202b9b2fec..7654acab550 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,8 +103,7 @@ repos: exclude: > (?x)^( modelopt/torch/quantization/utils/calib_utils.py| - modelopt/torch/quantization/ggml/iq1_s.py| - modelopt/torch/quantization/ggml/iq2_xs.py| + modelopt/torch/quantization/ggml/codebooks.py| modelopt/onnx/quantization/operators.py| modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| diff --git a/modelopt/torch/quantization/ggml/codebooks.py b/modelopt/torch/quantization/ggml/codebooks.py new file mode 100644 index 00000000000..04402f4b1d0 --- /dev/null +++ b/modelopt/torch/quantization/ggml/codebooks.py @@ -0,0 +1,178 @@ +# This file includes the IQ1_S and IQ2_XS codebooks adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canonical IQ codebooks, carried verbatim from the GGML tables. + +This module holds only data and its decoding. The tables are the one piece of GGML material +reproduced in this package, so keeping them here bounds the MIT-licensed surface to a single +file and leaves the codecs beside it as reviewable logic. +""" + +import base64 +import zlib +from functools import cache + +# zlib-compressed little-endian bytes of the canonical uint64_t table. The +# decoded int8 values are -1, 0, and 1. +_IQ1_S_GRID_ZLIB_B64 = ( + "eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv" + "c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY" + "4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX" + "bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8" + "Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W" + "YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH" + "9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL" + "2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E" + "UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg" + "MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0" + "jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+" + "LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs" + "fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2" + "a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g" + "n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq" + "ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ" + "q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/" + "1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2" + "kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM" + "T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8" + "0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT" + "oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy" + "i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44" + "rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg" + "O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO" + "k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI" + "Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8" + "73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn" + "gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6" + "hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+" + "EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA" + "0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH" + "OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD" + "FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn" + "h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6" + "DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB" + "MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI" + "4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh" + "eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ" + "zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep" + "V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5" + "/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C" + "qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T" + "rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD" + "hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr" + "x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO" + "umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0" + "Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA" + "MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ" + "jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8=" +) + +# Compact byte representation of the canonical [512, 8] grid. Values are only +# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids +# adding a pickle-backed torch.save artifact to the wheel. +_IQ2_XS_GRID_B64 = ( + "CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr" + "CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI" + "GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI" + "GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr" + "CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr" + "CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI" + "CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI" + "KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ" + "CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI" + "CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI" + "CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ" + "GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI" + "CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI" + "KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ" + "CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI" + "GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI" + "GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI" + "GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ" + "CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI" + "CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI" + "KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI" + "GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr" + "KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr" + "KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI" + "CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI" + "GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ" + "GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr" + "CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ" + "CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI" + "CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr" + "CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI" + "CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI" + "CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ" + "GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI" + "CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr" + "GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ" + "GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI" + "CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ" + "KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI" + "KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr" + "CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI" + "KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI" + "GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr" + "GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr" + "CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI" + "KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ" + "GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI" + "CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ" + "GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr" + "KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ" + "CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw==" +) + + +@cache +def iq1_s_grid_bytes() -> bytes: + """Decoded little-endian int8 bytes of the [2048, 8] IQ1_S ternary table.""" + return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64)) + + +@cache +def iq2_xs_grid_bytes() -> bytes: + """Decoded bytes of the [512, 8] IQ2_XS magnitude table.""" + return base64.b64decode(_IQ2_XS_GRID_B64) diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 55718087981..cde4857f60c 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -1,30 +1,5 @@ -# This file includes the IQ1_S codebook adapted from: -# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h -# -# MIT License -# -# Copyright (c) 2023-2026 The ggml authors -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 AND MIT +# SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -52,17 +27,15 @@ Each metadata word describes four consecutive eight-value vectors. Bits 0..11 hold the three high index bits, bits 12..14 select one of eight local scales, and bit 15 selects the shared -0.125 rather than +0.125 delta. The canonical -2048 x 8 ternary grid below comes from llama.cpp ``ggml-common.h`` revision +2048 x 8 ternary grid lives in :mod:`.codebooks`, carried from llama.cpp +``ggml-common.h`` revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. """ -import base64 -import zlib -from functools import cache - import torch from ..extensions import get_cuda_ext_ggml +from .codebooks import iq1_s_grid_bytes from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, @@ -90,76 +63,17 @@ # At 1024 blocks, each largest IQ1_S search temporary is about 16 MiB in FP32. _DEFAULT_BLOCK_CHUNK_SIZE = 1024 -# zlib-compressed little-endian bytes of the canonical uint64_t table. The -# decoded int8 values are -1, 0, and 1. -_IQ1_S_GRID_ZLIB_B64 = ( - "eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv" - "c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY" - "4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX" - "bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8" - "Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W" - "YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH" - "9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL" - "2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E" - "UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg" - "MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0" - "jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+" - "LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs" - "fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2" - "a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g" - "n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq" - "ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ" - "q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/" - "1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2" - "kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM" - "T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8" - "0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT" - "oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy" - "i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44" - "rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg" - "O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO" - "k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI" - "Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8" - "73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn" - "gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6" - "hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+" - "EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA" - "0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH" - "OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD" - "FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn" - "h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6" - "DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB" - "MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI" - "4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh" - "eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ" - "zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep" - "V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5" - "/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C" - "qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T" - "rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD" - "hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr" - "x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO" - "umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0" - "Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA" - "MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ" - "jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8=" -) _GRID_CACHE: dict[torch.device, torch.Tensor] = {} -@cache -def _grid_bytes() -> bytes: - return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64)) - - def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: """Return the canonical IQ1_S ternary grid as float32.""" resolved_device = torch.device(device or "cpu") if resolved_device.type == "cuda" and resolved_device.index is None: resolved_device = torch.device("cuda", torch.cuda.current_device()) if resolved_device not in _GRID_CACHE: - raw = torch.tensor(list(_grid_bytes()), dtype=torch.uint8).view(torch.int8) + raw = torch.tensor(list(iq1_s_grid_bytes()), dtype=torch.uint8).view(torch.int8) _GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to( device=resolved_device, dtype=torch.float32 ) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 265746c47be..c85728cf8aa 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -1,30 +1,5 @@ -# This file includes the IQ2_XS codebook adapted from: -# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h -# -# MIT License -# -# Copyright (c) 2023-2026 The ggml authors -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 AND MIT +# SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -49,18 +24,16 @@ * bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign) * bytes 66..73: 16 four-bit local scales, two per byte -The canonical 512 x 8 magnitude grid below comes from llama.cpp -ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +The canonical 512 x 8 magnitude grid lives in :mod:`.codebooks`, carried from +llama.cpp ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. The matching dequantization formula is in ggml-quants.c at the same revision: https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-quants.c#L2516-L2538 """ -import base64 -from functools import cache - import torch from ..extensions import get_cuda_ext_ggml +from .codebooks import iq2_xs_grid_bytes from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, @@ -90,82 +63,17 @@ _DEFAULT_BLOCK_CHUNK_SIZE = 256 _SCALE_BLOCK_CHUNK_SIZE = 4096 -# Compact byte representation of the canonical [512, 8] grid. Values are only -# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids -# adding a pickle-backed torch.save artifact to the wheel. -_IQ2_XS_GRID_B64 = ( - "CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr" - "CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI" - "CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI" - "GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI" - "GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI" - "CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr" - "CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr" - "CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI" - "CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI" - "KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ" - "CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI" - "CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI" - "CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI" - "CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ" - "GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI" - "CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI" - "KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ" - "CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI" - "GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI" - "GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI" - "GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ" - "CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI" - "CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI" - "CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI" - "KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI" - "GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr" - "KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr" - "KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI" - "CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI" - "GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ" - "GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr" - "CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ" - "CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI" - "CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr" - "CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI" - "CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI" - "CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ" - "GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI" - "CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr" - "GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ" - "GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI" - "CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ" - "KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI" - "KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr" - "CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI" - "KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI" - "GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr" - "GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr" - "CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI" - "KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ" - "GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI" - "CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ" - "GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr" - "KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ" - "CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw==" -) _GRID_CACHE: dict[torch.device, torch.Tensor] = {} -@cache -def _grid_bytes() -> bytes: - return base64.b64decode(_IQ2_XS_GRID_B64) - - def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: """Return the canonical IQ2_XS magnitude grid as float32.""" resolved_device = torch.device(device or "cpu") if resolved_device.type == "cuda" and resolved_device.index is None: resolved_device = torch.device("cuda", torch.cuda.current_device()) if resolved_device not in _GRID_CACHE: - values = torch.tensor(list(_grid_bytes()), dtype=torch.float32) + values = torch.tensor(list(iq2_xs_grid_bytes()), dtype=torch.float32) _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) return _GRID_CACHE[resolved_device] From f342df43ad483a0b5184be2a111adc7a846039dd Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 17 Sep 2026 22:03:46 +0000 Subject: [PATCH 9/9] Saturate finite out-of-range float64 weights in the reference encoders The CPU encoders converted with ``blocks.float()`` before ``nan_to_num``, so a finite float64 above the float32 range became infinity and then zero: a large weight silently encoded as nothing. Sanitize at the source precision instead and clamp to the float32 range, via a shared ``narrow_to_float32`` used by both ``_encode_blocks`` and ``_predict_iq2_xs_scales``. This closes a divergence introduced when the CUDA ``load_float`` helper gained the same fix: before that the two paths agreed by both zeroing, and afterwards only the extension saturated. For a [1, 256] float64 weight with one 1e100 element the two encoders disagreed outright -- CPU decoded it as -0.18, CUDA as 73728. The clamp applies only to float64. The float32 bounds do not fit in float16 or bfloat16, so clamping those would raise rather than no-op, and their finite range is already inside float32. Tests, guarding both the policy and the CPU/CUDA agreement: - narrow_to_float32 maps non-finite to zero and saturates out-of-range values - both CPU encoders pack such a weight identically to one clamped by hand, and differently from one zeroed by hand - both CUDA encoders do the same - float64 weights inside the float32 range pack byte-identically on both paths Byte parity is deliberately not asserted for out-of-range values: at those magnitudes the squared-error objective overflows to infinity in float32, every codebook candidate ties, and the two searches break the tie differently. The residual difference is two bytes in the affected block. Saturation is the property both paths can agree on. 39 unit tests and 37 GPU tests pass on an RTX PRO 6000 Blackwell. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/common.py | 18 +++++++++++ modelopt/torch/quantization/ggml/iq1_s.py | 3 +- modelopt/torch/quantization/ggml/iq2_xs.py | 5 +-- .../gpu/torch/quantization/test_iq1_s_cuda.py | 31 +++++++++++++++++++ .../torch/quantization/test_iq2_xs_cuda.py | 31 +++++++++++++++++++ .../torch/quantization/test_ggml_backend.py | 16 ++++++++++ tests/unit/torch/quantization/test_iq1_s.py | 19 ++++++++++++ tests/unit/torch/quantization/test_iq2_xs.py | 19 ++++++++++++ 8 files changed, 139 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 88689e6e087..80e8cf2e2e2 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -95,6 +95,24 @@ def fake_quantize_with_cache( return inputs + (reconstructed - inputs).detach() +def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: + """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does. + + Non-finite elements become zero, and finite elements outside the float32 range saturate + instead of overflowing to infinity and then being zeroed. Sanitizing at the source precision + is what keeps the reference encoders byte-identical to the extension for float64 weights; + converting first would turn a finite 1e100 into zero on this path and into the float32 + maximum on the CUDA one. + """ + finite = torch.nan_to_num(blocks, nan=0.0, posinf=0.0, neginf=0.0) + if finite.dtype == torch.float64: + # Only float64 can hold a finite value the narrowing would overflow. The float32 bounds + # do not fit in the narrower dtypes, so clamping them would raise rather than no-op. + info = torch.finfo(torch.float32) + finite = finite.clamp(info.min, info.max) + return finite.float() + + def validate_weight(weight: torch.Tensor, format_name: str) -> None: """Validate weight metadata accepted by the current GGML block encoders.""" if weight.numel() == 0: diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index cde4857f60c..804d27af058 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -39,6 +39,7 @@ from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, + narrow_to_float32, validate_block_chunk_size, validate_packed_weights, validate_weight, @@ -82,7 +83,7 @@ def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: """Encode a moderate-size batch of flattened 256-value blocks.""" - x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) + x = narrow_to_float32(blocks) block_count = x.shape[0] vectors = x.reshape(block_count, 32, 8) xnorm = vectors.square().sum(dim=-1) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index c85728cf8aa..2f43e71d479 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -37,6 +37,7 @@ from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, + narrow_to_float32, validate_block_chunk_size, validate_packed_weights, validate_weight, @@ -80,7 +81,7 @@ def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: def _predict_iq2_xs_scales(blocks: torch.Tensor) -> torch.Tensor: """Predict one FP16 super-block scale for each flattened block.""" - x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) + x = narrow_to_float32(blocks) amax = x.abs().amax(dim=1) rms = x.square().mean(dim=1).sqrt() peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) @@ -96,7 +97,7 @@ def _encode_blocks( blocks: torch.Tensor, grid: torch.Tensor, scales: torch.Tensor | None = None ) -> torch.Tensor: """Encode a moderate-size batch of flattened 256-value blocks.""" - x = torch.nan_to_num(blocks.float(), nan=0.0, posinf=0.0, neginf=0.0) + x = narrow_to_float32(blocks) block_count = x.shape[0] vectors = x.reshape(block_count, 32, 8) magnitudes = vectors.abs() diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py index 9788b34ffb1..4f05daf8632 100644 --- a/tests/gpu/torch/quantization/test_iq1_s_cuda.py +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -79,3 +79,34 @@ def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch): assert packed.shape == (2, 1, 50) assert normalized_mse < 0.25 + + +def test_iq1_s_cuda_float64_matches_pytorch_encoder(): + """float64 weights inside the float32 range must pack identically on both paths.""" + weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7)) + + reference, _ = quantize_iq1_s(weight) + packed, _ = quantize_iq1_s(weight.cuda()) + + assert torch.equal(reference, packed.cpu()) + + +def test_iq1_s_cuda_saturates_finite_values_above_the_float32_range(): + """The extension saturates such values rather than dropping them to zero. + + Byte parity with the reference encoder is not asserted here: at these magnitudes the + squared-error objective overflows to infinity in float32, so every codebook candidate ties + and the two search implementations break that tie differently. The saturation policy is + what both paths must agree on. + """ + weight = torch.randn(1, 256, dtype=torch.float64, device="cuda") + weight[0, 7] = 1e100 + saturated = weight.clone() + saturated[0, 7] = torch.finfo(torch.float32).max + zeroed = weight.clone() + zeroed[0, 7] = 0.0 + + packed, _ = quantize_iq1_s(weight) + + assert torch.equal(packed, quantize_iq1_s(saturated)[0]) + assert not torch.equal(packed, quantize_iq1_s(zeroed)[0]) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py index aeeb40171d9..6bf5b84e781 100644 --- a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -92,3 +92,34 @@ def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch): assert packed.shape == (2, 1, 74) assert normalized_mse < 0.1 + + +def test_iq2_xs_cuda_float64_matches_pytorch_encoder(): + """float64 weights inside the float32 range must pack identically on both paths.""" + weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7)) + + reference, _ = quantize_iq2_xs(weight) + packed, _ = quantize_iq2_xs(weight.cuda()) + + assert torch.equal(reference, packed.cpu()) + + +def test_iq2_xs_cuda_saturates_finite_values_above_the_float32_range(): + """The extension saturates such values rather than dropping them to zero. + + Byte parity with the reference encoder is not asserted here: at these magnitudes the + squared-error objective overflows to infinity in float32, so every codebook candidate ties + and the two search implementations break that tie differently. The saturation policy is + what both paths must agree on. + """ + weight = torch.randn(1, 256, dtype=torch.float64, device="cuda") + weight[0, 7] = 1e100 + saturated = weight.clone() + saturated[0, 7] = torch.finfo(torch.float32).max + zeroed = weight.clone() + zeroed[0, 7] = 0.0 + + packed, _ = quantize_iq2_xs(weight) + + assert torch.equal(packed, quantize_iq2_xs(saturated)[0]) + assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0]) diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 869e602e4bb..ee78bab5713 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -23,6 +23,7 @@ import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module from modelopt.torch.quantization.ggml.backend import ggml_fake_quant +from modelopt.torch.quantization.ggml.common import narrow_to_float32 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) @@ -119,3 +120,18 @@ def counted_quantize(*args, **kwargs): weight.add_(0.01) fake_quant(weight, quantizer, block_chunk_size=2) assert call_count == 3 + + +def test_narrow_to_float32_matches_the_cuda_load_float_policy(): + """Non-finite elements become zero; finite out-of-range elements saturate.""" + largest = torch.finfo(torch.float32).max + values = torch.tensor( + [torch.nan, torch.inf, -torch.inf, 1e100, -1e100, 1.5], dtype=torch.float64 + ) + + narrowed = narrow_to_float32(values) + + assert narrowed.dtype is torch.float32 + assert torch.equal( + narrowed, torch.tensor([0.0, 0.0, 0.0, largest, -largest, 1.5], dtype=torch.float32) + ) diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py index 42c32942abc..0ce9e8f0d28 100644 --- a/tests/unit/torch/quantization/test_iq1_s.py +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -162,3 +162,22 @@ class Quantizer: output.sum().backward() assert torch.equal(weight.grad, torch.ones_like(weight)) + + +def test_iq1_s_saturates_finite_values_above_the_float32_range(): + """float64 weights are accepted, so a finite value too large for float32 must saturate. + + Converting before sanitizing would turn it into infinity and then zero, which silently + encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy. + """ + weight = torch.randn(1, 256, dtype=torch.float64) + weight[0, 7] = 1e100 + saturated = weight.clone() + saturated[0, 7] = torch.finfo(torch.float32).max + zeroed = weight.clone() + zeroed[0, 7] = 0.0 + + packed, _ = quantize_iq1_s(weight) + + assert torch.equal(packed, quantize_iq1_s(saturated)[0]) + assert not torch.equal(packed, quantize_iq1_s(zeroed)[0]) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py index 644bda6a906..73a1b352c8a 100644 --- a/tests/unit/torch/quantization/test_iq2_xs.py +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -147,3 +147,22 @@ class Quantizer: output.sum().backward() assert torch.equal(weight.grad, torch.ones_like(weight)) + + +def test_iq2_xs_saturates_finite_values_above_the_float32_range(): + """float64 weights are accepted, so a finite value too large for float32 must saturate. + + Converting before sanitizing would turn it into infinity and then zero, which silently + encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy. + """ + weight = torch.randn(1, 256, dtype=torch.float64) + weight[0, 7] = 1e100 + saturated = weight.clone() + saturated[0, 7] = torch.finfo(torch.float32).max + zeroed = weight.clone() + zeroed[0, 7] = 0.0 + + packed, _ = quantize_iq2_xs(weight) + + assert torch.equal(packed, quantize_iq2_xs(saturated)[0]) + assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0])