diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2f4fdd595e3..7654acab550 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,6 +103,7 @@ repos: exclude: > (?x)^( modelopt/torch/quantization/utils/calib_utils.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/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..29f42e8f72b 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,9 @@ from .model_quant import * from .nn.modules.quant_module import QuantModuleRegistry from .utils import update_quant_cfg_with_kv_cache_quant + +# 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/__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..97da3332398 --- /dev/null +++ b/modelopt/torch/quantization/ggml/backend.py @@ -0,0 +1,39 @@ +# 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 weight-only 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) + 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, **extra_args) + if num_bits == "iq2_xs": + return iq2_xs_fake_quant(inputs, quantizer, **extra_args) + 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/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/common.py b/modelopt/torch/quantization/ggml/common.py new file mode 100644 index 00000000000..80e8cf2e2e2 --- /dev/null +++ b/modelopt/torch/quantization/ggml/common.py @@ -0,0 +1,163 @@ +# 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 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, + block_chunk_size=block_chunk_size, + ) + 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: + 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}") + + +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, + *, + block_bytes: int, + format_name: str, +) -> tuple[int, ...]: + """Validate a packed payload and return its logical shape.""" + 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 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: + 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..804d27af058 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -0,0 +1,248 @@ +# 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. + +"""IQ1_S fake quantization and GGML-compatible block packing. + +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 +* 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 lives in :mod:`.codebooks`, carried from llama.cpp +``ggml-common.h`` revision +9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +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, + narrow_to_float32, + validate_block_chunk_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 +_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 + + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +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(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 + ) + 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 = narrow_to_float32(blocks) + 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) + # 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, 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) + + # 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 = _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]``. The packed payload remains on the weight's device; + the logical-shape metadata is kept on CPU. Non-finite input elements are + treated as zero during packing. + """ + validate_weight(weight, "IQ1_S") + 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) + grid = iq1_s_grid(weight.device) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + packed = extension.iq1_s_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, + 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) + shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) + 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( + 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'") + 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 new file mode 100644 index 00000000000..2f43e71d479 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -0,0 +1,262 @@ +# 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. + +"""IQ2_XS fake quantization and GGML-compatible block packing. + +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) +* bytes 66..73: 16 four-bit local scales, two per byte + +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 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, + narrow_to_float32, + validate_block_chunk_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 +_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 16 MiB in FP32. +_DEFAULT_BLOCK_CHUNK_SIZE = 256 +_SCALE_BLOCK_CHUNK_SIZE = 4096 + + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +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(iq2_xs_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) + return _GRID_CACHE[resolved_device] + + +def _predict_iq2_xs_scales(blocks: torch.Tensor) -> torch.Tensor: + """Predict one FP16 super-block scale for each flattened block.""" + 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)) + # 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 + ) + 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 = narrow_to_float32(blocks) + 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) + 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 = _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]``. The packed payload remains on the weight's device; + the logical-shape metadata is kept on CPU. Non-finite input elements are + treated as zero during packing. + """ + validate_weight(weight, "IQ2_XS") + 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) + grid = iq2_xs_grid(weight.device) + if weight.is_cuda: + 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.iq2_xs_pack(blocks, grid, torch.cat(scale_chunks)) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XS_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + 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, + 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, + 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) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + 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( + 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'") + 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 new file mode 100644 index 00000000000..4f05daf8632 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -0,0 +1,112 @@ +# 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.ggml.iq1_s as iq1_s_module +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_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +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().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) + + assert packed.shape == (8, 1, 50) + assert torch.equal(packed, packed_again) + assert torch.equal(packed, reference) + assert shape.device.type == "cpu" + 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().iq1_s_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_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().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_ggml", 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 + + +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 new file mode 100644 index 00000000000..6bf5b84e781 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -0,0 +1,125 @@ +# 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.ggml.iq2_xs as iq2_xs_module +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_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +def _pack(weight): + blocks = weight.contiguous().reshape(-1, 256) + scales = iq2_xs_module._predict_iq2_xs_scales(blocks) + return _extension().iq2_xs_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 = _pack(weight).reshape(8, 2, 74) + packed_again = _pack(weight).reshape(8, 2, 74) + monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", 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, reference) + assert shape.device.type == "cpu" + 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 = _pack(weight).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 = _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_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_ggml", 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 + + +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 new file mode 100644 index 00000000000..ee78bab5713 --- /dev/null +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -0,0 +1,137 @@ +# 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 +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 +from modelopt.torch.quantization.ggml.common import narrow_to_float32 + + +@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) + unquantized_output = model(inputs).detach() + 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() + assert not torch.equal(output, unquantized_output) + + +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 + + +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 new file mode 100644 index 00000000000..0ce9e8f0d28 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -0,0 +1,183 @@ +# 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 + +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, + 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_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) + + 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, 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() + 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_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)) + + +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" + + 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)) + + +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 new file mode 100644 index 00000000000..73a1b352c8a --- /dev/null +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -0,0 +1,168 @@ +# 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 + +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, + 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_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) + + 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, 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() + 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_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)) + + +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" + + 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)) + + +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])