From 489ca8ad503b5b01df64c1943b88a31ff5083225 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Tue, 22 Sep 2026 15:13:27 -0700 Subject: [PATCH 1/4] [OMNIML-5899] Add Q8_0 quantization codec Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/quantization/ggml/__init__.py | 7 +- modelopt/torch/quantization/ggml/backend.py | 14 +- modelopt/torch/quantization/ggml/common.py | 24 ++- modelopt/torch/quantization/ggml/iq1_s.py | 2 +- modelopt/torch/quantization/ggml/iq2_xs.py | 2 +- modelopt/torch/quantization/ggml/iq2_xxs.py | 2 +- modelopt/torch/quantization/ggml/q8_0.py | 145 ++++++++++++++++++ modelopt/torch/quantization/ggml/registry.py | 23 ++- .../gpu/torch/quantization/test_q8_0_cuda.py | 80 ++++++++++ .../torch/quantization/test_ggml_backend.py | 38 +++-- tests/unit/torch/quantization/test_q8_0.py | 96 ++++++++++++ 11 files changed, 392 insertions(+), 41 deletions(-) create mode 100644 modelopt/torch/quantization/ggml/q8_0.py create mode 100644 tests/gpu/torch/quantization/test_q8_0_cuda.py create mode 100644 tests/unit/torch/quantization/test_q8_0.py diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index db8a3473067..b13b21ac5a0 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -23,12 +23,17 @@ from .iq2_xs import __all__ as _iq2_xs_all from .iq2_xxs import * from .iq2_xxs import __all__ as _iq2_xxs_all -from .registry import IQ_FORMAT_REGISTRY, IQFormat +from .q8_0 import * +from .q8_0 import __all__ as _q8_0_all +from .registry import GGML_FORMAT_REGISTRY, IQ_FORMAT_REGISTRY, GGMLFormat, IQFormat __all__ = [ # noqa: PLE0604 *_iq1_s_all, *_iq2_xs_all, *_iq2_xxs_all, + *_q8_0_all, + "GGML_FORMAT_REGISTRY", + "GGMLFormat", "IQ_FORMAT_REGISTRY", "IQFormat", ] diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py index 1f4a83eeff6..e8d9ecfd0cc 100644 --- a/modelopt/torch/quantization/ggml/backend.py +++ b/modelopt/torch/quantization/ggml/backend.py @@ -13,16 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TensorQuantizer backend dispatch for GGML-compatible weight-only IQ formats.""" +"""TensorQuantizer backend dispatch for GGML-compatible weight-only formats.""" import torch from ..nn.modules.tensor_quantizer import register_quant_backend -from .registry import IQ_FORMAT_REGISTRY +from .registry import GGML_FORMAT_REGISTRY def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: - """Dispatch an IQ quantizer to its format-specific implementation.""" + """Dispatch a GGML 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", "decode_chunk_size"} @@ -30,11 +30,11 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}") # num_bits arrives untyped from the quantizer and is a tuple for scalar formats, # so narrow before the lookup rather than relying on the dict to reject it. - iq_format = IQ_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None - if iq_format is None: - supported = ", ".join(repr(name) for name in sorted(IQ_FORMAT_REGISTRY)) + ggml_format = GGML_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None + if ggml_format is None: + supported = ", ".join(repr(name) for name in sorted(GGML_FORMAT_REGISTRY)) raise ValueError(f"The ggml backend requires num_bits in ({supported})") - return iq_format.fake_quant(inputs, quantizer, **extra_args) + return ggml_format.fake_quant(inputs, quantizer, **extra_args) register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 2b39d3e73b7..0110df39717 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -122,11 +122,11 @@ def fake_quantize_with_cache( @dataclass(frozen=True) -class IQFormat: - """Everything backend dispatch and export need to know about one IQ format. +class GGMLFormat: + """Everything backend dispatch and export need to know about one GGML format. Each format module declares one of these beside its encoder and decoder, and - :data:`~modelopt.torch.quantization.ggml.registry.IQ_FORMAT_REGISTRY` lists them. The + :data:`~modelopt.torch.quantization.ggml.registry.GGML_FORMAT_REGISTRY` lists them. The per-format pieces -- codebook, search, payload layout -- stay in the format's module; what lives here is the part every format does the same way. """ @@ -174,6 +174,11 @@ def fake_quant( ) +# Compatibility alias for callers that imported the record type before the registry was +# generalized from the IQ family to every supported GGML block format. +IQFormat = GGMLFormat + + def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does. @@ -192,14 +197,16 @@ def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: return finite.float() -def validate_weight(weight: torch.Tensor, format_name: str) -> None: +def validate_weight( + weight: torch.Tensor, format_name: str, *, block_size: int = GGML_BLOCK_SIZE +) -> 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: + if weight.dim() == 0 or weight.shape[-1] % 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)}" + f"{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}") @@ -218,6 +225,7 @@ def validate_packed_weights( weight_shape: torch.Tensor, *, block_bytes: int, + block_size: int = GGML_BLOCK_SIZE, format_name: str, ) -> tuple[int, ...]: """Validate a packed payload and return its logical shape.""" @@ -234,9 +242,9 @@ def validate_packed_weights( 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: + if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % block_size: raise ValueError(f"invalid {format_name} logical weight shape: {shape}") - expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes + expected_payload_values = math.prod(shape) // 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 index 956c52458ee..86960f5618c 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -247,5 +247,5 @@ def dequantize_iq1_s( # Kept for callers of the per-format entry point. The record captured quantize_iq1_s and # dequantize_iq1_s when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq1_s_fake_quant = IQ1_S_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index a4b7d505bc0..bb10cdb3fad 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -267,5 +267,5 @@ def dequantize_iq2_xs( # Kept for callers of the per-format entry point. The record captured quantize_iq2_xs and # dequantize_iq2_xs when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq2_xs_fake_quant = IQ2_XS_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py index a27c2679632..a019da79cc3 100644 --- a/modelopt/torch/quantization/ggml/iq2_xxs.py +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -277,5 +277,5 @@ def dequantize_iq2_xxs( # Kept for callers of the per-format entry point. The record captured quantize_iq2_xxs and # dequantize_iq2_xxs when it was built, so patching those module functions changes neither backend -# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY. iq2_xxs_fake_quant = IQ2_XXS_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/q8_0.py b/modelopt/torch/quantization/ggml/q8_0.py new file mode 100644 index 00000000000..3f5d46e753a --- /dev/null +++ b/modelopt/torch/quantization/ggml/q8_0.py @@ -0,0 +1,145 @@ +# 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. + +"""Q8_0 fake quantization and GGML-compatible block packing. + +Every 32 logical values become one 34-byte ``block_q8_0`` payload: a +little-endian FP16 block scale followed by 32 signed int8 quants. +""" + +import torch + +from ..extensions import get_cuda_ext_ggml +from .common import ( + GGMLFormat, + narrow_to_float32, + validate_block_chunk_size, + validate_packed_weights, + validate_weight, +) + +__all__ = [ + "Q8_0_BLOCK_BYTES", + "Q8_0_BLOCK_SIZE", + "Q8_0_EFFECTIVE_BITS", + "Q8_0_FORMAT", + "dequantize_q8_0", + "q8_0_fake_quant", + "quantize_q8_0", +] + +Q8_0_BLOCK_SIZE = 32 +Q8_0_BLOCK_BYTES = 34 +Q8_0_EFFECTIVE_BITS = Q8_0_BLOCK_BYTES * 8 / Q8_0_BLOCK_SIZE +_Q8_0_MAX_QUANT = 127 +_DEFAULT_BLOCK_CHUNK_SIZE = 4096 +_DEFAULT_DECODE_CHUNK_SIZE = 16384 + + +def _encode_blocks(blocks: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 32-value blocks.""" + x = narrow_to_float32(blocks) + block_count = x.shape[0] + amax = x.abs().amax(dim=1) + # Keep the serialized fp16 scale finite. This extends the existing GGML backend policy for + # finite float64 values outside the float32/fp16 scale range. + d_float = (amax / _Q8_0_MAX_QUANT).clamp(max=65504.0) + d = d_float.to(torch.float16) + inverse = torch.where(d_float > 0, d_float.reciprocal(), torch.zeros_like(d_float)) + normalized = x * inverse.unsqueeze(1) + # C roundf, used by the canonical Q8_0 encoder, rounds half-way cases away from zero. + rounded = normalized.sign() * (normalized.abs() + 0.5).floor() + quants = rounded.clamp(-_Q8_0_MAX_QUANT, _Q8_0_MAX_QUANT).to(torch.int8) + + packed = torch.empty((block_count, Q8_0_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:] = quants.view(torch.uint8) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_q8_0( + weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible Q8_0 blocks.""" + validate_weight(weight, "Q8_0", block_size=Q8_0_BLOCK_SIZE) + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, Q8_0_BLOCK_SIZE) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + packed = extension.q8_0_pack(blocks) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // Q8_0_BLOCK_SIZE, + Q8_0_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size]) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // Q8_0_BLOCK_SIZE, + Q8_0_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_q8_0( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, + block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """Decode GGML-compatible Q8_0 payload bytes.""" + shape = validate_packed_weights( + packed_weights, + weight_shape, + block_bytes=Q8_0_BLOCK_BYTES, + block_size=Q8_0_BLOCK_SIZE, + format_name="Q8_0", + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, Q8_0_BLOCK_BYTES) + decoded = torch.empty((blocks.shape[0], Q8_0_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() + quants = block_chunk[:, 2:].contiguous().view(torch.int8).float() + decoded[start:stop] = (d.unsqueeze(1) * quants).to(dtype) + return decoded.reshape(shape) + + +Q8_0_FORMAT = GGMLFormat( + name="q8_0", + block_size=Q8_0_BLOCK_SIZE, + block_bytes=Q8_0_BLOCK_BYTES, + quantize=quantize_q8_0, + dequantize=dequantize_q8_0, + block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, +) + +# Keep the public per-format entry point on the same record used by backend dispatch. +q8_0_fake_quant = Q8_0_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/registry.py b/modelopt/torch/quantization/ggml/registry.py index 5faca3c0377..f31fd90233d 100644 --- a/modelopt/torch/quantization/ggml/registry.py +++ b/modelopt/torch/quantization/ggml/registry.py @@ -13,21 +13,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The GGML IQ formats, listed once for backend dispatch and export.""" +"""GGML block formats, listed once for backend dispatch and export.""" -from .common import IQFormat +from .common import GGMLFormat, IQFormat from .iq1_s import IQ1_S_FORMAT from .iq2_xs import IQ2_XS_FORMAT from .iq2_xxs import IQ2_XXS_FORMAT +from .q8_0 import Q8_0_FORMAT -__all__ = ["IQ_FORMAT_REGISTRY", "IQFormat"] +__all__ = ["GGML_FORMAT_REGISTRY", "IQ_FORMAT_REGISTRY", "GGMLFormat", "IQFormat"] -# Every IQ format, keyed by the name a quantizer's num_bits carries, in increasing bits per -# weight. Backend dispatch and both exporters read this mapping and export's IQ_FORMATS is derived -# from it, so adding a format is one entry here rather than a row in several parallel tables. +# Every GGML block format, keyed by the name a quantizer's num_bits carries, in increasing bits +# per weight. Backend dispatch reads this mapping, and export will derive its supported formats +# from it rather than maintaining parallel metadata and packer tables. # # It is an explicit list rather than formats registering themselves on import, so its contents # never depend on which modules happen to have been imported first. -IQ_FORMAT_REGISTRY: dict[str, IQFormat] = { - fmt.name: fmt for fmt in (IQ1_S_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT) +GGML_FORMAT_REGISTRY: dict[str, GGMLFormat] = { + fmt.name: fmt for fmt in (IQ1_S_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, Q8_0_FORMAT) +} + +# Compatibility view for the already-merged IQ export path. Export migrates to the general +# registry in the Q8_0 export PR; IQ-specific conformance tests continue to use this subset. +IQ_FORMAT_REGISTRY: dict[str, GGMLFormat] = { + name: fmt for name, fmt in GGML_FORMAT_REGISTRY.items() if name.startswith("iq") } diff --git a/tests/gpu/torch/quantization/test_q8_0_cuda.py b/tests/gpu/torch/quantization/test_q8_0_cuda.py new file mode 100644 index 00000000000..3bd03bd657b --- /dev/null +++ b/tests/gpu/torch/quantization/test_q8_0_cuda.py @@ -0,0 +1,80 @@ +# 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.q8_0 as q8_0_module +from modelopt.torch.quantization.extensions import get_cuda_ext_ggml +from modelopt.torch.quantization.ggml.q8_0 import dequantize_q8_0, quantize_q8_0 + + +def _extension(): + extension = get_cuda_ext_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +def test_q8_0_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 64), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _extension().q8_0_pack(weight).reshape(8, 2, 34) + packed_again = _extension().q8_0_pack(weight).reshape(8, 2, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, shape = quantize_q8_0(weight) + reconstructed = dequantize_q8_0(packed, shape) + + assert packed.shape == (8, 2, 34) + 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 < 1e-4 + + +def test_q8_0_cuda_zero_and_nonfinite_encoding_match_reference(monkeypatch): + weight = torch.zeros((2, 32), device="cuda", dtype=torch.bfloat16) + weight[1, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda") + + packed = _extension().q8_0_pack(weight).reshape(2, 1, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, shape = quantize_q8_0(weight) + + assert not packed.any() + assert torch.equal(packed, reference) + assert torch.equal(dequantize_q8_0(packed, shape), torch.zeros_like(weight)) + + +def test_q8_0_cuda_falls_back_to_pytorch_encoder(monkeypatch): + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((2, 32), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight) + reconstructed = dequantize_q8_0(packed, shape) + + assert packed.shape == (2, 1, 34) + assert torch.isfinite(reconstructed).all() + + +def test_q8_0_cuda_float64_matches_pytorch_encoder(): + weight = torch.randn(4, 32, dtype=torch.float64, generator=torch.Generator().manual_seed(7)) + + reference, _ = quantize_q8_0(weight) + packed, _ = quantize_q8_0(weight.cuda()) + + assert torch.equal(reference, packed.cpu()) diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index cffea63cb74..a2506d562d6 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -25,15 +25,16 @@ import modelopt.torch.quantization.ggml as ggml import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +import modelopt.torch.quantization.ggml.q8_0 as q8_0_module from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.backend import ggml_fake_quant from modelopt.torch.quantization.ggml.common import narrow_to_float32 from modelopt.torch.quantization.nn import TensorQuantizer # Every registered format, paired with its own module. The module -- not the registry -- supplies # what a test expects, so a mis-wired registry entry cannot make both sides of an assertion agree. -FORMAT_NAMES = sorted(IQ_FORMAT_REGISTRY) +FORMAT_NAMES = sorted(GGML_FORMAT_REGISTRY) FORMAT_MODULES = [ (name, importlib.import_module(f"modelopt.torch.quantization.ggml.{name}")) for name in FORMAT_NAMES @@ -43,11 +44,11 @@ def _substitute(monkeypatch, name, **changes): """Replace parts of one registered format for the duration of a test. - Dispatch reads IQ_FORMAT_REGISTRY, so that is where a substitute has to go: patching the + Dispatch reads GGML_FORMAT_REGISTRY, so that is where a substitute has to go: patching the format module's function would not reach it. """ monkeypatch.setitem( - IQ_FORMAT_REGISTRY, name, dataclasses.replace(IQ_FORMAT_REGISTRY[name], **changes) + GGML_FORMAT_REGISTRY, name, dataclasses.replace(GGML_FORMAT_REGISTRY[name], **changes) ) @@ -89,6 +90,7 @@ def test_ggml_backend_rejects_unknown_format(): 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 + assert mtq.quantize_q8_0 is q8_0_module.quantize_q8_0 @pytest.mark.parametrize( @@ -108,7 +110,7 @@ def fake_quant(inputs, _quantizer, **kwargs): return inputs # The dispatcher resolves through the registry, so that is the seam to patch. - monkeypatch.setitem(IQ_FORMAT_REGISTRY, "iq1_s", SimpleNamespace(fake_quant=fake_quant)) + monkeypatch.setitem(GGML_FORMAT_REGISTRY, "iq1_s", SimpleNamespace(fake_quant=fake_quant)) inputs = torch.ones(1, 256) quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args=extra_args) @@ -125,7 +127,8 @@ def test_ggml_backend_rejects_unknown_extra_arg(): @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_ggml_backend_caches_packed_weight_and_invalidates_on_change(monkeypatch, num_bits, module): - weight = torch.randn(1, 256) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size + weight = torch.randn(1, block_size) quantizer = SimpleNamespace(num_bits=num_bits, _quantizer_cache=None) original_quantize = getattr(module, f"quantize_{num_bits}") call_count = 0 @@ -136,7 +139,7 @@ def counted_quantize(*args, **kwargs): return original_quantize(*args, **kwargs) _substitute(monkeypatch, num_bits, quantize=counted_quantize) - fake_quant = IQ_FORMAT_REGISTRY[num_bits].fake_quant + fake_quant = GGML_FORMAT_REGISTRY[num_bits].fake_quant fake_quant(weight, quantizer, block_chunk_size=1) fake_quant(weight, quantizer, block_chunk_size=1) @@ -183,8 +186,9 @@ def counting(weight, **kwargs): return original(weight, **kwargs) _substitute(monkeypatch, num_bits, quantize=counting) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: block_size}, backend="ggml") ) weight = torch.randn(4, 256) @@ -192,7 +196,9 @@ def counting(weight, **kwargs): for _ in range(5): quantizer(weight) - assert calls == [(4, 256)], f"expected one pack, got {len(calls)}" + assert calls == [(weight.numel() // block_size, block_size)], ( + f"expected one pack, got {len(calls)}" + ) @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) @@ -212,8 +218,9 @@ def recording(packed_weights, weight_shape, **kwargs): return original(packed_weights, weight_shape, **kwargs) _substitute(monkeypatch, num_bits, dequantize=recording) + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: block_size}, backend="ggml") ) quantizer(torch.randn(4, 256)) @@ -226,13 +233,13 @@ def test_registry_lists_every_exported_encoder(): encoders = { name.removeprefix("quantize_") for name in ggml.__all__ if name.startswith("quantize_") } - assert encoders == set(IQ_FORMAT_REGISTRY) + assert encoders == set(GGML_FORMAT_REGISTRY) @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_registry_record_is_wired_to_its_own_codec(num_bits, module): """Each record points at its own format's encoder, decoder, geometry and chunk defaults.""" - record = IQ_FORMAT_REGISTRY[num_bits] + record = GGML_FORMAT_REGISTRY[num_bits] upper = num_bits.upper() assert record.name == num_bits @@ -248,7 +255,7 @@ def test_registry_record_is_wired_to_its_own_codec(num_bits, module): @pytest.mark.parametrize(("num_bits", "module"), FORMAT_MODULES) def test_public_fake_quant_is_the_registered_record(num_bits, module): """The per-format entry point and backend dispatch run the same code path.""" - assert getattr(module, f"{num_bits}_fake_quant").__self__ is IQ_FORMAT_REGISTRY[num_bits] + assert getattr(module, f"{num_bits}_fake_quant").__self__ is GGML_FORMAT_REGISTRY[num_bits] @pytest.mark.parametrize("num_bits", FORMAT_NAMES) @@ -260,6 +267,9 @@ def test_format_fake_quant_rejects_another_formats_quantizer(num_bits): """ other = next(name for name in FORMAT_NAMES if name != num_bits) expected = f"The ggml {num_bits.upper()} backend requires num_bits={num_bits!r}" + block_size = GGML_FORMAT_REGISTRY[num_bits].block_size with pytest.raises(ValueError, match=re.escape(expected)): - IQ_FORMAT_REGISTRY[num_bits].fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits=other)) + GGML_FORMAT_REGISTRY[num_bits].fake_quant( + torch.ones(1, block_size), SimpleNamespace(num_bits=other) + ) diff --git a/tests/unit/torch/quantization/test_q8_0.py b/tests/unit/torch/quantization/test_q8_0.py new file mode 100644 index 00000000000..5b76cadd991 --- /dev/null +++ b/tests/unit/torch/quantization/test_q8_0.py @@ -0,0 +1,96 @@ +# 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.q8_0 import ( + Q8_0_BLOCK_BYTES, + dequantize_q8_0, + q8_0_fake_quant, + quantize_q8_0, +) + + +def test_q8_0_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 32), dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight) + + assert packed.shape == (2, 1, Q8_0_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert torch.equal(dequantize_q8_0(packed, shape), weight) + + +def test_q8_0_uses_roundf_ties_away_from_zero(): + weight = torch.zeros((1, 32), dtype=torch.float32) + weight[0, :5] = torch.tensor([127.0, 0.5, -0.5, 1.5, -1.5]) + + packed, shape = quantize_q8_0(weight) + + assert packed[0, 0, :2].contiguous().view(torch.float16).item() == 1.0 + assert packed[0, 0, 2:7].contiguous().view(torch.int8).tolist() == [127, 1, -1, 2, -2] + expected = weight.clone() + expected[0, 1:5] = torch.tensor([1.0, -1.0, 2.0, -2.0]) + assert torch.equal(dequantize_q8_0(packed, shape, dtype=torch.float32), expected) + + +def test_q8_0_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 64), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_q8_0(weight, block_chunk_size=1) + reconstructed = dequantize_q8_0(packed, shape, dtype=torch.float32, block_chunk_size=1) + + assert packed.shape == (2, 2, 34) + normalized_mse = ( + reconstructed - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 1e-4 + assert torch.all(packed.reshape(-1, 34)[:, :2].contiguous().view(torch.float16) > 0) + + +def test_q8_0_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="divisible by 32"): + quantize_q8_0(torch.ones(2, 33)) + + +def test_q8_0_treats_nonfinite_values_as_zero(): + weight = torch.randn(1, 32) + weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf]) + + packed, _ = quantize_q8_0(weight) + expected, _ = quantize_q8_0(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0)) + + assert torch.equal(packed, expected) + + +@pytest.mark.parametrize("shape", [[1, 31], [0, 32]]) +def test_q8_0_rejects_invalid_shape_metadata(shape): + packed = torch.zeros((1, 1, 34), dtype=torch.uint8) + + with pytest.raises(ValueError, match="logical weight shape"): + dequantize_q8_0(packed, torch.tensor(shape)) + + +def test_q8_0_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "q8_0" + + weight = torch.randn(1, 32, requires_grad=True) + q8_0_fake_quant(weight, Quantizer()).sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) From 0b5461f5ad80f76a2379cdf751a4efb0b75899da Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Mon, 28 Sep 2026 09:30:43 -0700 Subject: [PATCH 2/4] [OMNIML-5899] Address Q8_0 review feedback Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/quantization/ggml/q8_0.py | 7 ++++++- tests/unit/torch/quantization/test_q8_0.py | 13 ++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/ggml/q8_0.py b/modelopt/torch/quantization/ggml/q8_0.py index 3f5d46e753a..0e16e819113 100644 --- a/modelopt/torch/quantization/ggml/q8_0.py +++ b/modelopt/torch/quantization/ggml/q8_0.py @@ -60,7 +60,12 @@ def _encode_blocks(blocks: torch.Tensor) -> torch.Tensor: inverse = torch.where(d_float > 0, d_float.reciprocal(), torch.zeros_like(d_float)) normalized = x * inverse.unsqueeze(1) # C roundf, used by the canonical Q8_0 encoder, rounds half-way cases away from zero. - rounded = normalized.sign() * (normalized.abs() + 0.5).floor() + magnitude = normalized.abs() + rounded = normalized.sign() * torch.where( + magnitude.frac() == 0.5, + magnitude.floor() + 1, + magnitude.round(), + ) quants = rounded.clamp(-_Q8_0_MAX_QUANT, _Q8_0_MAX_QUANT).to(torch.int8) packed = torch.empty((block_count, Q8_0_BLOCK_BYTES), dtype=torch.uint8, device=x.device) diff --git a/tests/unit/torch/quantization/test_q8_0.py b/tests/unit/torch/quantization/test_q8_0.py index 5b76cadd991..020ce0ae03d 100644 --- a/tests/unit/torch/quantization/test_q8_0.py +++ b/tests/unit/torch/quantization/test_q8_0.py @@ -37,14 +37,25 @@ def test_q8_0_zero_block_has_canonical_zero_encoding(): def test_q8_0_uses_roundf_ties_away_from_zero(): weight = torch.zeros((1, 32), dtype=torch.float32) + below_half = torch.nextafter(torch.tensor(0.5), torch.tensor(0.0)) weight[0, :5] = torch.tensor([127.0, 0.5, -0.5, 1.5, -1.5]) + weight[0, 5:7] = torch.stack((below_half, -below_half)) packed, shape = quantize_q8_0(weight) assert packed[0, 0, :2].contiguous().view(torch.float16).item() == 1.0 - assert packed[0, 0, 2:7].contiguous().view(torch.int8).tolist() == [127, 1, -1, 2, -2] + assert packed[0, 0, 2:9].contiguous().view(torch.int8).tolist() == [ + 127, + 1, + -1, + 2, + -2, + 0, + 0, + ] expected = weight.clone() expected[0, 1:5] = torch.tensor([1.0, -1.0, 2.0, -2.0]) + expected[0, 5:7] = 0 assert torch.equal(dequantize_q8_0(packed, shape, dtype=torch.float32), expected) From 6d43a4c062460a754030ea159e3b2dc6e2dbd335 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Mon, 28 Sep 2026 15:30:38 -0700 Subject: [PATCH 3/4] fix: align Q8_0 CUDA rounding with reference Signed-off-by: Hung-Yueh Chiang --- modelopt/torch/kernels/quantization/ggml/q8_0.cu | 3 ++- tests/gpu/torch/quantization/test_q8_0_cuda.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/kernels/quantization/ggml/q8_0.cu b/modelopt/torch/kernels/quantization/ggml/q8_0.cu index 5f9707bfcc7..a5bd9be02b1 100644 --- a/modelopt/torch/kernels/quantization/ggml/q8_0.cu +++ b/modelopt/torch/kernels/quantization/ggml/q8_0.cu @@ -53,7 +53,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, uint8_t *outpu payload[1] = static_cast(d_bits >> 8); } d = __shfl_sync(0xffffffff, d, 0); - const float rounded = d > 0.0f ? roundf(value / d) : 0.0f; + const float inverse = d > 0.0f ? 1.0f / d : 0.0f; + const float rounded = roundf(value * inverse); const int quant = static_cast(fminf(fmaxf(rounded, -kMaxQuant), kMaxQuant)); payload[kScaleBytes + lane] = static_cast(static_cast(quant)); } diff --git a/tests/gpu/torch/quantization/test_q8_0_cuda.py b/tests/gpu/torch/quantization/test_q8_0_cuda.py index 3bd03bd657b..eef6fe28703 100644 --- a/tests/gpu/torch/quantization/test_q8_0_cuda.py +++ b/tests/gpu/torch/quantization/test_q8_0_cuda.py @@ -78,3 +78,15 @@ def test_q8_0_cuda_float64_matches_pytorch_encoder(): packed, _ = quantize_q8_0(weight.cuda()) assert torch.equal(reference, packed.cpu()) + + +def test_q8_0_cuda_reciprocal_rounding_boundary_matches_pytorch_encoder(monkeypatch): + weight = torch.zeros((1, 32), device="cuda", dtype=torch.float32) + weight[0, :3] = torch.tensor([5207.0, 20.5, -20.5], device="cuda") + + packed = _extension().q8_0_pack(weight).reshape(1, 1, 34) + monkeypatch.setattr(q8_0_module, "get_cuda_ext_ggml", lambda: None) + reference, _ = quantize_q8_0(weight) + + assert torch.equal(packed, reference) + assert packed[0, 0, 2:5].view(torch.int8).tolist() == [127, 0, 0] From 49f8d1eda88a178c5957532dd73a5f23779b6245 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 29 Sep 2026 22:46:09 +0000 Subject: [PATCH 4/4] Add the IQ1_M CUDA encoder and register the format Second of two changes adding IQ1_M. The previous change landed the PyTorch codec; this one adds its CUDA encoder and makes the format reachable. In the kernel the delta shift is free per group, so it sits above the entry index in the sort key: a tie still prefers the lower shift and then the lower entry, as the reference encoder does. The 2048-entry grid IQ1_M shares with IQ1_S is 64 KiB, past the 48 KiB static shared-memory limit, so both kernels read it from global memory and rely on the cache. On a 5632x2048 weight the encoder runs at 318 M elem/s against the torch search's 5.6, and its packed bytes match the PyTorch encoder's. The two IQ1 kernels load each vector, score it against a grid entry and apply the +/-1/8 shift the same way, so those three steps move into common.cuh as load_vector, grid_terms and shifted_error, and IQ1_S uses them too. IQ1_S's packed bytes are unchanged. IQ1_M gets an IQFormat record and one IQ_FORMAT_REGISTRY entry, so backend dispatch, both exporters and convert_hf_config take it from there. The ggml package exports it, and a general/ptq/iq1_m recipe uses it. Registering it brings it under every registry-driven test with no IQ1_M-specific test code, and the hf_ptq example test gains an IQ1_M case. IQ1_M covers 25 tensors and 1.2B parameters of the mixed-precision checkpoint IQ2_XXS measured. With it, ModelOpt supports all five GGML IQ formats at one and two bits. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + modelopt/torch/export/quant_format.py | 1 + .../kernels/quantization/ggml/common.cuh | 40 ++++ .../torch/kernels/quantization/ggml/ggml.cpp | 24 ++ .../torch/kernels/quantization/ggml/iq1_m.cu | 205 ++++++++++++++++++ .../torch/kernels/quantization/ggml/iq1_s.cu | 62 ++---- modelopt/torch/quantization/extensions.py | 4 +- modelopt/torch/quantization/ggml/__init__.py | 3 + modelopt/torch/quantization/ggml/iq1_m.py | 30 +++ modelopt/torch/quantization/ggml/registry.py | 4 +- modelopt_recipes/configs/numerics/iq1_m.yaml | 26 +++ .../configs/ptq/presets/model/iq1_m.yaml | 32 +++ modelopt_recipes/general/ptq/iq1_m.yaml | 27 +++ modelopt_recipes/ptq.md | 7 +- tests/examples/hf_ptq/test_llm_ptq.py | 5 +- .../quantization/test_iq_formats_cuda.py | 3 + tests/unit/recipe/test_presets.py | 3 + 17 files changed, 421 insertions(+), 56 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_m.cu create mode 100644 modelopt_recipes/configs/numerics/iq1_m.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml create mode 100644 modelopt_recipes/general/ptq/iq1_m.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 86502659238..ea80a8be01a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -20,6 +20,7 @@ Changelog - Add IQ1_S and IQ2_XS weight-only quantization with GGML-compatible 256-value block encoders, built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes, and unified HF and Megatron export of the packed blocks. Quantized weights must have a final dimension divisible by 256, and Megatron export requires tensor and pipeline parallel sizes of 1. - Add ``iq2_xxs`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 2.0625 bits per weight between ``iq1_s`` and ``iq2_xs``. The same 256-value block constraint applies. - Add ``iq2_s`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 2.5625 bits per weight above ``iq2_xs``. The same 256-value block constraint applies. +- Add ``iq1_m`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 1.75 bits per weight between ``iq1_s`` and ``iq2_xxs``. The same 256-value block constraint applies. - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it. - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index 665ef272add..55e7d25f398 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -39,6 +39,7 @@ QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" QUANTIZATION_IQ1_S = "iq1_s" +QUANTIZATION_IQ1_M = "iq1_m" QUANTIZATION_IQ2_XXS = "iq2_xxs" QUANTIZATION_IQ2_XS = "iq2_xs" QUANTIZATION_IQ2_S = "iq2_s" diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index f80a212433e..5eccec6cd03 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -109,6 +109,46 @@ __device__ __forceinline__ float clamped_quant_error(float xnorm, float dot, flo return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f); } +// The IQ1 kernels approximate each 8-value vector by scale * (q + delta), with q from the ternary +// grid and delta = +/- 1/8. The three helpers below are theirs. + +// Loads one 8-value vector with its squared norm and its sum. +template +__device__ __forceinline__ void load_vector(const scalar_t *source, float (&x)[kVectorSize], + float &xnorm, float &xsum) { + xnorm = 0.0f; + xsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + j); + xnorm = fmaf(x[j], x[j], xnorm); + xsum += x[j]; + } +} + +// Accumulates x . q, |q|^2 and sum(q) for one codebook vector. +__device__ __forceinline__ void grid_terms(const float *x, const float *q, float &dot, float &qnorm, + float &qsum) { + dot = 0.0f; + qnorm = 0.0f; + qsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + dot = fmaf(x[j], q[j], dot); + qnorm = fmaf(q[j], q[j], qnorm); + qsum += q[j]; + } +} + +// Squared error of approximating x by scale * (q + delta): the offset shifts the dot and norm +// that grid_terms computed for q alone. +__device__ __forceinline__ float shifted_error(float xnorm, float xsum, float dot, float qnorm, + float qsum, float scale, float delta) { + const float shifted_dot = dot + delta * xsum; + const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; + return clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale); +} + // Orders candidates by error first and codebook index second, so the lowest index wins a tie -- // the rule the PyTorch reference encoder applies. __device__ __forceinline__ unsigned long long error_key(float error, int entry) { diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index 17075d8ba7d..e3d16b2c462 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -25,6 +25,7 @@ at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid); at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); +at::Tensor iq1_m_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); at::Tensor q8_0_pack_cuda(at::Tensor input); namespace { @@ -88,6 +89,21 @@ at::Tensor iq2_s_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { return iq2_s_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); } +at::Tensor iq1_m_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_cuda(), "IQ1_M packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ1_M packing requires a CUDA grid"); + TORCH_CHECK(scales.is_cuda(), "IQ1_M packing requires CUDA scales"); + modelopt::ggml::check_pack_inputs("IQ1_M", input, grid, modelopt::ggml::kIq1sEntries); + const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == num_blocks, + "scales must be float16 [numel / 256]"); + TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), + "scales must be finite and non-negative"); + TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); + return iq1_m_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); +} + at::Tensor q8_0_pack(at::Tensor input) { TORCH_CHECK(input.is_cuda(), "Q8_0 packing requires a CUDA input"); modelopt::ggml::check_scalar_pack_input("Q8_0", input, 32); @@ -126,6 +142,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 82] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); + module.def("iq1_m_pack", &iq1_m_pack, + "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " + "dimension is a multiple of 256. The grid must be float32 [2048, 8], and scales must " + "be finite non-negative float16 [numel / 256]; IQ1_M carries that scale in the four " + "packed scale words rather than a leading field. " + "Returns uint8 [numel / 256, 56] on the input device. Non-finite input elements are " + "treated as zero during packing, and finite elements outside the float32 range " + "saturate."); module.def("q8_0_pack", &q8_0_pack, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 32. Returns uint8 [numel / 32, 34] on the input device. " diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_m.cu b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu new file mode 100644 index 00000000000..b5eadc85677 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu @@ -0,0 +1,205 @@ +/* + * 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. + */ + +#include "common.cuh" + +namespace { + +using namespace modelopt::ggml; + +// The IQ1_M packed payload layout and format constants below follow the GGML +// definition at: +// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +// IQ1_M shares the IQ1_S ternary grid. At 64 KiB it is past the 48 KiB static shared-memory +// limit, so both kernels read it from global memory and rely on the cache. +constexpr int kEntries = kIq1sEntries; +constexpr int kGroups = 16; // one 3-bit local scale per 16 values +constexpr int kVectorsPerGroup = 2; +constexpr int kVectors = kGroups * kVectorsPerGroup; +constexpr int kLocalScales = 8; +constexpr int kSubBlocks = 8; +constexpr int kScaleWords = 4; +constexpr int kLowOffset = 0; // no leading block scale in this format +constexpr int kHighOffset = kLowOffset + kVectors; // 16 bytes, two vectors per byte +constexpr int kScaleWordOffset = kHighOffset + 2 * kSubBlocks; +constexpr int kPayloadBytes = kScaleWordOffset + 2 * kScaleWords; +constexpr float kDelta = 0.125f; +constexpr int kEntryBits = 11; // 2048 entries; the shift bit sits above them in the sort key + +static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); +static_assert(kPayloadBytes == 56, "IQ1_M blocks are 56 bytes"); + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + __shared__ float warp_best[kWarps * kLocalScales]; + __shared__ float group_error[kLocalScales]; + __shared__ unsigned long long warp_keys[kWarps]; + __shared__ int selected_local; + __shared__ uint8_t locals[kGroups]; + + const int tid = threadIdx.x; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * kPayloadBytes; + const __half d_half = scales[block]; + const uint16_t d_bits = __half_as_ushort(d_half); + const float d = __half2float(d_half); + // IQ1_M has no leading scale field, so the shared helper does not apply: a zero block is all + // zero bytes, which decodes to a zero scale and therefore zero values. + if ((d_bits & 0x7FFF) == 0) { + if (tid < kPayloadBytes) + payload[tid] = 0; + return; + } + if (tid < kHighOffset + 2 * kSubBlocks) + payload[tid] = 0; // index-high and shift nibbles are OR-ed into below + __syncthreads(); + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kLocalScales) + group_error[tid] = 0.0f; + __syncthreads(); + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm, xsum; + load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, + xnorm, xsum); + float local_best[kLocalScales]; +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) + local_best[local] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); + // The shift is free per vector here, unlike IQ1_S where it is shared, so take the + // better of the two before the local scale is chosen. +#pragma unroll + for (int shift = 0; shift < 2; ++shift) { + const float delta = shift ? -kDelta : kDelta; +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + const float scale = d * (2 * local + 1); + local_best[local] = fminf(local_best[local], + shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, delta)); + } + } + } + block_min_accumulate(local_best, warp_best, group_error); + } + + if (tid == 0) { + selected_local = 0; + float best = group_error[0]; +#pragma unroll + for (int local = 1; local < kLocalScales; ++local) { + if (group_error[local] < best) { + best = group_error[local]; + selected_local = local; + } + } + locals[group] = static_cast(selected_local); + } + __syncthreads(); + const float selected_scale = d * (2 * selected_local + 1); + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm, xsum; + load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, + xnorm, xsum); + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); +#pragma unroll + for (int shift = 0; shift < 2; ++shift) { + const float delta = shift ? -kDelta : kDelta; + const float error = shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, delta); + // Shift occupies the bits above the entry index so that a tie prefers the lower + // shift first and then the lower entry, matching the reference encoder. + const unsigned long long candidate = error_key(error, (shift << kEntryBits) | entry); + key = candidate < key ? candidate : key; + } + } + key = block_min_key(key, warp_keys); + if (tid == 0) { + const int packed_index = static_cast(key & 0xFFFFFFFFULL); + const int entry = packed_index & ((1 << kEntryBits) - 1); + const int shift = (packed_index >> kEntryBits) & 1; + const int slot = group * kVectorsPerGroup + vector; // 0..31 + payload[kLowOffset + slot] = static_cast(entry & 0xFF); + // Two vectors share a qh byte: low nibble first, each holding three index-high bits + // and one shift bit. + const int nibble = (slot % 4) % 2; + const int qh_index = 2 * (slot / 4) + (slot % 4) / 2; + payload[kHighOffset + qh_index] |= + static_cast((((entry >> 8) & 0x7) | (shift << 3)) << (4 * nibble)); + } + __syncthreads(); + } + } + + // Four scale words: each carries two sub-blocks' 3-bit local scales in bits 0..11 and one + // nibble of the FP16 block scale in bits 12..15. + if (tid < kScaleWords) { + uint32_t word = 0; +#pragma unroll + for (int parity = 0; parity < 2; ++parity) { + const int sub = 2 * tid + parity; + const int base = 6 * parity; + word |= static_cast(locals[2 * sub]) << base; + word |= static_cast(locals[2 * sub + 1]) << (base + 3); + } + word |= static_cast((d_bits >> (4 * tid)) & 0xF) << 12; + payload[kScaleWordOffset + 2 * tid] = static_cast(word); + payload[kScaleWordOffset + 2 * tid + 1] = static_cast(word >> 8); + } +} + +} // namespace + +at::Tensor iq1_m_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(), + "inputs must be contiguous"); + check_pack_inputs("IQ1_M", input, grid, kEntries); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == num_blocks, + "scales must be float16 [numel / 256]"); + TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_m_pack", [&] { + encode<<(num_blocks), kThreads, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + reinterpret_cast(scales.data_ptr()), + output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu index 666176e951f..cfe0aba39f3 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu @@ -41,22 +41,6 @@ constexpr float kScaleAnchor = 0.61f; static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two"); -__device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x, - const float *q, float scale, float delta) { - float dot = 0.0f; - float qnorm = 0.0f; - float qsum = 0.0f; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - dot = fmaf(x[j], q[j], dot); - qnorm = fmaf(q[j], q[j], qnorm); - qsum += q[j]; - } - const float shifted_dot = dot + delta * xsum; - const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; - return clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale); -} - template __global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; @@ -104,39 +88,23 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g #pragma unroll for (int vector = 0; vector < kVectorsPerGroup; ++vector) { float x[kVectorSize]; - float xnorm = 0.0f; - float xsum = 0.0f; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - xsum += x[j]; - } + float xnorm, xsum; + load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, + xnorm, xsum); float local_best[kChoices]; #pragma unroll for (int choice = 0; choice < kChoices; ++choice) local_best[choice] = FLT_MAX; for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float *q = grid + entry * kVectorSize; - float dot = 0.0f; - float qnorm = 0.0f; - float qsum = 0.0f; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - dot = fmaf(x[j], q[j], dot); - qnorm = fmaf(q[j], q[j], qnorm); - qsum += q[j]; - } + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); #pragma unroll for (int choice = 0; choice < kChoices; ++choice) { const int local = choice & 7; const float delta = choice < 8 ? kDelta : -kDelta; const float scale = d * (2 * local + 1); - const float shifted_dot = dot + delta * xsum; - const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; - local_best[choice] = fminf(local_best[choice], - clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale)); + local_best[choice] = + fminf(local_best[choice], shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, delta)); } } block_min_accumulate(local_best, warp_best, group_error); @@ -161,19 +129,15 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g #pragma unroll for (int vector = 0; vector < kVectorsPerGroup; ++vector) { float x[kVectorSize]; - float xnorm = 0.0f; - float xsum = 0.0f; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - xsum += x[j]; - } + float xnorm, xsum; + load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, + xnorm, xsum); unsigned long long key = ~0ULL; for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); const float error = - quant_error(xnorm, xsum, x, grid + entry * kVectorSize, selected_scale, selected_delta); + shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, selected_delta); const unsigned long long candidate = error_key(error, entry); key = candidate < key ? candidate : key; } diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 087ad33c4d6..590569a6f63 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -85,7 +85,8 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False): The formats share their packing helpers, CUDA version requirement, and build flags, so they build as a single extension: ``iq1_s_pack(input, grid)``, ``iq2_xs_pack(input, grid, scales)``, ``iq2_xxs_pack(input, grid, scales)``, - ``iq2_s_pack(input, grid, scales)``, and ``q8_0_pack(input)``. + ``iq2_s_pack(input, grid, scales)``, ``iq1_m_pack(input, grid, scales)``, and + ``q8_0_pack(input)``. """ if not hasattr(get_cuda_ext_ggml, "extension") or ( raise_if_failed and get_cuda_ext_ggml.extension is None @@ -98,6 +99,7 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False): kernels_ggml / "iq2_xs.cu", kernels_ggml / "iq2_xxs.cu", kernels_ggml / "iq2_s.cu", + kernels_ggml / "iq1_m.cu", kernels_ggml / "q8_0.cu", ], cuda_version_specifiers=">=11.8", diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index 4843e7e5f05..736b01ea7a9 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -17,6 +17,8 @@ # Importing the backend installs its TensorQuantizer dispatch entry. from . import backend as _backend +from .iq1_m import * +from .iq1_m import __all__ as _iq1_m_all from .iq1_s import * from .iq1_s import __all__ as _iq1_s_all from .iq2_s import * @@ -28,6 +30,7 @@ from .registry import IQ_FORMAT_REGISTRY, IQFormat __all__ = [ # noqa: PLE0604 + *_iq1_m_all, *_iq1_s_all, *_iq2_s_all, *_iq2_xs_all, diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index 1c91b8f50ad..31fde77c4c2 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -39,8 +39,10 @@ import torch +from ..extensions import get_cuda_ext_ggml from .common import ( GGML_BLOCK_SIZE, + IQFormat, narrow_to_float32, validate_block_chunk_size, validate_packed_weights, @@ -53,6 +55,7 @@ "IQ1_M_BLOCK_SIZE", "IQ1_M_EFFECTIVE_BITS", "dequantize_iq1_m", + "iq1_m_fake_quant", "iq1_m_grid", "quantize_iq1_m", ] @@ -74,6 +77,7 @@ _IQ1_M_SUBBLOCKS = 8 _DEFAULT_BLOCK_CHUNK_SIZE = 1024 _DEFAULT_DECODE_CHUNK_SIZE = 4096 +_SCALE_BLOCK_CHUNK_SIZE = 4096 def iq1_m_grid(device: torch.device | str | None = None) -> torch.Tensor: @@ -156,6 +160,16 @@ def quantize_iq1_m( blocks = weight.contiguous().reshape(-1, IQ1_M_BLOCK_SIZE) grid = iq1_s_grid(weight.device) packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ1_M_BLOCK_SIZE, IQ1_M_BLOCK_BYTES) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + scale_chunks = [ + _predict_iq1_m_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) + for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) + ] + packed = extension.iq1_m_pack(blocks, grid, torch.cat(scale_chunks)) + 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) @@ -217,3 +231,19 @@ def dequantize_iq1_m( chunk_decoded = values * group_scale.unsqueeze(-1) decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_M_BLOCK_SIZE) return decoded.reshape(shape) + + +IQ1_M_FORMAT = IQFormat( + name="iq1_m", + block_size=IQ1_M_BLOCK_SIZE, + block_bytes=IQ1_M_BLOCK_BYTES, + quantize=quantize_iq1_m, + dequantize=dequantize_iq1_m, + block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, +) + +# Kept for callers of the per-format entry point. The record captured quantize_iq1_m and +# dequantize_iq1_m when it was built, so patching those module functions changes neither backend +# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY. +iq1_m_fake_quant = IQ1_M_FORMAT.fake_quant diff --git a/modelopt/torch/quantization/ggml/registry.py b/modelopt/torch/quantization/ggml/registry.py index c1e880aa957..8183b3a93ea 100644 --- a/modelopt/torch/quantization/ggml/registry.py +++ b/modelopt/torch/quantization/ggml/registry.py @@ -16,6 +16,7 @@ """The GGML IQ formats, listed once for backend dispatch and export.""" from .common import IQFormat +from .iq1_m import IQ1_M_FORMAT from .iq1_s import IQ1_S_FORMAT from .iq2_s import IQ2_S_FORMAT from .iq2_xs import IQ2_XS_FORMAT @@ -30,5 +31,6 @@ # It is an explicit list rather than formats registering themselves on import, so its contents # never depend on which modules happen to have been imported first. IQ_FORMAT_REGISTRY: dict[str, IQFormat] = { - fmt.name: fmt for fmt in (IQ1_S_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, IQ2_S_FORMAT) + fmt.name: fmt + for fmt in (IQ1_S_FORMAT, IQ1_M_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, IQ2_S_FORMAT) } diff --git a/modelopt_recipes/configs/numerics/iq1_m.yaml b/modelopt_recipes/configs/numerics/iq1_m.yaml new file mode 100644 index 00000000000..606aae477ff --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq1_m.yaml @@ -0,0 +1,26 @@ +# 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_M weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq1_m +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq1_m", so the generic estimator cannot derive the +# storage cost: 56 packed bytes * 8 / 256 weights. Keep in sync with IQ1_M_BLOCK_BYTES. +effective_bits: 1.75 +block_sizes: + -1: 256 +backend: ggml diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml new file mode 100644 index 00000000000..3c152baa0ec --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml @@ -0,0 +1,32 @@ +# 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. + +# QuantizeConfig preset for IQ1_M weight-only quantization. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq1_m: configs/numerics/iq1_m + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq1_m + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq1_m.yaml b/modelopt_recipes/general/ptq/iq1_m.yaml new file mode 100644 index 00000000000..40e8bbf7c41 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_m.yaml @@ -0,0 +1,27 @@ +# 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_M weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq1_m + +metadata: + description: >- + Applies uniform GGML-compatible IQ1_M weight-only quantization to eligible linear layers. + This is not a mixed per-tensor precision preset. No calibration data is required. +quantize: + $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index d8a59188a9b..66ad8759a0b 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 30 general/ptq/ recipes (click to expand) +All 31 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -59,6 +59,7 @@ supported combinations. | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | | `iq1_s` | IQ1_S W1A16, eligible linears | none | none (no calibration) | +| `iq1_m` | IQ1_M W1A16 (1.75 bpw), eligible linears | none | none (no calibration) | | `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), eligible linears | none | none (no calibration) | | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), eligible linears | none | none (no calibration) | | `iq2_s` | IQ2_S W2A16 (2.56 bpw), eligible linears | none | none (no calibration) | @@ -140,10 +141,10 @@ activations and tensor-core math are what deliver the throughput. - **`mxfp4_mlp_weight_only`** — MXFP4 weights on MLP/MoE layers only, BF16 activations. Needs no calibration forward pass; the QAT starting point for the GPT-OSS family (see `examples/gpt-oss`). -- **`iq1_s` / `iq2_xxs` / `iq2_xs` / `iq2_s`** — GGML-compatible IQ weights +- **`iq1_s` / `iq1_m` / `iq2_xxs` / `iq2_xs` / `iq2_s`** — GGML-compatible IQ weights on the eligible linear layers, with BF16 activations; `lm_head`, MoE routers, `conv1d` and the vision branch stay in BF16 like every other preset. The formats - trade size against accuracy in order: 1.56, 2.06, 2.31 and 2.56 bits per weight. No calibration data is + trade size against accuracy in order: 1.56, 1.75, 2.06, 2.31 and 2.56 bits per weight. No calibration data is required. Quantized weights must have a final dimension divisible by 256. Unified HF export writes the packed GGML blocks; Megatron export additionally requires tensor and pipeline parallel sizes of 1, and does not support diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 1242e9e0226..2282e55d81a 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -76,13 +76,14 @@ def test_ptq_whisper(command): PTQCommand(quant="int8_weight_only", kv_cache_quant="none"), PTQCommand(quant="int4_awq", kv_cache_quant="none"), PTQCommand(quant="w4a8_awq_beta", kv_cache_quant="none"), - # GGML IQ weight-only, recipe-driven: four formats between 1.56 and 2.56 bits - # per weight. These encoders require every weight's input dimension to be a multiple of + # GGML IQ weight-only, recipe-driven: the five formats between 1.56 and 2.56 bits per + # weight. These encoders require every weight's input dimension to be a multiple of # 256; TinyLlama's 2048 and 5632 both are. None of them calibrates -- every recipe # sets algorithm: null -- so the only IQ-specific cost is packing each weight once on # a CUDA encoder and decoding it on each forward, which fits the 300s # tests/examples default. PTQCommand(recipe="general/ptq/iq1_s", kv_cache_quant="none"), + PTQCommand(recipe="general/ptq/iq1_m", kv_cache_quant="none"), PTQCommand(recipe="general/ptq/iq2_xxs", kv_cache_quant="none"), PTQCommand(recipe="general/ptq/iq2_xs", kv_cache_quant="none"), PTQCommand(recipe="general/ptq/iq2_s", kv_cache_quant="none"), diff --git a/tests/gpu/torch/quantization/test_iq_formats_cuda.py b/tests/gpu/torch/quantization/test_iq_formats_cuda.py index 8dfb99372ad..d0756cd9b6b 100644 --- a/tests/gpu/torch/quantization/test_iq_formats_cuda.py +++ b/tests/gpu/torch/quantization/test_iq_formats_cuda.py @@ -23,12 +23,14 @@ import pytest import torch +import modelopt.torch.quantization.ggml.iq1_m as iq1_m_module import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module import modelopt.torch.quantization.ggml.iq2_s as iq2_s_module import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module import modelopt.torch.quantization.ggml.iq2_xxs as iq2_xxs_module from modelopt.torch.quantization.extensions import get_cuda_ext_ggml from modelopt.torch.quantization.ggml import ( + IQ1_M_BLOCK_BYTES, IQ1_S_BLOCK_BYTES, IQ2_S_BLOCK_BYTES, IQ2_XS_BLOCK_BYTES, @@ -39,6 +41,7 @@ # module, packer name, per-block payload size, whether the packer takes precomputed scales FORMATS = { "iq1_s": (iq1_s_module, "iq1_s_pack", IQ1_S_BLOCK_BYTES, False), + "iq1_m": (iq1_m_module, "iq1_m_pack", IQ1_M_BLOCK_BYTES, True), "iq2_xxs": (iq2_xxs_module, "iq2_xxs_pack", IQ2_XXS_BLOCK_BYTES, True), "iq2_xs": (iq2_xs_module, "iq2_xs_pack", IQ2_XS_BLOCK_BYTES, True), "iq2_s": (iq2_s_module, "iq2_s_pack", IQ2_S_BLOCK_BYTES, True), diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index ff6d6060c60..f002bcc3c59 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -34,6 +34,8 @@ from modelopt.torch.opt.config_loader import BUILTIN_CONFIG_ROOT from modelopt.torch.quantization.config import LocalHessianCalibConfig, QuantizeConfig from modelopt.torch.quantization.ggml import ( + IQ1_M_BLOCK_SIZE, + IQ1_M_EFFECTIVE_BITS, IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS, IQ2_S_BLOCK_SIZE, @@ -139,6 +141,7 @@ def test_mlp_weight_only_recipe_matches_its_mtq_cfg(recipe_name, cfg_name): ("qformat", "block_size", "effective_bits"), [ ("iq1_s", IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS), + ("iq1_m", IQ1_M_BLOCK_SIZE, IQ1_M_EFFECTIVE_BITS), ("iq2_xxs", IQ2_XXS_BLOCK_SIZE, IQ2_XXS_EFFECTIVE_BITS), ("iq2_xs", IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS), ("iq2_s", IQ2_S_BLOCK_SIZE, IQ2_S_EFFECTIVE_BITS),