From 0ec9d71299242090b601b260110ff746392f4648 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 18:33:41 +0000 Subject: [PATCH 1/6] Add the IQ1_M, IQ2_XXS and IQ2_S weight-only formats llama.cpp defines five IQ formats at one and two bits; we shipped two of them. On a real mixed-precision checkpoint (unsloth/Qwen3.8-27B-GGUF) the three missing ones cover 93 tensors and 4.66B parameters, 17.3% of the file, so a reader limited to IQ1_S and IQ2_XS cannot consume it. Each format follows the existing single-pass grid search at a fixed anchored super-block scale. They differ in ways worth naming: * IQ2_XXS packs a 4-bit sub-block scale into the same word as four 7-bit sign indices, and its 256-entry grid needs no high index bits. * IQ2_S stores a full 8-bit sign mask per group rather than the 7-bit parity-coded index, so the encoder takes the input signs directly instead of flipping the weakest element to fix parity. * IQ1_M has no ``d`` field at all: the FP16 super-block scale is reassembled from the top nibble of each of four scale words. It is also finer grained than IQ1_S, with a local scale per two groups and a delta shift per group. Export previously spelled the family as a two-element tuple at nine sites. Replace those with an IQ_FORMATS frozenset plus per-format packer and block geometry tables, so a sixth format is a row rather than a sweep. Decoders are validated against the llama.cpp checkpoint itself, not just round-tripped: IQ2_XXS over 11,100,160 blocks, IQ1_M over 4,730,880 and IQ2_S over 2,355,200, all bit-identical to dequantize_row_* output. The new codebooks match the ggml-common.h tables entry for entry. Blocks lifted from that checkpoint ship as conformance vectors so CI keeps checking bytes we did not produce; mutation testing confirms they catch a mis-set scale nibble and a wrong sign-field width. There is no CUDA encoder for these three yet, so packing runs the torch search. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/quant_format.py | 20 +- modelopt/torch/export/quant_utils.py | 39 ++- modelopt/torch/export/unified_export_hf.py | 28 +- .../torch/export/unified_export_megatron.py | 42 ++- modelopt/torch/quantization/ggml/__init__.py | 14 +- modelopt/torch/quantization/ggml/backend.py | 24 +- modelopt/torch/quantization/ggml/codebooks.py | 52 ++++ modelopt/torch/quantization/ggml/iq1_m.py | 271 +++++++++++++++++ modelopt/torch/quantization/ggml/iq2_s.py | 241 +++++++++++++++ modelopt/torch/quantization/ggml/iq2_xxs.py | 275 ++++++++++++++++++ .../quantization/iq_llama_cpp_vectors.py | 239 +++++++++++++++ .../torch/quantization/test_ggml_backend.py | 3 +- .../test_iq_llama_cpp_conformance.py | 149 ++++++++++ 13 files changed, 1359 insertions(+), 38 deletions(-) create mode 100644 modelopt/torch/quantization/ggml/iq1_m.py create mode 100644 modelopt/torch/quantization/ggml/iq2_s.py create mode 100644 modelopt/torch/quantization/ggml/iq2_xxs.py create mode 100644 tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py create mode 100644 tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index d1bb213b13d..6146c6746da 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -37,16 +37,30 @@ 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" + +# Every GGML IQ format. They share the weight-only, 256-value-block, per-module-scale +# shape, so export treats them as one family; adding a format means adding it here +# rather than extending a tuple at each use site. +IQ_FORMATS = frozenset( + { + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ1_M, + QUANTIZATION_IQ2_XXS, + QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_S, + } +) # Formats whose scales are purely per-module, so export never merges them across the q/k/v # and gate/up groups that share an input. Every other format unifies input_amax (and, for # NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. -FUSION_FREE_FORMATS = frozenset( +FUSION_FREE_FORMATS = IQ_FORMATS | frozenset( { QUANTIZATION_FP8, - QUANTIZATION_IQ1_S, - QUANTIZATION_IQ2_XS, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL, } diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index c36369bab19..16caaa8b3c7 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -28,12 +28,21 @@ from modelopt import __version__ from modelopt.torch.models import get_spec, list_all_possible from modelopt.torch.quantization.ggml import ( + IQ1_M_BLOCK_BYTES, + IQ1_M_BLOCK_SIZE, + IQ1_M_EFFECTIVE_BITS, IQ1_S_BLOCK_BYTES, IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS, + IQ2_S_BLOCK_BYTES, + IQ2_S_BLOCK_SIZE, + IQ2_S_EFFECTIVE_BITS, IQ2_XS_BLOCK_BYTES, IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS, + IQ2_XXS_BLOCK_BYTES, + IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_EFFECTIVE_BITS, ) from modelopt.torch.quantization.model_calib import ( enable_stats_collection, @@ -59,6 +68,7 @@ from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer from .model_utils import TiedWeightMap, get_language_model_from_vl from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_INT8, @@ -71,8 +81,11 @@ QUANTIZATION_INT4_AWQ, QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO, + QUANTIZATION_IQ1_M, QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -88,6 +101,16 @@ logger = logging.getLogger(__name__) +# Block geometry per GGML IQ format, used to describe the packed payload in the checkpoint. +_IQ_BLOCK_METADATA = { + QUANTIZATION_IQ1_S: (IQ1_S_BLOCK_SIZE, IQ1_S_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS), + QUANTIZATION_IQ1_M: (IQ1_M_BLOCK_SIZE, IQ1_M_BLOCK_BYTES, IQ1_M_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XXS: (IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, IQ2_XXS_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XS: (IQ2_XS_BLOCK_SIZE, IQ2_XS_BLOCK_BYTES, IQ2_XS_EFFECTIVE_BITS), + QUANTIZATION_IQ2_S: (IQ2_S_BLOCK_SIZE, IQ2_S_BLOCK_BYTES, IQ2_S_EFFECTIVE_BITS), +} + + def _has_large_fp8_scale(value: torch.Tensor) -> bool: """Return whether an FP8 scale exceeds the recommended threshold.""" return bool(torch.any(value > 0.5)) @@ -474,8 +497,7 @@ def uses_iq_quantization(module) -> bool: if ( weight_quantizer is not None and weight_quantizer.is_enabled - and getattr(weight_quantizer, "num_bits", None) - in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + and getattr(weight_quantizer, "num_bits", None) in IQ_FORMATS ): return True return any(uses_iq_quantization(child) for _, child in module.named_children()) @@ -515,7 +537,7 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames return QUANTIZATION_W4A8_AWQ # Handle individual num_bits cases - if weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if weight_quantizer.num_bits in IQ_FORMATS: if weight_quantizer.backend != "ggml": raise ValueError("IQ formats require the built-in 'ggml' quantization backend") # Both exporters return before collecting input_scale and before the pre_quant_scale @@ -781,15 +803,8 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } - elif v in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): - if v == QUANTIZATION_IQ1_S: - block_size = IQ1_S_BLOCK_SIZE - payload_bytes = IQ1_S_BLOCK_BYTES - effective_bits = IQ1_S_EFFECTIVE_BITS - else: - block_size = IQ2_XS_BLOCK_SIZE - payload_bytes = IQ2_XS_BLOCK_BYTES - effective_bits = IQ2_XS_EFFECTIVE_BITS + elif v in IQ_FORMATS: + block_size, payload_bytes, effective_bits = _IQ_BLOCK_METADATA[v] if block_size_value != block_size: raise ValueError( f"{v.upper()} requires block size {block_size}, got {block_size_value}" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 28edade5d7b..48144954abb 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -67,7 +67,13 @@ from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs +from modelopt.torch.quantization.ggml import ( + quantize_iq1_m, + quantize_iq1_s, + quantize_iq2_s, + quantize_iq2_xs, + quantize_iq2_xxs, +) from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper @@ -101,11 +107,15 @@ ) from .quant_format import ( FUSION_FREE_FORMATS, + IQ_FORMATS, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, + QUANTIZATION_IQ1_M, QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP8, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -134,6 +144,16 @@ ) from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry +# One packer per GGML IQ format; the checkpoint stores the packed blocks directly. +_IQ_PACKERS = { + QUANTIZATION_IQ1_S: quantize_iq1_s, + QUANTIZATION_IQ1_M: quantize_iq1_m, + QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, + QUANTIZATION_IQ2_XS: quantize_iq2_xs, + QUANTIZATION_IQ2_S: quantize_iq2_s, +} + + __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] @@ -630,15 +650,13 @@ def _export_quantized_weight( "which dispatches to the streaming writer that materialises weights layer-by-layer." ) - if quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if quantization_format in IQ_FORMATS: if weight_name != "weight": raise NotImplementedError( "IQ unified export currently supports modules with a standard 'weight' " f"attribute, got {weight_name!r} on {type(sub_module).__name__}" ) - quantize_iq = ( - quantize_iq1_s if quantization_format == QUANTIZATION_IQ1_S else quantize_iq2_xs - ) + quantize_iq = _IQ_PACKERS[quantization_format] packed_weight, _ = quantize_iq(weight.to(dtype)) setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) maybe_clear_cuda_cache() diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 215a62ff85e..0be18fda79b 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,7 +35,13 @@ from safetensors.torch import save_file from modelopt import __version__ -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs +from modelopt.torch.quantization.ggml import ( + quantize_iq1_m, + quantize_iq1_s, + quantize_iq2_s, + quantize_iq2_xs, + quantize_iq2_xxs, +) from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -57,13 +63,17 @@ ) from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PB_WO, + QUANTIZATION_IQ1_M, QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, @@ -85,6 +95,16 @@ import transformers from transformers import AutoProcessor +# One packer per GGML IQ format, mirroring the HF exporter's table. +_IQ_PACKERS = { + QUANTIZATION_IQ1_S: quantize_iq1_s, + QUANTIZATION_IQ1_M: quantize_iq1_m, + QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, + QUANTIZATION_IQ2_XS: quantize_iq2_xs, + QUANTIZATION_IQ2_S: quantize_iq2_s, +} + + has_mcore = False with import_plugin("megatron"): from megatron.core.models.gpt import GPTModel @@ -351,7 +371,7 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" - elif quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + elif quantization_format in IQ_FORMATS: quantization = quantization_format.upper() if is_last_stage_main_rank: @@ -1115,7 +1135,7 @@ def _get_quantized_state( self._record_excluded_module(prefix) block_size = get_weight_block_size(module) - is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + is_iq = qformat in IQ_FORMATS name_to_value = self._get_weight_bias( module, dtype, name_to_value, keep_weight_device=is_iq ) @@ -1185,7 +1205,7 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str): @staticmethod def _pack_iq_weight(weight: torch.Tensor, qformat: str) -> torch.Tensor: """Pack one ``[out, in]`` weight and return its CPU payload.""" - quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs + quantize_iq = _IQ_PACKERS[qformat] packed_weight, _ = quantize_iq(weight) return packed_weight.detach().cpu() @@ -1210,7 +1230,7 @@ def _reject_unsupported_fused_iq_export(qformat: str) -> None: The one gap left is a rank holding no local expert at all, which needs expert-parallel size to exceed the expert count. Worth revisiting if that becomes a supported topology. """ - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: raise NotImplementedError( "Fused-MoE IQ export requires a deployment loader that supports " "[num_experts, out_features, in_features // 256, payload_bytes]" @@ -1280,7 +1300,7 @@ def _name_remapping( weight = weight + 1.0 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat)) elif weight_scale is None: self._state_dict[prefix + "weight"] = weight @@ -1327,7 +1347,7 @@ def _gated_mlp_slicing( gate_proj_weight = weight[:ffn_hidden_size, :] up_proj_weight = weight[ffn_hidden_size:, :] - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: self._state_dict.update( self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat) ) @@ -1501,7 +1521,7 @@ def _grouped_mlp_slicing( seen_qformat, seen_block_size = qformat, block_size weight = state_dict[weight_key].to(self.dtype) - if qformat not in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat not in IQ_FORMATS: weight = weight.cpu() weight_scale_cpu = ( weight_scale.detach().cpu().clone() if weight_scale is not None else None @@ -1533,7 +1553,7 @@ def _grouped_mlp_slicing( ] for shard_prefix, shard_weight, shard_scale in shards: - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: local_expert_state.update( self._get_iq_weight_state( shard_prefix + "weight", shard_weight, qformat @@ -1702,7 +1722,7 @@ def _take(tensor, index, last_dim, with_gate=False): proj_weights = [_take(weight, s, hidden_size, g) for s, g in zip(slices, gated)] proj_keys = [p + "weight" for p in prefixes] - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: for key, weight in zip(proj_keys, proj_weights): self._state_dict.update(self._get_iq_weight_state(key, weight, qformat)) elif weight_scale is None: @@ -1820,7 +1840,7 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): proj_keys = [p + "weight" for p in proj_prefixes] weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: for proj_prefix, proj_weight in zip(proj_prefixes, proj_weights): if proj_prefix in keep_bf16: self._state_dict[proj_prefix + "weight"] = proj_weight.cpu() diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index 0fbb63ee405..f56d22fd5d4 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -17,9 +17,21 @@ # 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 * +from .iq2_s import __all__ as _iq2_s_all from .iq2_xs import * from .iq2_xs import __all__ as _iq2_xs_all +from .iq2_xxs import * +from .iq2_xxs import __all__ as _iq2_xxs_all -__all__ = [*_iq1_s_all, *_iq2_xs_all] # noqa: PLE0604 +__all__ = [ # noqa: PLE0604 + *_iq1_m_all, + *_iq1_s_all, + *_iq2_s_all, + *_iq2_xs_all, + *_iq2_xxs_all, +] diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py index dc8e15ffb04..1d6fa055d3d 100644 --- a/modelopt/torch/quantization/ggml/backend.py +++ b/modelopt/torch/quantization/ggml/backend.py @@ -18,8 +18,20 @@ import torch from ..nn.modules.tensor_quantizer import register_quant_backend +from .iq1_m import iq1_m_fake_quant from .iq1_s import iq1_s_fake_quant +from .iq2_s import iq2_s_fake_quant from .iq2_xs import iq2_xs_fake_quant +from .iq2_xxs import iq2_xxs_fake_quant + +# One entry per GGML IQ format; adding a format is adding a row here. +_FAKE_QUANTS = { + "iq1_m": iq1_m_fake_quant, + "iq1_s": iq1_s_fake_quant, + "iq2_s": iq2_s_fake_quant, + "iq2_xs": iq2_xs_fake_quant, + "iq2_xxs": iq2_xxs_fake_quant, +} def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: @@ -29,11 +41,13 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: unknown_args = set(extra_args) - {"block_chunk_size", "decode_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'") + # 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. + fake_quant = _FAKE_QUANTS.get(num_bits) if isinstance(num_bits, str) else None + if fake_quant is None: + supported = ", ".join(repr(name) for name in sorted(_FAKE_QUANTS)) + raise ValueError(f"The ggml backend requires num_bits in ({supported})") + return fake_quant(inputs, quantizer, **extra_args) register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/codebooks.py b/modelopt/torch/quantization/ggml/codebooks.py index 04402f4b1d0..22b14decb17 100644 --- a/modelopt/torch/quantization/ggml/codebooks.py +++ b/modelopt/torch/quantization/ggml/codebooks.py @@ -166,6 +166,46 @@ ) +# Compact byte representation of the canonical [256, 8] IQ2_XXS grid. Same +# 8/25/43 magnitude alphabet as IQ2_XS, so it compresses well. +_IQ2_XXS_GRID_ZLIB_B64 = ( + "eNqFVVuS5CAM++UKOoPuf78ZLMmYLLXT1SkSYvyQZGct/egVuDfoFQtPA34M6RXLB74HqRW/N9rm5ZBZmWc5hK+vY+ATII55Oz4X" + "c1+OMZ2PANgWtccrIBOYqYzwEdoH7cfP5EyQSZRJmNpHlbLNMZHCBghKgy/kfh/pOLPgxOgrwUkXoRKrCjo7OpsJlGOXAY+LMhvQ" + "fIGM8/OST2CBvwAWw4JlAk5ZIEC/AZZ7Gz0A54Z5AJ4rkTr1hYuAijSIEEdlXMqryIOQ7paoHJEMMaVPr5LNJGRq90isoTxEVbGC" + "6xCHFaiptZQ2CXVH7Jcidiuqkf2HaGMR8XTsdXUBrOaA1wwdgeAjFOca+ZsFwOu65FtxRlUllEtQfAsLlCKiobobAotS+sSnwyXK" + "0elKQHnT7Wo7wXjB2T3XGr/hOQJmhKxuiaqLvuNeOMKh91qGAl9aAOZkESSOW4KnIS9sJL3/NQJ5NYRiUX9pGrNBgEysAb4ahXru" + "STY/O/mcdAJRbDOSXlzd1WuOymbMktGQnRrkNQF7LBgFz106e/rDkjZmRpFm1KMx3WvRuqemoa19esbHuBgw7VKggELYclb7VH91" + "Xf5hzoD21KEk6CpbkT/vBJbU" +) + +# Compact byte representation of the canonical [1024, 8] IQ2_S grid. +_IQ2_S_GRID_ZLIB_B64 = ( + "eNqFmWF25DAIg//6CjoD979fG5sPI2bS7r48tttJYoOQhGet8ycySv4fkVFL9oHIqPAbIqPkD9DKB+SDIuPvx+zBkVHyF2mFvTAy" + "KnwBkfH3LluQlr4u7PfHc1/kz3mx4Mj43NU3EBmfZfYNacXXjUXG5+N9o1qyDUdGhSdAykRE/rwyISMxkVHKuOJrorieh/XEaelr" + "Ap/0RKalJzQyPtvtCX620RP9PO48xxMf6yRePflsjpe1gkTGnZ1WoL3aVqjIuN/SChcZd9laISOjlqywkXGnvxU6Mu40t6UBgMio" + "cEDstJx151Yi7wsDDLXhAkCRUUsGqMi4l9MAtl+nk3i11AG8yLg/1oAYGSUHplZ8BWhklBKwkT8vfQUwpaPWXAoHtuQAlxzoWnLA" + "JzPsVP3RAGCe3HIBktU3Ffny1igb3XH6IqGn3kCRcVe5NVRk3NmMzF7kbiN32Rpv/zayrSJ5qkF9NmRk3LBtDbrh2Bp0wyUSJpFw" + "sMYN9VbaC2sNrSVr7FhhDR6j0Z+7eysuIy/VHpaxqAqsgEXGQsrFXYKgFJIThgz9oBh0XiKJ1VEFii51LKvWJRzJiUeZpUlAQERy" + "QiKn/IWgIuOhmUtYkfHQySUw7iQFkfHQwCW2yHja+xIdb5yEB6RPO14CZGUQISsk9WCSmh7Y3x1MgjywvTuDMNnhG2HSYnAiPUxG" + "VicnSKQRK00BwZLB1Tddi7vEGxlPGi8Rk/lJyFAAt5403ApNwqZiq3Ehvd2JnNxyQexUehI8FMTSIXwQMYkfhLwJAFuHM0HSFASQ" + "9SYMD7915NGaoWY2mnCwVy6QunqTxi3JTlRD8hWYMGQXFTRKVispmkAHTGGiI6ZAPSp0nhd/CtWT3C5YSAJQQovouNVFt1zgFTbI" + "PtJ10aFT8AAloJjCd9TnCmCkANLhUwjp9cLickFE4qollr4KJJbqTShLg/LvWWikb42UmC6k+Yna0vkkTAS1h5q5bYJLT2FaEPu6" + "7hsv+TaBhkymUJ/kd8F2Zw1FUAKYcgo6zDmFPXbbXoHHil6hlwm+Nlyv8Bc1rWZGGzNPY3CTfH6AuUsadzridEf3DmAV43AfsaN8" + "LGUMZegog0FF5WMkvVXKsdzklBEBCfLxDvEpxVlOipBaGReQVGTg41YZGopFMUhmJcvHopIGWk8+xqDic0xBdVVY8bGiqE4+NhTy" + "5eNAKes0VHQIbAu7yu17GS25Da+t0WHSMGAyO1zSSBfJ7W0ZNLmNLYMGGuR2s6Aqt5OlZXKbWJ0vt4OV4+W2rlp5uV3DnpURlNuu" + "MoQwy7BX006VUZTbo5Iwuf0pppLbmuLsaSypgdx+1COm4dxq35zTmwGFIeUyX0uaxlQuz60FNtC2unZnRibl8liMjEfCQ+DkpsEF" + "e0PuprzdE4IXIwzzy+WlUl7FyUXK5eDjpCFpfNL3q6FGceQ0WaVeTnOT1kpr5TRWXLKcliYN1Zbk9FKmbRp4FBHzgLWQ00CR9fK2" + "LZJY3o4FWkAmb6tKBi+Rt0ddc0CQw/djUJDDrFppDg5yWFTuHoHrg4S8fGVJ5WUp6MjTXZ5sDh6wnDw9tec5kMi38zGYgBr5a4s6" + "5sAi//XH4MI/5gATW5E/BxmcETdiQXnAHHR44Bx4zjBzTTkv+m8A2q72vD/6QMQCl51+3wGJVmfhzDRsYA5ObAguYmPvA5WfzGFS" + "SACiTCLmwEVi3gYv7S19HcCaa82EtMGMxDISkOA5qJHwt4GNQrwNblhUigSlVjVOx1UhS2nS2nDvTl0b9LaXjeNJesGXf/3C1y01" + "CAIEKIOlMKsVkpd9XVGDItwAYMAqNeGaAEIhARJKiblafpxeQKO0AK7Ey4+1C4iQHICk6bkKXFnkI+d3UKVD83i2gAylAuhS8NQQ" + "AL76SN4Av7pktgYAsjTC8mO9agw0nwYpLswlsrc7KPsJMmpKQ6F2NFZJRG6VRnsbrCObiuLPQVun2z4G7kLNsmODatzaba6SSZCG" + "LseTZmGTYBsXD4ceDvQj5TACwDExvi0fz4og3gZ5JlGIg5VBIHPQB3NvA/8kGqQZwpmDf9rfIqK3AwCYYvUvLa/9qkxCYPOgAEJ7" + "OzA4hwGX8P47QIAIIZs4m9+02IlxHjDESWXNuDrVrOE7V1EnAkmLdQABwfKtLXRW31LmlxAQ7zyoKCgdZihzelIZ9ztJjhxWPyWK" + "PCXi5nm6cf8Jwb8deBykxb2yrCUlmYaoL1/rjOq8XmeEyR6tI6z6mK+2LS0/ln9+AIHIjsQ=" +) + + @cache def iq1_s_grid_bytes() -> bytes: """Decoded little-endian int8 bytes of the [2048, 8] IQ1_S ternary table.""" @@ -176,3 +216,15 @@ def iq1_s_grid_bytes() -> bytes: def iq2_xs_grid_bytes() -> bytes: """Decoded bytes of the [512, 8] IQ2_XS magnitude table.""" return base64.b64decode(_IQ2_XS_GRID_B64) + + +@cache +def iq2_xxs_grid_bytes() -> bytes: + """Decoded bytes of the [256, 8] IQ2_XXS magnitude table.""" + return zlib.decompress(base64.b64decode(_IQ2_XXS_GRID_ZLIB_B64)) + + +@cache +def iq2_s_grid_bytes() -> bytes: + """Decoded bytes of the [1024, 8] IQ2_S magnitude table.""" + return zlib.decompress(base64.b64decode(_IQ2_S_GRID_ZLIB_B64)) diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py new file mode 100644 index 00000000000..c76de4d4618 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -0,0 +1,271 @@ +# 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 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, mirroring :mod:`.iq1_s`. Every 256 +logical values become one 56-byte block_iq1_m payload: + +* bytes 0..31: 32 low bytes of the grid index, four per sub-block +* bytes 32..47: 16 bytes holding, per nibble, three grid-index high bits and + one delta-shift bit -- two groups per byte +* bytes 48..55: four little-endian uint16 holding four 3-bit local scales each + in bits 0..11, and one nibble of the FP16 super-block scale in bits 12..15 + +IQ1_M has no dedicated ``d`` field: the FP16 super-block scale is reassembled +from the top nibble of each of the four scale words. It is also finer-grained +than IQ1_S -- a local scale covers two groups rather than four, and the delta +shift is chosen per group rather than per sub-block -- which is where its extra +0.1875 bits per weight go. + +The grid is the same canonical 2048 x 8 ternary table IQ1_S uses, 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#L2573-L2620 +""" + +import torch + +from .common import ( + GGML_BLOCK_SIZE, + fake_quantize_with_cache, + narrow_to_float32, + validate_block_chunk_size, + validate_packed_weights, + validate_weight, +) +from .iq1_s import iq1_s_grid + +__all__ = [ + "IQ1_M_BLOCK_BYTES", + "IQ1_M_BLOCK_SIZE", + "IQ1_M_EFFECTIVE_BITS", + "dequantize_iq1_m", + "iq1_m_fake_quant", + "quantize_iq1_m", +] + +IQ1_M_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ1_M_BLOCK_BYTES = 56 +IQ1_M_EFFECTIVE_BITS = IQ1_M_BLOCK_BYTES * 8 / IQ1_M_BLOCK_SIZE +_IQ1_M_DELTA = 0.125 +# Largest representable magnitude: (1 + delta) at local scale 7 -> 15 * 1.125. +_IQ1_M_NATIVE_MAX = 16.875 +_IQ1_M_SCALE_ANCHOR = 0.61 +_IQ1_M_GROUPS = 32 +_IQ1_M_SUBBLOCKS = 8 +_DEFAULT_BLOCK_CHUNK_SIZE = 1024 +_DEFAULT_DECODE_CHUNK_SIZE = 4096 + + +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, _IQ1_M_GROUPS, 8) + xnorm = vectors.square().sum(dim=-1) + xsum = vectors.sum(dim=-1) + + amax = x.abs().amax(dim=1) + d = ((amax / _IQ1_M_NATIVE_MAX) * _IQ1_M_SCALE_ANCHOR).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + # Choice index is shift * 8 + local, as in IQ1_S, but here the shift is free per group. + shape = (block_count, _IQ1_M_GROUPS, 16) + best_error = torch.full(shape, torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros(shape, dtype=torch.int64, device=x.device) + grid_norm = grid.square().sum(dim=-1) + grid_sum = grid.sum(dim=-1) + + 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_M_DELTA if shift else _IQ1_M_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] + ) + + # Each group picks its own shift; only the 3-bit local scale is shared, over two groups. + per_shift_error = best_error.reshape(block_count, _IQ1_M_GROUPS, 2, 8) + per_shift_entry = best_entry.reshape(block_count, _IQ1_M_GROUPS, 2, 8) + local_error, shift_index = per_shift_error.min(dim=2) + local_entry = per_shift_entry.gather(2, shift_index.unsqueeze(2)).squeeze(2) + + pair_error = local_error.reshape(block_count, 16, 2, 8).sum(dim=2) + selected_local = pair_error.argmin(dim=-1) + group_local = selected_local.repeat_interleave(2, dim=1) + selected_entry = local_entry.gather(2, group_local.unsqueeze(-1)).squeeze(-1) + selected_shift = shift_index.gather(2, group_local.unsqueeze(-1)).squeeze(-1) + + packed = torch.empty((block_count, IQ1_M_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :32] = (selected_entry & 0xFF).to(torch.uint8) + + high = (selected_entry >> 8).reshape(block_count, _IQ1_M_SUBBLOCKS, 4) + shift_bits = selected_shift.reshape(block_count, _IQ1_M_SUBBLOCKS, 4) + qh = torch.empty((block_count, _IQ1_M_SUBBLOCKS, 2), dtype=torch.int64, device=x.device) + for half in range(2): + low_group, high_group = 2 * half, 2 * half + 1 + qh[:, :, half] = ( + (high[:, :, low_group] & 0x7) + | (shift_bits[:, :, low_group] << 3) + | ((high[:, :, high_group] & 0x7) << 4) + | (shift_bits[:, :, high_group] << 7) + ) + packed[:, 32:48] = qh.reshape(block_count, 16).to(torch.uint8) + + # Four scale words, each carrying two sub-blocks' local scales plus one nibble of d. + d_bits = d.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF + locals_by_sub = selected_local.reshape(block_count, _IQ1_M_SUBBLOCKS, 2) + words = torch.zeros((block_count, 4), dtype=torch.int64, device=x.device) + for word in range(4): + for parity in range(2): + sub = 2 * word + parity + base = 6 * parity + words[:, word] |= locals_by_sub[:, sub, 0] << base + words[:, word] |= locals_by_sub[:, sub, 1] << (base + 3) + words[:, word] |= ((d_bits >> (4 * word)) & 0xF) << 12 + packed[:, 48:56:2] = (words & 0xFF).to(torch.uint8) + packed[:, 49:56:2] = ((words >> 8) & 0xFF).to(torch.uint8) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_iq1_m( + 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_M blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 56]`` + and ``[weight.ndim]``. There is no CUDA encoder for this format yet, so + packing always runs the torch search. + """ + validate_weight(weight, "IQ1_M") + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, IQ1_M_BLOCK_SIZE) + grid = iq1_s_grid(weight.device) + 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_M_BLOCK_SIZE, IQ1_M_BLOCK_BYTES) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq1_m( + 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 IQ1_M payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ1_M_BLOCK_BYTES, format_name="IQ1_M" + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, IQ1_M_BLOCK_BYTES) + grid = iq1_s_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ1_M_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] + count = block_chunk.shape[0] + low = block_chunk[:, :32].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 4) + qh = block_chunk[:, 32:48].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 2) + words = block_chunk[:, 48:56:2].to(torch.int64) | ( + block_chunk[:, 49:56:2].to(torch.int64) << 8 + ) + + # The FP16 super-block scale is the four words' top nibbles, low word first. + d_bits = ( + (words[:, 0] >> 12) + | ((words[:, 1] >> 8) & 0x00F0) + | ((words[:, 2] >> 4) & 0x0F00) + | (words[:, 3] & 0xF000) + ) + d = d_bits.to(torch.int16).view(torch.float16).float() + + entries = torch.empty((count, _IQ1_M_SUBBLOCKS, 4), dtype=torch.int64, device=blocks.device) + deltas = torch.empty( + (count, _IQ1_M_SUBBLOCKS, 4), dtype=torch.float32, device=blocks.device + ) + for group in range(4): + byte = qh[:, :, group // 2] + nibble_shift = 4 * (group % 2) + entries[:, :, group] = low[:, :, group] | (((byte >> nibble_shift) & 0x7) << 8) + negative = (byte >> (nibble_shift + 3)) & 1 + deltas[:, :, group] = torch.where( + negative.bool(), -_IQ1_M_DELTA, torch.tensor(_IQ1_M_DELTA, device=blocks.device) + ) + + scales = torch.empty( + (count, _IQ1_M_SUBBLOCKS, 2), dtype=torch.float32, device=blocks.device + ) + for sub in range(_IQ1_M_SUBBLOCKS): + word = words[:, sub // 2] + base = 6 * (sub % 2) + scales[:, sub, 0] = d * (2 * ((word >> base) & 0x7) + 1).float() + scales[:, sub, 1] = d * (2 * ((word >> (base + 3)) & 0x7) + 1).float() + + values = grid[entries] + deltas.unsqueeze(-1) + # Groups 0,1 take the first local scale and groups 2,3 the second, so the repeat + # runs along the half axis: [h0, h0, h1, h1], not [h0, h1, h0, h1]. + group_scale = scales.repeat_interleave(2, dim=2) + chunk_decoded = values * group_scale.unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_M_BLOCK_SIZE) + return decoded.reshape(shape) + + +def iq1_m_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """IQ1_M weight backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq1_m": + raise ValueError("The ggml IQ1_M backend requires num_bits='iq1_m'") + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq1_m", + block_chunk_size=block_chunk_size, + decode_chunk_size=decode_chunk_size, + quantize=quantize_iq1_m, + dequantize=dequantize_iq1_m, + ) diff --git a/modelopt/torch/quantization/ggml/iq2_s.py b/modelopt/torch/quantization/ggml/iq2_s.py new file mode 100644 index 00000000000..25aff6232d0 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_s.py @@ -0,0 +1,241 @@ +# 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_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, mirroring :mod:`.iq2_xs`. Every 256 +logical values become one 82-byte block_iq2_s payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..33: 32 low bytes of the grid index, four per sub-block +* bytes 34..65: 32 sign masks, four per sub-block +* bytes 66..73: eight bytes holding the grid index high 2 bits, four per byte +* bytes 74..81: 16 four-bit local scales, two per byte + +IQ2_S stores a full eight-bit sign mask per group rather than the seven-bit +parity-coded index used by IQ2_XS and IQ2_XXS, so the encoder can take the +input signs directly instead of flipping the weakest element to fix parity. + +The canonical 1024 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#L2540-L2571 +""" + +import torch + +from .codebooks import iq2_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__ = [ + "IQ2_S_BLOCK_BYTES", + "IQ2_S_BLOCK_SIZE", + "IQ2_S_EFFECTIVE_BITS", + "dequantize_iq2_s", + "iq2_s_fake_quant", + "iq2_s_grid", + "quantize_iq2_s", +] + +IQ2_S_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ2_S_BLOCK_BYTES = 82 +IQ2_S_EFFECTIVE_BITS = IQ2_S_BLOCK_BYTES * 8 / IQ2_S_BLOCK_SIZE +_IQ2_S_GRID_ENTRIES = 1024 +_IQ2_S_LOCAL_SCALES = 16 +_IQ2_S_GROUPS = 32 +_IQ2_S_SUBBLOCKS = 8 +_IQ2_S_NATIVE_MAX = 43 * 31 / 8 +_IQ2_S_SCALE_ANCHOR_MIN = 0.65 +_IQ2_S_SCALE_ANCHOR_MAX = 0.92 +_IQ2_S_PEAK_TO_RMS_TAPER = 0.035 +# The grid is twice IQ2_XS's, so the same search tile costs twice the memory. +_DEFAULT_BLOCK_CHUNK_SIZE = 128 +_DEFAULT_DECODE_CHUNK_SIZE = 4096 + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +def iq2_s_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_S 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_s_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(_IQ2_S_GRID_ENTRIES, 8).to( + device=resolved_device + ) + return _GRID_CACHE[resolved_device] + + +def _predict_iq2_s_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)) + anchor_ratio = (1.0 - _IQ2_S_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( + _IQ2_S_SCALE_ANCHOR_MIN, _IQ2_S_SCALE_ANCHOR_MAX + ) + return ((amax / _IQ2_S_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + + +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, _IQ2_S_GROUPS, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + + d = _predict_iq2_s_scales(x) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + shape = (block_count, _IQ2_S_GROUPS, _IQ2_S_LOCAL_SCALES) + best_error = torch.full(shape, torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros(shape, dtype=torch.int64, device=x.device) + # All eight signs are storable, so the search compares magnitudes directly. + for entry_start in range(0, _IQ2_S_GRID_ENTRIES, 64): + grid_tile = grid[entry_start : entry_start + 64] + dot = (magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8)).sum(dim=-1) + tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1) + + for local in range(_IQ2_S_LOCAL_SCALES): + 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] + ) + + # One local scale covers two groups (16 values), as in IQ2_XS. + pair_error = best_error.reshape(block_count, 16, 2, _IQ2_S_LOCAL_SCALES).sum(dim=2) + selected_local = pair_error.argmin(dim=-1) + group_local = selected_local.repeat_interleave(2, dim=1) + selected_entry = best_entry.gather(2, group_local.unsqueeze(-1)).squeeze(-1) + + sign_bits = torch.arange(8, dtype=torch.int64, device=x.device) + sign_mask = (negative.to(torch.int64) << sign_bits).sum(dim=-1) + + packed = torch.empty((block_count, IQ2_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:66] = sign_mask.to(torch.uint8) + high = (selected_entry >> 8).reshape(block_count, _IQ2_S_SUBBLOCKS, 4) + packed[:, 66:74] = ( + high[:, :, 0] | (high[:, :, 1] << 2) | (high[:, :, 2] << 4) | (high[:, :, 3] << 6) + ).to(torch.uint8) + packed[:, 74:] = (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_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 IQ2_S blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 82]`` + and ``[weight.ndim]``. There is no CUDA encoder for this format yet, so + packing always runs the torch search. + """ + validate_weight(weight, "IQ2_S") + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, IQ2_S_BLOCK_SIZE) + grid = iq2_s_grid(weight.device) + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ2_S_BLOCK_SIZE, IQ2_S_BLOCK_BYTES) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_s( + 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 IQ2_S payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_S_BLOCK_BYTES, format_name="IQ2_S" + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_S_BLOCK_BYTES) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + high_shifts = torch.tensor([0, 2, 4, 6], dtype=torch.int64, device=blocks.device) + grid = iq2_s_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ2_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] + count = block_chunk.shape[0] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = block_chunk[:, 2:34].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) + sign_mask = block_chunk[:, 34:66].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) + qh = block_chunk[:, 66:74].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS) + scale_bytes = block_chunk[:, 74:].to(torch.int64) + + entries = low | (((qh.unsqueeze(-1) >> high_shifts) & 0x3) << 8) + local = torch.empty((count, 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + scales = d.unsqueeze(-1) * (0.5 + local.float()) * 0.25 + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + values = grid[entries] * signs + chunk_decoded = values.reshape(count, 16, 2, 8) * scales.unsqueeze(-1).unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_S_BLOCK_SIZE) + return decoded.reshape(shape) + + +def iq2_s_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """IQ2_S weight backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_s": + raise ValueError("The ggml IQ2_S backend requires num_bits='iq2_s'") + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq2_s", + block_chunk_size=block_chunk_size, + decode_chunk_size=decode_chunk_size, + quantize=quantize_iq2_s, + dequantize=dequantize_iq2_s, + ) diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py new file mode 100644 index 00000000000..106a1c2c9af --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -0,0 +1,275 @@ +# 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_XXS 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, mirroring :mod:`.iq2_xs`. Every 256 +logical values become one 66-byte block_iq2_xxs payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..65: eight 8-byte sub-block records, each holding four 8-bit grid + indices followed by a little-endian uint32 of four 7-bit sign indices + (bits 0..27) and one 4-bit local scale (bits 28..31) + +IQ2_XXS differs from IQ2_XS in three ways: the grid is 256 entries rather than +512 so an index needs no high bits, one local scale covers a whole 32-value +sub-block rather than 16 values, and the scale shares a word with the signs +instead of living in a trailing array. + +The canonical 256 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#L2489-L2514 +""" + +import torch + +from .codebooks import iq2_xxs_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_XXS_BLOCK_BYTES", + "IQ2_XXS_BLOCK_SIZE", + "IQ2_XXS_EFFECTIVE_BITS", + "dequantize_iq2_xxs", + "iq2_xxs_fake_quant", + "iq2_xxs_grid", + "quantize_iq2_xxs", +] + +IQ2_XXS_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ2_XXS_BLOCK_BYTES = 66 +IQ2_XXS_EFFECTIVE_BITS = IQ2_XXS_BLOCK_BYTES * 8 / IQ2_XXS_BLOCK_SIZE +_IQ2_XXS_GRID_ENTRIES = 256 +_IQ2_XXS_LOCAL_SCALES = 16 +_IQ2_XXS_GROUPS = 32 +_IQ2_XXS_GROUPS_PER_SUBBLOCK = 4 +_IQ2_XXS_SUBBLOCKS = _IQ2_XXS_GROUPS // _IQ2_XXS_GROUPS_PER_SUBBLOCK +# Largest representable magnitude: grid entry 43 at local scale 15 -> (0.5 + 15) * 0.25. +_IQ2_XXS_NATIVE_MAX = 43 * 31 / 8 +_IQ2_XXS_SCALE_ANCHOR_MIN = 0.65 +_IQ2_XXS_SCALE_ANCHOR_MAX = 0.92 +_IQ2_XXS_PEAK_TO_RMS_TAPER = 0.035 +# Bounds the encode search temporaries; see the note in .iq2_xs. The grid is half the size +# of IQ2_XS's, so the same chunk holds half the search tile. +_DEFAULT_BLOCK_CHUNK_SIZE = 512 +# The decode runs on every forward and is launch-bound, so it takes a much larger chunk. +_DEFAULT_DECODE_CHUNK_SIZE = 4096 + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +def iq2_xxs_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_XXS 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_xxs_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(_IQ2_XXS_GRID_ENTRIES, 8).to( + device=resolved_device + ) + return _GRID_CACHE[resolved_device] + + +def _predict_iq2_xxs_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)) + anchor_ratio = (1.0 - _IQ2_XXS_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( + _IQ2_XXS_SCALE_ANCHOR_MIN, _IQ2_XXS_SCALE_ANCHOR_MAX + ) + return ((amax / _IQ2_XXS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + + +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, _IQ2_XXS_GROUPS, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + # Only seven sign bits are stored; the eighth is their parity, so an odd sign pattern + # has to flip one element. Account for that cost while searching, not after. + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + d = _predict_iq2_xxs_scales(x) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + shape = (block_count, _IQ2_XXS_GROUPS, _IQ2_XXS_LOCAL_SCALES) + best_error = torch.full(shape, torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros(shape, 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. + for entry_start in range(0, _IQ2_XXS_GRID_ENTRIES, 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(_IQ2_XXS_LOCAL_SCALES): + 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] + ) + + # One local scale covers four groups here, against two for IQ2_XS. + subblock_error = best_error.reshape( + block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK, _IQ2_XXS_LOCAL_SCALES + ).sum(dim=2) + selected_local = subblock_error.argmin(dim=-1) + group_local = selected_local.repeat_interleave(_IQ2_XXS_GROUPS_PER_SUBBLOCK, dim=1) + selected_entry = best_entry.gather(2, group_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) & 0x7F + + signs = sign_mask.reshape(block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK) + aux = ( + signs[:, :, 0] + | (signs[:, :, 1] << 7) + | (signs[:, :, 2] << 14) + | (signs[:, :, 3] << 21) + | (selected_local << 28) + ) + + body = torch.empty((block_count, _IQ2_XXS_SUBBLOCKS, 8), dtype=torch.uint8, device=x.device) + body[:, :, 0:4] = ( + selected_entry.reshape(block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK) + ).to(torch.uint8) + for byte in range(4): + body[:, :, 4 + byte] = ((aux >> (8 * byte)) & 0xFF).to(torch.uint8) + + packed = torch.empty((block_count, IQ2_XXS_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:] = body.reshape(block_count, -1) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_iq2_xxs( + 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_XXS blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 66]`` + 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. There is no CUDA encoder for this format + yet, so packing always runs the torch search. + """ + validate_weight(weight, "IQ2_XXS") + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, IQ2_XXS_BLOCK_SIZE) + grid = iq2_xxs_grid(weight.device) + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_xxs( + 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 IQ2_XXS payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_XXS_BLOCK_BYTES, format_name="IQ2_XXS" + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_XXS_BLOCK_BYTES) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + sign_shifts = 7 * torch.arange( + _IQ2_XXS_GROUPS_PER_SUBBLOCK, dtype=torch.int64, device=blocks.device + ) + grid = iq2_xxs_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ2_XXS_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] + count = block_chunk.shape[0] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + body = block_chunk[:, 2:].reshape(count, _IQ2_XXS_SUBBLOCKS, 8).to(torch.int64) + entries = body[:, :, 0:4] + aux = body[:, :, 4] | (body[:, :, 5] << 8) | (body[:, :, 6] << 16) | (body[:, :, 7] << 24) + # Top nibble is the sub-block scale; the low 28 bits are four 7-bit sign indices. + scales = d.unsqueeze(-1) * (0.5 + ((aux >> 28) & 0xF).float()) * 0.25 + sign_index = (aux.unsqueeze(-1) >> sign_shifts) & 0x7F + folded = sign_index ^ (sign_index >> 4) + folded ^= folded >> 2 + folded ^= folded >> 1 + sign_mask = sign_index | ((folded & 1) << 7) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + values = grid[entries] * signs + chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XXS_BLOCK_SIZE) + return decoded.reshape(shape) + + +def iq2_xxs_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """IQ2_XXS weight backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_xxs": + raise ValueError("The ggml IQ2_XXS backend requires num_bits='iq2_xxs'") + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq2_xxs", + block_chunk_size=block_chunk_size, + decode_chunk_size=decode_chunk_size, + quantize=quantize_iq2_xxs, + dequantize=dequantize_iq2_xxs, + ) diff --git a/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py b/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py new file mode 100644 index 00000000000..d37d794602f --- /dev/null +++ b/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py @@ -0,0 +1,239 @@ +# 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. + +"""Conformance vectors captured from a real llama.cpp-quantized checkpoint. + +Each entry holds packed block bytes lifted verbatim from +unsloth/Qwen3.8-27B-GGUF (Qwen3.8-27B-UD-IQ1_S.gguf) together with the values +llama.cpp's own dequantize_row_* produces for them, computed from +ggml-quants.c revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. + +These pin our decoders against bytes we did not produce. A decoder that drifts +from the GGML layout -- a mis-set high bit, a swapped scale nibble, a sign +parity mistake -- fails here even though a round-trip test against our own +encoder would still pass. +""" + +import base64 +import zlib + +import numpy as np + +_VECTORS = { + "iq1_s": { + "source": "blk.0.attn_qkv.weight", + "block_bytes": 50, + "blocks": ( + "eNoBLAHT/sgYPIuVHR5unIa25z3nTMBrrDsowEZpEmhZYzdCh6ybWSOReiY2MD9xWx/BZF4J3Z/LTxnt0w2T+DYRwT0N" + "IIwJoKkC4mejy6P3sqbJWgv0FgfWRjO1zs0h+ZpZymO6MyQ9Ms9aF+DLPrjwY54mJRLDZo2emJfnx3B7rQMfrk8bWdoJ" + "Hc7n5ue6duVmzL9SXOzWw0SmcOsXFEu0ESDUFxuXCM0YGhrzGpYf6sJ0hTc5bGk4hFUG9NtWPKTaovj92oteeEPl9Fc9" + "8BYIR2j/8Hv38eeyZ9fzxt60f8/DEfda/DyI7XWEZ6CXKNxq2nbhUvHVwm4FckftMeVoGA2YAjcSngU6wBLrdmQNavPg" + "js7fDIkqa+Domf9cZoy23tsEc05IMcBUzWTS8th0Vso1lo8=" + ), + "expected": ( + "eNqFVzGIHVUUHTCRgGCaLQJugoFUQYLoFuK8T3YtxDKIBASRDRZiRLCwsRtCYrGsWBjEwuAvAiJJJTEE5wW+pAsJNmJC" + "QNlOkG3EJiiic+7OmT3v/jv/F4d77znn3vf+vPdgtzrwT64+/KwO4TXWiKpFPvVGnrePTyy+dj5XX32RDMhZL+N3DtYF" + "6EGuPeQYfc41xtb0Pl1zzBftTXldc+VWNmy+XlfXd5NFQjXWqsPPHkSds2gWc3iefa4dcPXH/d+PWnPW9CBqr86gx/cQ" + "9AD+d3Nf0d49H30zP4t9GpnveVvD9d1sceVWyZFnVIBTv+/3oKZr+W+v39p/Z+qei7DsnIZZK23z0jdZowJcc/hC8rxp" + "wo96+n6dbdz+mqlaPd0SWjNv/vslA+rzfOFl7me7Gaa/8nU7hvWTL95mPnvzqckiL/Qxj+e1Xn/niUl1+b2asTl3L7Nm" + "ztp7fO7ryD+HJ7dCNFt/5oifvfrMxHsiL3zqJQcvwLy7AzVh54F4+ALuTNbc0OlzUJ/Uxbye1xnao371RXvyPdEcj2Kf" + "fr3+vkZ3Uzn1Rve88DpfMYO/n/WhB7k60ekdmtMd3+daw0MMNTXxsFYtrHXNv2e52n6rBpqjj9OQ93yhCwcv/Vb3nNZz" + "fehhL7lDD7C32kP5bs/ZMOKln76C7/vIz3mOriXDw/s/DIjqz9/NA88Ijh7q5HSu5k5vLn6aq3y8bv76PiES5KmhZlTd" + "uB5+TjEjyPt56M2IHdda3eesmVvdeSNedc4Mda/he7xxoGZsfrucgIEXrYiEr5XX/ojDeu6eF3eYuYsew/0feyvCz72t" + "G7++3Py7sRfPf5wAcqwLKN/3RX7lOLuI1F441QKz979MAGsFtfXHdzJqRPWS9z2LfAO/9cEcplcO3gZYzx7+lCKf6gT7" + "tJ95OEf2HO5Pfgt1avwurKPvsGx+xyHWjECnZc1Zw+M18tqvmu/X9Wze2WmqPnm6RmQOdHcD9yOzVg/4KNd+nct5Omvg" + "5Pz0nPSsNIfHa9pLzp+5+vR+8T0U70jekmpzb1XeV+Qr3lrgtXpzra527+7FZaAP8ecr2XJGArVyWus64FCv3uz+PzlT" + "G078kYacvOoRxx5E5RfNVNy4lgzbl7KBuUafA9Mj9QBomvs68nDm9Ejb5dmwfSkZmCNCVx41AY45Z6hfdeXU68+G5+XP" + "NTpP9fBORPfBz9d7FJ0TzyriNKfH3we9J9GZ64zdu9hzKqDc5lo7RAI6o9foJx95yY316jrarzW9vh6by171oN7ZaAtc" + "/TZZXH2ULQeQs6bm+9hLD2ud4zXE5y92fw/9ng03VyZWK8D7HD72eH+EMR/4e1vdOzhWGzQHPvouGVTzfuiLIn0E5+6v" + "AT5ZnB5rLUcEyCmoax+i9lFnv/d43q9NnlzUp7nfH3X2Rr+NNe+S3o1liLy8o57Tu+x7VY/m+Xfg1/Wzovehb0i1/f1k" + "w85GPeSrj5LVBGpy6iev/eS19nOE/x+gaS2i" + ), + }, + "iq1_m": { + "source": "blk.1.attn_gate.weight", + "block_bytes": 56, + "blocks": ( + "eNoBUAGv/sgI5+PIe/BwwioFcBUzxw3iyphnqNoBMNQxtkYxKIg/RpRL53G82jgiUKrRvh5d/dzUmweXtBQV+zXsYEc4" + "BhC3bF7wWvQFzy2AI93zFgUGNb4wPvSYWoJISJPERQmGNC8t1ESzwwUl0VqLjpLMmxRBBIezwirprpBVD7/r/Ig7oFIl" + "9J2LtXeKAZqOhjjGIDxTHmLssDKbvAWf9um/y85ttT0H26bzHeDrAJbuisGPnSGlUT2c0uWVv/620cLVhOoO+ztPFia9" + "WPgjYsisuSabXH64WhO5WhRt5Z1kl3MXjvFr3vd3GQWsYMSRdfH1s7aUT2yDlAT+RMHIacwBCG0GIkBFwJylQDLB9wbg" + "zU483Lojdx+p3BogUaPTo7/iBfE6P39ha5fQr7yx4Uj2l/ngDgV3qsfCM5LfckI5HeHRM+6Q5AU8LF+SvtLEksaZEln9" + "qfk=" + ), + "expected": ( + "eNqVV01oXkUUfWBAFNIWoVKhFNFVXViRuDAzDz8QhWjpQlcWVIqCRiy4aFy0FF+kRBdStBR3FgtCNkJWEsnMK6mgDdZV" + "LbRQq9BFEd3Unyq6EOfOe2e+M/fNl+ricO49d+bO3Dv3fSTVxaOmued7L6iCzRid21IzF/H6E7LfRj7zqeSxYNECXPBt" + "xnJesOOZwYfGnPSQq89jYFNuU5170wuaT05ZMCNotpqadsGWmIctHGM6t9yd66Dz01rU27H7r0DNWY3PrEaM7L4WtmD9" + "lXtr0YQF0DgebdU7zt9z6c6p3mrp2yJG2x9ttbb+x8EaeopT//kNoCdMTRtwWisavUfk8Cb8VvrtWOvZImfi/Ly0B7OQ" + "2RePutiTwBpaD99BCx22zp/dU86gHlDNBvfjM0o2vkv9DcIf9KY70+pas5lHTPjG5bXqQm2rLYf8wBa+UPvANgL2jcuz" + "CXtCTRszLsOB8z7xgfM2YU+4g0DsjRkTseNoQvN8qAm8cNX3uhMErYXNWrP6Yl1995XZBK6pHqwFYg9A5yfI2QF8D9wt" + "6Vg77pNNLBr3T/rEvWMNfUBvANG6vnW9lD6j19zfD2+3/wfNkwfbTENufj/4wngr8PCOfvD+nIP3oVbtw9a1d3FHs+QH" + "+ad+Npth9PCxerN4aeZK85fmTc/heJ2ZtD/bs3DVZvv5W9D90d8MfNaxvrSWY3oN9v56zDGa5R+t0mSNCboXHr2/vWZf" + "xyfZWhNILj4v+G3hLtkefWbz5zu2+vyXWXDC4rM+cUCKQ8eeQt2FHiQdkLvGdU99ZgSj356uhdc/3tUqiO4EYU0rgC9o" + "2t9ts/+Mjwh29d5zRiB+iTWa6ZO2mtntYjzwwA5x9hnxTIlTDvaLe/o4cgaWtUYDeljr4Q+0jt3Ee/fn0FkeMbqv5zMH" + "vtqfcqo6s/MKPeVzOVe16yM3OnxHvX72oVYgPjQBfEa2rp8Hngs9I6zxPEXQWZwX+i3vs2NuDHPND9hcC79XcybyzRUf" + "bYHYor1wxN8Sbz9uI974YS36wtA4L+cHSnHh3m5eur+udv/jI9sPjNhF7m2sFz/aGwuuGd3ZMldbL1mxox9saMlm6N7d" + "XLER0GEjzmtEO7nsJqFZsm3mn76r1nqquwDEimvQF31PfVeOd/PQgX3Yug+dbwfviveTmO4no+u5j9hYMOE9akFmv7rT" + "xVzCV961KTd00QDWhQHEWRvnFDYRUhN8sQOneeI5028wIRZ1XS/5xfnr/TSruLeuYVyvSfOGOqQGgWh7501z3zcevD57" + "Ww1fbO3r9dWhR4wg2Bac8OUpn7TeTuspxlq2j3xA76lOHDfNW395ZsHo5fD/cW9vCnV3vm8PmxDqjQh2WOci9s670Jc2" + "7u1t9oWhiQ+N1saeMtD33nd6H/ZGm3u5Sf/026T+vfZFwvrOB1oBbObR6v6a16V9/Kacd3y+nO3QM90/njmeN+rBoGfw" + "e20wv3qWucc8z3qu+XyA+89vmO6EOaDa9JzoerP56ePo00CnuS/ixHEX5t4KlyDfhHwL+B7E5++kOrsc/h5+zESshN8W" + "truYC2wzBlYu+YwZslYgMWbY8HEeny/n4l6wEWMWfL21iOb0kXZSLIOuS9953AdT6I1JtaAHw16ZDFxLsJu/766rKz/5" + "iJklE/3AEZ0e/uZcchHBDvFWAH+QX/eR6+L6xjXaibXxnfW7jJHn51nQ81Ly8Y4MnIP7c2/1+/R9495lPUSfqIeZPb/Y" + "Ydv12Yh9cz6xjkFnINckhPvpN+Y7F3uSa+M34bcBzy/6iG3X18J9bATbXAPX0tv/ArwdXdM=" + ), + }, + "iq2_xxs": { + "source": "blk.0.attn_gate.weight", + "block_bytes": 66, + "blocks": ( + "eNoBjAFz/h4NhxlMgsOHqpTEeFmmMeZw4V0WDG0g4Cb0pTZTgvOM3r2oaw9MQEyjfPeuF/uVJwiZksqBaESGaYExuAFT" + "KMEBk0MNCEkDIWIGt8+8CRrWkLA2d7ZC/0m3vOOCDi0WUT6Mgv8VisTSWq625wLzKca0Lpy+i2qXT7x15YoTItH5Rzf0" + "hmMMII18mY4im+1dZIo77zVm6MUboI8D79T31BB1QJuOy5qNKSG6TWwbwyDmyNH/Kx6iY4bTDpD+o9HE4ACHp1TGvQwM" + "yGBr0/KEtaKcEJDJ4GDO/xjDFkyPmje7B5AlhYYbWP8QBfPTejik75531FEJNlH+mGkNxp5IHM3KzeC8Sr+o1pkM9A8t" + "sqJ02ZIZDS8LwnHA040JCXza+pedDYUCN94HQvK5kLGPRoUCkZtRJ4dVbQm0J+VwazuGl3caymqF/qftjjYMkfIiNBP2" + "4c7htG3cwS21qXiRkebCWsy4lJiFhE8AP6HXO76mrJpLsAYAURyv5J/4toWGKjGmD6mo9CS+VbBguEDGxLA=" + ), + "expected": ( + "eNp1VE2IVmUUvhhhCMlQQi7Shv60gppVi865NauiRSBRKUg0i4wkSBfiJnJRhn9BhgsJUobQfikGixx9z4XRLHSiEnUR" + "ljpEmJnJV0EOktY5r+fceeb1tng45zznec77c9/vq1btT72dM5rhs1u5WrWfAEk5sRiwukNH3ifXh5fBx9bTdWrr23qe" + "p+q5OTJ06hfWGBBD3/JHauUtZ++LI9fab4AT8Oe+e6Xw5Wizo+4/cYFH7nqyrt6Zz5Z7FM8F+SKXgOnDF56Azm6wDm+s" + "WT37WzL0n95tOVk+sGB+45x4n7yfe6gFT42ewisGn0sWY41q7VuEGDy5goce/ZY1T5qL1cFbrT0DAy8WjY8ckBzk8xjy" + "K1p/O/Ym4g3Z27A341wqQKGD90bxLuHdte8w3h28y6k3+M/HZBge/1UsTpxbx86lyLXHmls/WQ5cqw1/zHJdArQzwZvi" + "rHAm7jhXe444C/ymmvJuyt+ln7eG33R7Xwf3D9TVpmvEIzsEwAMbvkGOgcuIGYWmjYXf6lY3vGM1V8e3UaB38zExziHK" + "JeXYALoEINeFV0CTOZyNPpvZt6wn1eldyTDyxSZWCObeJ+yhx3vTZmg/ELOiT+DhXM+7jSfvXdosOnFZDFoHGGJAXIcc" + "uy/6XHgYZrVrROy9cJ6rrXNFkaPWbe6Rix5qc6/w5PzMrMeacnaBrDl8U8OGatafSUEe23x0+62N5R5zzz1inHvJNbX3" + "pZhD4Qu+9eo3sG9VfCf8Rsm/ef6e/n2lQys+S/w7l+9CYh3smT7eifcJ30nk+IZwLnpiv7hXXL/rPL0vv+LqxcsJQM4R" + "cmf+uLsBPve0FqydY+NN7746OOfr/9G280LbsQ7uRaIOlByuizrcQ3VpphgWvaT/L5dmtphM99fBad5oLhFdk33God/6" + "nk/jUOvI/MC/pzKqXe+SIjksJ+XZ6j56hU2jUbAPCF8Kj8fcD5/NQU3uP/VDUtDBp29sBpZ+JJY7ktZsCM40wbsvddTk" + "c3Ktnto5Rn9gZe8NqR4cTYgubtGWC6I8RwxOIwFaL+g5+uDLvGk39+3h6qdzBEjfL7ildj45SGuBOsM0qm0Kvp1lc2Ce" + "+Ixpa1Xjh9LhndvEooIcaXTLzFp5Dt5r01FE5ZrS11FnXbHG1LqTH5AieaSxQ2uNS4PLfhbN2aLxGtl6xgVcbzFFz/Sh" + "dV+5RrtW7p25SBNbFtYWA4PXHjA+aeSIpvE6dAmjaxqvW84hCOubNvdWvCmBauHepJE1EiCN3XHRuGS15mIwXWjD634J" + "H4AKHbU6P4PvnTqQ/A7E7wnPJ8W9Md5fcZ8E3nZmNXqWxp7ZLcOfzqstap0QyrFrciz75jM+YngMyjWGLl/rhX3GHcA5" + "GL75tHOH1u4E7mXavcG7Sh1v68qdLBmn/r/0d7RkPI3N2S6eZ15rNs4icKapgUugZa8DhDBveGydXK8/TkNHb2gG937I" + "mifNa+DEeIvWC1gNHgpf1K4n5yQ0HXO5+vwTGlizXjy2teWbD5zl6GGuSA4qo/lcF3zuqV86dQ9PkiJZHPl7HxuQ85g6" + "em3/uu/urD2/aobnMUPAdwVT+yv3nfw+cq/Y/zRYz/td90Loi3O389bNZoX0dp3MsDy4APDRy1C+BXgYgLOy7vnXHmpA" + "z3BW2xPjt8Hv7e9B8Eyuy3zcVdnDd9Xhl2rN0aQgjCsf2MkQy555KKJx78+dXaMeexatD7NaPuvOb0iHF++xSBOvLs61" + "R7LcIwWvWtbIxkGeZ4TfezGLY7b7xbWZrx5/m/oHV4nGNDD7iETuoJLXnI0HrtVYzzjvcfRgJsHslDWvH0kOMgy9vEO6" + "ctQZb+hbfX0dfOHj8DifQh/eds0fR5KCAhP3bGSD8cPv/R5RrOcxuU4MrsE5qdCmmAfeqfmf7aPhJ7aKxYDW3Lt9RuN1" + "wujayBn61NGPucnm+dw6fHmNqX3lc/reyPcd5+bYN5w3zkHF+RKcW2JGcUdT/unnSHgm37MU50xwNwR3kwofu4+K+7J+" + "3d7PxhkyfOHrDMsVXMTIuSO/CjGnmNd6Vy6/r8Y63ha8PYk3E3dXvp/y+0CeineG/niTOK/zffn9xLegjiigQw+hrvwu" + "AX+DecZ/P39gQw==" + ), + }, + "iq2_xs": { + "source": "blk.2.ssm_out.weight", + "block_bytes": 74, + "blocks": ( + "eNoBvAFD/ooQA6AAiCN6lmuU9VeqmJ0A5LvON22xpJNoA/QAWAGcDZaWQKyDoBePECMoCX4R2nrRr2MDSLwulHU1HAI0" + "CJIA7qhpU211f1h7XxAAzCTiWUXSqgJfCmAn5Z2fsULmOlXie4P+2Hs6K4JXIkNMR9BMRi0ZH7TGIaAzha65Ev8svrLU" + "1ifnCjik2AR9VzVCZmaZ9mj0FNMeyOaKjppuTNaS4wbCEv+Gcmfqfj5EnRP95Ooudk/83EyYhB7dAMgAhGrkfgw9mb+0" + "iVCEeyeZATYADOqTF9s0MyQ0hTxVjwANNM9T0JRdV5xbUJAyTRtpVWhT61gGwqYlAPxNcKRXB3CIfKqRCfyjf2kPFAEX" + "LmoX/nqif6ITk2ALpVceb5N4Tbpntlv4epaLABIJsJWJzbQAtif3Y3IAVhfoXfpC8GUkXYI35fkDxVr/SmoyAAJJMjgb" + "H0hKjGpgJVmrJd94HzoDLeSgUQN6OSomyKRERK+Iw6ljEmOmdWwqnHp6sDpbIynDWhCVbdJGRFvGr1757kJ49dKyb1ZJ" + "vCaNS9cwmkYajOIXLCaNWoLNhkaeOkv/ABP8eMYiEiIh58/YfcDzwlE=" + ), + "expected": ( + "eNp9Vl+IncUVHymxoUVYIZYktXpRCLQNuiw0LX5n9BZDjSS1C81Dgwveh6ASKElpK5GK3hppI0RYTRM01nYfUtiaNmyb" + "kj/OmeU++CeBPIS2lH0odNOq9UFwzUsjVew5M+fMd+bbmz78+J3zO7+Zb+bMuey6nRthefAd74gF2GH4v/nxu5EAvZV1" + "kWNl1oRTPPmFYRyj+eXR9d49sRmVLebefCRKDMrki8YL/YfPglv7fmBQjMSNIAyO3OI1trpB4yafDMMj30JmQkMxcN7R" + "sH/+VdRYask7+uwvwZ19Kwg3GhPQ6pyr1+jB3fxjuHTyeGQm4BhApwaW5Y71ncydpSc2rupzd17vFd07CKNoBVrjNW7z" + "Why9/lckBolBYuj9d0vUmGF8ID50pw8BY/D3B6NlA/bg7LoYNRbkOu8he9tzGNTfkzqfjf00+5HnuH/7h2WuKU6zrTX5" + "feDY30L7nmjf2bw7VO/dIvXeXZwGAiom7j+kWtIHt97uCbHjK3W+A93Fd+7X7YHtRZ3f9hkGClblC1+bjMy9dy6CMFb+" + "naOGEOa+eTTFwikXoOioXuZ937gumjxIPVgPa+Tz4vfWL3Hjzv+IEYRT3Id7UBiYB48dR4X1cT56kfqwdEdDDMwSo8mD" + "ahIHrvWv7Mn1mVfC6B/0W5x5pZn74V9AcmRmzXCjPvEi54PDfwC3418NIQhKLLXCEuPEmfVefA3FUWP2dfazrB40nqJZ" + "n+RprX7LnMH6LdRTrWeN97DfMD6Yfp3edeMmJAAz5VFjZqmDQdFXph6KssaPWQ9clxztN2R91s29tZedPnb7U/fppz9n" + "gAC/d9PpOHvuu54QVRPdG2/ypfxX6xQgwOUt76Hy9J/u86wTx3Fe8478VtHMDJq74Jg3y+efnsLhA5+PDI4JwBi+vNdL" + "jJXO3pf3Fv9w5q7oFj6H7tQSEkNBm6OB6uy/Vq316L7ZC9396Nu+7HUNjzlbe1ar5/t379rVQDV794QLx9D13gyM4ZfW" + "RM6JPeWNaN7WWFeNcjB7p752zyDf6p6j9dp+5l61/RTuvFHtWdnFwOG9X4kcE3vO+99/IbFoqSYaWM09Os9Axmj7P2P/" + "+QNec4qj1NVT+RPm/wMEZQu8Vry8e9Ni0cwcpXmgXOYCVs0jQT2lD3IXc+8Ul/u19y39kH5pL9Awai+5V5090PTTq6+3" + "84FFd2hDJHhhjS2upXvuue0r5dx/YF1qqumb1O8zQX8T3Z+DoKEcipa5SVh/gy+5XbNwCpx7lxkTu3ebwgc3epMH8UFZ" + "w7UB/R11R5VD4sEeTBpj9kIbZz1U/t52+k1tp/jpIGgSWOMaxwsnUfRQ9LberhlXXzgJsh+O0bt3albd1/Yj9ySa2I/x" + "12tbT1N6nDmvW/4qxTPMkGJmzt1M4y4dhoRcUz2UOvOlZ8Dte5v2eDVUYD3HTcK+t7OPOWshx19koNs2HRPXgCqfvJKh" + "tW3TPs3Q+htiZ/7AzGFTgfUqTvOQMXsByny0CEXjWcrzBWXOdEba+WnnqJ2PUGI7M2kGFilepP78IseZm4TpT6HkzOqp" + "fGYm7O+i/k2N+/3IrJk30pjf5dIzKADzrk31dlx3OyBh4ll6ix0ogCoe3OoJsXi1ntYc5Z6j4RzPXoirtLHxUhA0JR7N" + "ozBI3JR678ZYa0+gAMawBf39/K03fvHIt3lf/bZ+n3U+g/2+PZ/V1ZfP113Dui/75Xq7nvdr61jd3+6t69mbEaveaV/q" + "M4RV32vvRv7XGCDI+cpHvsrbehdY+rq8te1tjmHV2yxvjammdc55ttp5A8nbudR84tm4SrO+ei5rL89p1tDMOeSZnqc+" + "XKaeXwYBJi3r4PoHMljvH8hw8yHV2DPYU895m4P8DnxHx1JjrTdF9d1ggEmbfCmWnHnypdrHnt5ULPPNvbT9zm9Sa7n3" + "0f4mRkfuQ7f/qUBoGHPHT7AGVuOYfaqzx9Qa8VqPao3sX3nM+iD76ffUA919ZC/Q/dQzHH4QCI3btqGANeIg3Iyufhuk" + "lrzGX87YPaf0Qr8XzHkae2a9i/StxJ3zBnOP8s3ky+fVswU6K2qsOZ9fzyxx0Dt13iqYnqK9l/EFe8bphZuiO3YPEntm" + "AuwPQ9ZAwbn6Ol50Xz8HBGa8+sdZ/97lq141yqPWjK/UUv2TN5AAH69Zszj/+2OR4DlXXZlrJm49J36N55/7m5/8yeOe" + "YuCcQXmUvGjsk7jVN9H/5pt2MePCb77sD77zvFeN8shsNGAPQ9aA+8HvUAAWS2f+HU0Nt539meZAsbc1WYOyzosXDPvu" + "/qrzW/DbKMq7yFtqTG/oGeZ9U/1/TYe49w==" + ), + }, + "iq2_s": { + "source": "blk.35.attn_output.weight", + "block_bytes": 82, + "blocks": ( + "eNoB7AET/q4M+OCIlAAzLJjmBQAJuLQHGSwziwAmnlIDogF04URRrgOutsGuCoscGISHWf5Z7PnTZCXDydppD389LinF" + "F+K8HgxMCGQcABGwhFvXRv1Y11zFCvZOQw2g3hUdMb0vDvZmARjtM7DOQZ5kTxxzz/H+AAeVcDR0leZIoJDxcfhrr7Uv" + "kpvnsBWJYWuM7FOcANKeHe00igDNHJJDAYl2po2NiMp/XwiRcq1apgHtYw6m6g/Aao/ThGi7zgieTN49ZMg3FACEA/WD" + "VoudZA4dZAYThwSiFspPvOiNjuUxu11ujtJevf5VIolkIiygEQJrv56qq32nzw0KowLwPnr0CswZwQKXCL5JdIKYjwKX" + "nkrYcesAH00mlyNWkgcISZ9synXVqLyJ7IuNpgFD4XvoXfHtkou3P9FjihCFiMAeQwlxiWa7act/5XglDa3SmBz14UIO" + "vwsbhiMZ3oIAM4AAgmkAAXRzAAlEAHWHs6+UevSuXXqwriyP2frZUU5JYkM/84aPBbjS5XTjweCoMK6QHAYBoCQ0M2bP" + "hbc+UA3t0EhwrUstChnOAAkLEQRd7in2AFXGGaCDeNvvEgC4IkeqzF+vBg0/wJbAXSu4r42iDMm+Oqdr4+mcwHeTesEG" + "wEgAQgMAVoImQ/hldEZUNmfm3b4=" + ), + "expected": ( + "eNp1V22IXtURPjbWajDtVhJZxKQvRUStP9ZoNGZm263YmlIDi0K0fq62frIFNVASpGQtrYa0DYuCFcMmK5WgIrqiqNmZ" + "VxfxYzERV1EjiLBGhFhFFiWQH8F05tyZc+fevPnx8DzznDnnnjNzzsvuzMQinJlYxEM33cxp41kkGiMLwDh7mqta8jHP" + "OX0nCUAx+c6HOLN9Y4b5ZDErNLYctjmULt0HgjZH0LHGpx7r66YD68BAQ592WBG8Nmj+qq15XBiz96d/wdB116Iwjexi" + "5YbnkJjVM99zOL1wmARgoBA3/IXZs7thvMbVG2jm9F+iMJjmzub/akzmZZjPLa09YesLmubQr9JX6yO0+xj3ODn6BiqO" + "cY72OUHONJhO+SkrOk/vR2F0Hli4pOtjGhvautwDvzONu1EDgqf3rL531ndlA0iMPXpOPVD1sOpno9d+F4wpAAL3rIv1" + "e7BVq67VNueKru7O4S8zhlct7wqDs2N+74toOdFjnxfOG++/atRz6z0PXuN95Hk7x0DxxIdPZ33XobUsTMYae456KHms" + "Y8U/51YHrX78UVR2aDw3dz77uGgfL3Om/jaK4/edxOkHAyQA0ahaPfMh+paPNo4DT67AdOrvYfaeLUULKLL6DsljG1Of" + "7VsY4Hsh34d+X79r+2DRZb/p7G3Q/6NXUbmvcwMrRFOEeeAcxiC99SJFbNyzDU2DaG5psJysb99+kKXGHGvudfUaa81D" + "D0ofSm/qbzfWDvugVk7Zs34/fAO9v7qu34XY/8iWC7p/1743j21uvC/gc1rrUets3Dqzf4ePWr/qh/bQ+5b7pLFz6VXd" + "Yyze4JPwxEV/R2EKiF7JOXTH5+hxyfnkPRDQGUdOHlTd/538rpsXdQvgvt5nu9ts993vPrTfgubYOym5facu5/7+B+T3" + "+HeqM5tmhccCUGiO5UH2L/98zchDz5EBlIMHoqdV973Uj6ptTP01Ob7vARLA+O4r2TSJxquXye+eaR13bXHJK/uo942+" + "zzAG8Sx23irn/r/CVN9vlckYVBvAYxljjzXvxPVPVbmnnUmK4dePQ4XqhZU3oscCNk8ZdEzZcwf+/BMcf+4vmB6+iASg" + "UM9hMStbXol9jsVk+RTX8LXjWJzjfQ7AcA9Kv2P9TIPUkf0bCl0zfI8dYa+NM2o8d9tqfPnhHSzMafEoCSCAglfGNL+M" + "/XwRTa0aQGHofLFXQQqPdVzhWnLZcsDyslbfdBxTH33NsC7ZuqWWXk/vU6+6e40a9ajvdr7/cu853nd/H37nA7K38Mo5" + "KOA20rZryHH7vgfRNBzF39wCC4+dzwrVAjpwr/x2C9sYhjFwv+jpdTS8W/6mED60Wv6mnpa/LQziN+JeOQO3XsmduYOg" + "SFcsJ4Fy9pXFV480Nq/k6Jh4GPJ9jNp59g2y2D1Kv/inApRnjyO0uOErD2z4A3ssGjU2qMaY22Muytpsa1BcQ2ur9dY6" + "t2p7VK2tTyj5bGNkNeB2Pdp1aOXVNa736XtunNF1GINwJky/+Q8ZYO7aUVSofvnSPWweuxfgc+jjrQ9yeohIAMYZ6m9Z" + "chjNhzAO4nPJvX+VAqbmlrCi7+sNrLH7yuKjaffJfE4rfkgCmP/+XWUanl+JAu4Raw64ZzGlTc86IDCEOOORt79iAR6V" + "L2/C3oW/Idb3FN4Oxbdkb8396gz1WSGcWWuBdvZSI+Pav+DijMmpPlSonv9U/ge74GIIup3Hzmn2nTWG6ZFLFqOyxkOv" + "T4KAlNXzWHLYcrKn68g3qrXkm8bkHPfhOb4P5ZlTPqN0ch90xn/NpnOsLB6q9pwwXvI0R3ywOXme5qkX5/lcyyWbV9bw" + "ea5tHoR5cW/12Gv/oM6RYRQG0+za4Bqm+HkOfoH4aGs05thaHkNcS8fynHt/puCA6OExvKzXrls/mCYSDh+cZWNUFvCh" + "98/tujZgC5V33c3UC/P7BlBhmk2D6rllj2jMOfd/z4CADNAC3bX4Kw7j7RxK6x9VQOCM8T++i5E9Z+TMu1lR5uz6GgTO" + "UVMAbDmyotvKqXDeMlR8vGltNzAbsKWx5XG8536v4n0Md63XHafU2RoB6f2r0Jgzb9+NRXuOcpVH6c0fYzp+F6SxD6aF" + "q3jsgzU5dqjvqOIqX3NnllLG6IEKM0sh7b+MM1c+mO+6go4rbx7jNPQrEFC6U+pxp9Sniiuop2OOKp9K3uQLNV76FkNM" + "Gcv/jUXXPhR/6gw8BjgAxy6/vhvjwjN3kAHSkglOOy7EElcait5xITfylbfv5ozYw2afuPSq2dPK89pVtcJSs81jWGqr" + "2uvYru/IThLI/8YfYdYnbMIcn7CJy1jFbVR++z7FfdZnqc+me6/OC7Z/KvuPd6E+E9ve6z37HajvAzTqUN8lbtSk+S3M" + "OvbHtfdG+xV7tmQi9pvynOqsVM5bnw16vrt2PfxdVO+lehP7L6vfVdQxx/j/CFEl4w==" + ), + }, +} + + +def packed_blocks(name: str) -> np.ndarray: + """Packed bytes for ``name`` as ``(n_blocks, block_bytes)`` uint8.""" + entry = _VECTORS[name] + raw = zlib.decompress(base64.b64decode(entry["blocks"])) + return np.frombuffer(raw, dtype=np.uint8).reshape(-1, entry["block_bytes"]).copy() + + +def expected_values(name: str) -> np.ndarray: + """llama.cpp's dequantized output for ``packed_blocks(name)``, as ``(n, 256)``.""" + raw = zlib.decompress(base64.b64decode(_VECTORS[name]["expected"])) + return np.frombuffer(raw, dtype=np.float32).reshape(-1, 256).copy() + + +def formats() -> list[str]: + """Every format with captured vectors.""" + return sorted(_VECTORS) diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 3cfc522d915..a9f93a7aee5 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -82,7 +82,8 @@ def fake_quant(inputs, _quantizer, **kwargs): received.update(kwargs) return inputs - monkeypatch.setattr(backend_module, "iq1_s_fake_quant", fake_quant) + # The dispatcher resolves through its registry, so that is the seam to patch. + monkeypatch.setitem(backend_module._FAKE_QUANTS, "iq1_s", fake_quant) inputs = torch.ones(1, 256) quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args=extra_args) diff --git a/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py b/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py new file mode 100644 index 00000000000..eba6ba655e7 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py @@ -0,0 +1,149 @@ +# 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. + +"""Conformance of the GGML IQ decoders against real llama.cpp-produced blocks.""" + +import numpy as np +import pytest +import torch +from _test_utils.torch.quantization.iq_llama_cpp_vectors import ( + expected_values, + formats, + packed_blocks, +) + +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import ( + IQ1_M_BLOCK_BYTES, + IQ1_S_BLOCK_BYTES, + IQ2_S_BLOCK_BYTES, + IQ2_XS_BLOCK_BYTES, + IQ2_XXS_BLOCK_BYTES, + dequantize_iq1_m, + dequantize_iq1_s, + dequantize_iq2_s, + dequantize_iq2_xs, + dequantize_iq2_xxs, + quantize_iq1_m, + quantize_iq2_s, + quantize_iq2_xxs, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +DECODERS = { + "iq1_s": (dequantize_iq1_s, IQ1_S_BLOCK_BYTES), + "iq1_m": (dequantize_iq1_m, IQ1_M_BLOCK_BYTES), + "iq2_xxs": (dequantize_iq2_xxs, IQ2_XXS_BLOCK_BYTES), + "iq2_xs": (dequantize_iq2_xs, IQ2_XS_BLOCK_BYTES), + "iq2_s": (dequantize_iq2_s, IQ2_S_BLOCK_BYTES), +} +NEW_FORMATS = { + "iq1_m": (quantize_iq1_m, dequantize_iq1_m, IQ1_M_BLOCK_BYTES), + "iq2_xxs": (quantize_iq2_xxs, dequantize_iq2_xxs, IQ2_XXS_BLOCK_BYTES), + "iq2_s": (quantize_iq2_s, dequantize_iq2_s, IQ2_S_BLOCK_BYTES), +} + + +@pytest.mark.parametrize("name", formats()) +def test_decoder_matches_llama_cpp_on_captured_blocks(name): + """Decode bytes we did not produce and match llama.cpp's own output exactly. + + A round-trip against our own encoder cannot catch a layout error that the encoder + makes symmetrically; these blocks come from a real checkpoint, so they can. + """ + decode, block_bytes = DECODERS[name] + blocks = packed_blocks(name) + assert blocks.shape[1] == block_bytes + + count = blocks.shape[0] + decoded = decode( + torch.from_numpy(blocks).reshape(count, 1, block_bytes), + torch.tensor([count, 256]), + dtype=torch.float32, + ).reshape(count, 256) + + assert np.array_equal(decoded.numpy(), expected_values(name)) + + +@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) +def test_round_trip_is_close_and_shaped(name): + quantize, dequantize, block_bytes = NEW_FORMATS[name] + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize(weight) + reconstructed = dequantize(packed, shape) + chunked = dequantize(packed, shape, block_chunk_size=1) + + assert packed.shape == (2, 2, block_bytes) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + # Chunking the decode is a memory bound, not a numerical choice. + assert torch.equal(reconstructed, chunked) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) +def test_zero_block_has_canonical_zero_encoding(name): + quantize, dequantize, block_bytes = NEW_FORMATS[name] + packed, shape = quantize(torch.zeros(1, 256)) + assert torch.equal(packed, torch.zeros_like(packed)) + assert torch.equal(dequantize(packed, shape), torch.zeros(1, 256, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) +def test_requires_complete_last_dimension_blocks(name): + quantize, _, _ = NEW_FORMATS[name] + with pytest.raises(ValueError, match="last weight dimension"): + quantize(torch.ones(2, 257)) + + +@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) +def test_nonfinite_values_are_treated_as_zero(name): + quantize, _, _ = NEW_FORMATS[name] + weight = torch.zeros(1, 256) + weight[0, 0] = float("nan") + weight[0, 1] = float("inf") + packed, _ = quantize(weight) + assert torch.equal(packed, torch.zeros_like(packed)) + + +@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) +def test_fake_quant_has_pass_through_gradient(name): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + weight = torch.randn(2, 256, requires_grad=True) + quantizer(weight).sum().backward() + assert torch.equal(weight.grad, torch.ones_like(weight)) + + +def test_error_decreases_with_bit_width(): + """More bits must buy less error, or a format's scale handling is wrong.""" + generator = torch.Generator().manual_seed(7) + weight = torch.randn((4, 1024), generator=generator) + errors = {} + for name in ("iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s"): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + reconstructed = quantizer(weight) + errors[name] = float((reconstructed - weight).square().mean()) + + ordered = [errors[n] for n in ("iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s")] + assert ordered == sorted(ordered, reverse=True), errors From 21eb1af07113e3880279ec51e4c2866ed330ddea Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 19:04:15 +0000 Subject: [PATCH 2/6] Note the new IQ formats in the changelog Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0a6afcb0017..67744b6c3ea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,7 @@ Changelog *Quantization* - 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 ``iq1_m``, ``iq2_xxs`` and ``iq2_s`` weight-only quantization, completing the GGML IQ formats at one and two bits. The same 256-value block constraint applies; these three pack with the PyTorch encoder rather than a CUDA kernel, so quantizing a large model with them is slower than with ``iq1_s`` or ``iq2_xs``. - 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. From 0243784729107eeb4eff0a146474ef76008bc9cb Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 20:17:58 +0000 Subject: [PATCH 3/6] Use the IQ1_M scale anchor rather than IQ1_S's IQ1_M inherited IQ1_S's flat 0.61 anchor when it was added. The reference predictor these encoders follow anchors IQ1_M differently: the ratio rises with a block's peak-to-RMS instead of being flat, and is allowed closer to full range, clamp(0.58 + 0.035 * peak_to_rms, 0.65, 0.95). Measured on 15 Qwen3.8-27B MLP weights, relative reconstruction MSE falls from 0.17372 to 0.17291, a consistent -0.47% on every tensor. Small, but the flat value was a guess and this one is not. The anchor is an encoder choice, so it changes quality without touching the layout: the llama.cpp conformance vectors decode identically either way, which is why they did not catch this. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/iq1_m.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index c76de4d4618..6bdcd31af74 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -64,13 +64,31 @@ _IQ1_M_DELTA = 0.125 # Largest representable magnitude: (1 + delta) at local scale 7 -> 15 * 1.125. _IQ1_M_NATIVE_MAX = 16.875 -_IQ1_M_SCALE_ANCHOR = 0.61 +# IQ1_M anchors differently from IQ1_S: the ratio rises with a block's peak-to-RMS +# instead of being flat, and it is allowed closer to full range. Values follow the +# reference predictor these encoders are derived from. +_IQ1_M_SCALE_ANCHOR_BASE = 0.58 +_IQ1_M_SCALE_ANCHOR_MIN = 0.65 +_IQ1_M_SCALE_ANCHOR_MAX = 0.95 +_IQ1_M_PEAK_TO_RMS_TAPER = 0.035 _IQ1_M_GROUPS = 32 _IQ1_M_SUBBLOCKS = 8 _DEFAULT_BLOCK_CHUNK_SIZE = 1024 _DEFAULT_DECODE_CHUNK_SIZE = 4096 +def _predict_iq1_m_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)) + anchor_ratio = (_IQ1_M_SCALE_ANCHOR_BASE + _IQ1_M_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( + _IQ1_M_SCALE_ANCHOR_MIN, _IQ1_M_SCALE_ANCHOR_MAX + ) + return ((amax / _IQ1_M_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + + 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) @@ -79,8 +97,7 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: xnorm = vectors.square().sum(dim=-1) xsum = vectors.sum(dim=-1) - amax = x.abs().amax(dim=1) - d = ((amax / _IQ1_M_NATIVE_MAX) * _IQ1_M_SCALE_ANCHOR).clamp(max=65504.0).to(torch.float16) + d = _predict_iq1_m_scales(x) d_float = d.float() # Choice index is shift * 8 + local, as in IQ1_S, but here the shift is free per group. From 178efd3444dbc82eac5a66f358d2529c68a09344 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 21:08:29 +0000 Subject: [PATCH 4/6] Add CUDA encoders for IQ1_M, IQ2_XXS and IQ2_S The PyTorch search these formats shipped with is correct but not usable at scale. Extrapolated to a 27B model it costs 81 minutes for IQ1_M, 42 for IQ2_XXS and about 9.8 hours for IQ2_S, against roughly a minute for the two formats that already had kernels -- IQ2_S is the worst because its search is the widest, 1024 grid entries against IQ2_XS's 512. Each kernel follows the existing per-block structure: one CUDA block encodes one GGML block, the codebook search is tiled across threads, and the tie-break key orders by error then codebook index so it matches the reference encoder. The format-specific parts are where they differ from IQ2_XS: * IQ2_XXS reuses the even-parity sign rule but packs a 4-bit sub-block scale into the same word as four 7-bit sign indices. * IQ2_S stores all eight signs, so it needs no parity correction and compares magnitudes directly. Its 1024-entry codebook exceeds the static shared memory limit, so the grid and its norms are dynamic shared memory. * IQ1_M has no leading scale field -- the FP16 scale is spread across the top nibble of four scale words -- and chooses its delta shift per group rather than per sub-block, so the shift sits above the entry index in the sort key to keep the reference tie-break. Its 2048-entry grid does not fit in shared memory at all, so it reads from global as IQ1_S does. Measured on a 5632x2048 weight: iq1_m 10.7 -> 309.8 M elem/s (29x) iq2_xxs 10.7 -> 1047.9 M elem/s (98x) iq2_s 0.8 -> 725.7 M elem/s (907x) CUDA and PyTorch outputs agree byte for byte on the large majority of blocks and diverge only where two local scales are within a float32 ULP of each other: over 12288 blocks, IQ2_XXS and IQ2_S at 0.000% and 0.008%, against 0.016% for the already-shipped IQ2_XS. Both encodings are equally good and decode identically, so this matches existing behaviour rather than adding to it. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../kernels/quantization/ggml/common.cuh | 2 + .../torch/kernels/quantization/ggml/ggml.cpp | 74 ++++++ .../torch/kernels/quantization/ggml/iq1_m.cu | 240 ++++++++++++++++++ .../torch/kernels/quantization/ggml/iq2_s.cu | 213 ++++++++++++++++ .../kernels/quantization/ggml/iq2_xxs.cu | 240 ++++++++++++++++++ modelopt/torch/quantization/extensions.py | 3 + modelopt/torch/quantization/ggml/iq1_m.py | 17 +- modelopt/torch/quantization/ggml/iq2_s.py | 17 +- modelopt/torch/quantization/ggml/iq2_xxs.py | 23 +- 9 files changed, 817 insertions(+), 12 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_m.cu create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_s.cu create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 30818f0da3b..9b36b1ed5bc 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -49,6 +49,8 @@ constexpr int kScaleBytes = 2; // them cannot drift apart. constexpr int kIq1sEntries = 2048; constexpr int kIq2xsEntries = 512; +constexpr int kIq2xxsEntries = 256; +constexpr int kIq2sEntries = 1024; // One CUDA block encodes one GGML block. The reductions below fold over exactly this many warps, // and each kernel static_asserts that its codebook divides evenly among the threads. diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index 4af30719f2f..99087cce8c9 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -23,6 +23,9 @@ 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); namespace { @@ -53,6 +56,53 @@ at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); } +at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_cuda(), "IQ2_XXS packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ2_XXS packing requires a CUDA grid"); + TORCH_CHECK(scales.is_cuda(), "IQ2_XXS packing requires CUDA scales"); + modelopt::ggml::check_pack_inputs("IQ2_XXS", input, grid, modelopt::ggml::kIq2xxsEntries); + 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]"); + // Same rule as IQ2_XS: these bits become the block scale verbatim, and a negative or non-finite + // one packs cleanly while decoding to garbage. + 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 iq2_xxs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); +} + +at::Tensor iq2_s_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_cuda(), "IQ2_S packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ2_S packing requires a CUDA grid"); + TORCH_CHECK(scales.is_cuda(), "IQ2_S packing requires CUDA scales"); + modelopt::ggml::check_pack_inputs("IQ2_S", input, grid, modelopt::ggml::kIq2sEntries); + 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 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()); +} + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { @@ -69,4 +119,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 74] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); + module.def("iq2_xxs_pack", &iq2_xxs_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 [256, 8] holding " + "non-negative codebook magnitudes, and scales must be finite non-negative float16 " + "[numel / 256]. " + "Returns uint8 [numel / 256, 66] on the input device. Non-finite input elements are " + "treated as zero during packing, and finite elements outside the float32 range " + "saturate."); + module.def("iq2_s_pack", &iq2_s_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 [1024, 8] holding " + "non-negative codebook magnitudes, and scales must be finite non-negative float16 " + "[numel / 256]. " + "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."); } 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..c8db5776916 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu @@ -0,0 +1,240 @@ +/* + * 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 +constexpr int kEntries = kIq1sEntries; // IQ1_M shares the IQ1_S ternary grid +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"); + +// Squared error of approximating x by scale * (q + delta). The grid is ternary and the format +// offsets it by +/- 1/8, so the dot and norm are shifted the way IQ1_S shifts them. +__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 + kVectorSize * delta * delta; + return clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale); +} + +// Accumulates x . q, |q|^2 and sum(q) for one codebook vector read straight from global memory. +// The 2048-entry IQ1 grid is 64 KiB, so it does not fit in shared memory; IQ1_S reads it the +// same way and relies on the cache. +__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]; + } +} + +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 = 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 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 = 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]; + } + 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/iq2_s.cu b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu new file mode 100644 index 00000000000..be718d16b7d --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu @@ -0,0 +1,213 @@ +/* + * 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 IQ2_S 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 +constexpr int kEntries = kIq2sEntries; +constexpr int kGroups = 16; // one 4-bit local scale per 16 values +constexpr int kVectorsPerGroup = 2; +constexpr int kVectors = kGroups * kVectorsPerGroup; // 32 codebook vectors per block +constexpr int kLocalScales = 16; +constexpr int kSubBlocks = 8; +constexpr int kLowOffset = kScaleBytes; // 32 low index bytes +constexpr int kSignOffset = kLowOffset + kVectors; // 32 sign masks +constexpr int kHighOffset = kSignOffset + kVectors; // 8 bytes, two high bits per vector +constexpr int kLocalScaleOffset = kHighOffset + kSubBlocks; +constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; +constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. + +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"); +static_assert(kPayloadBytes == 82, "IQ2_S blocks are 82 bytes"); + +// IQ2_S stores all eight sign bits per vector rather than seven plus a parity bit, so unlike +// IQ2_XS and IQ2_XXS the encoder never has to flip a coordinate: the best signs are simply the +// input's, and the search compares magnitudes directly. +__device__ __forceinline__ float magnitude_dot(const float *x, const float *q) { + float dot = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + dot = fmaf(fabsf(x[j]), q[j], dot); + return dot; +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + extern __shared__ float shared_grid[]; // kEntries * kVectorSize, then kEntries norms + float *grid_norm = shared_grid + kEntries * kVectorSize; + __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; + + for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x) + shared_grid[i] = grid[i]; + __syncthreads(); + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float norm = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float q = shared_grid[entry * kVectorSize + j]; + norm = fmaf(q, q, norm); + } + grid_norm[entry] = norm; + } + __syncthreads(); + + 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); + if (!store_block_scale(payload, d_bits)) + return; + // The high-bit bytes are OR-ed into below, so they must start clear. + if (tid < kSubBlocks) + payload[kHighOffset + tid] = 0; + __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 = 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); + } + 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) { + const float dot = magnitude_dot(x, shared_grid + entry * kVectorSize); +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + const float scale = d * (2 * local + 1) * kLocalScaleStep; + local_best[local] = + fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } + 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) * kLocalScaleStep; + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm = 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); + } + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = + clamped_quant_error(xnorm, magnitude_dot(x, shared_grid + entry * kVectorSize), + grid_norm[entry], selected_scale); + const unsigned long long candidate = error_key(error, entry); + key = candidate < key ? candidate : key; + } + key = block_min_key(key, warp_keys); + if (tid == 0) { + const int entry = static_cast(key & (kEntries - 1)); + int sign_mask = 0; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + sign_mask |= static_cast(x[j] < 0.0f) << j; + const int slot = group * kVectorsPerGroup + vector; // 0..31, GGML's qs index + payload[kLowOffset + slot] = static_cast(entry & 0xFF); + payload[kSignOffset + slot] = static_cast(sign_mask); + // Two high index bits per vector, packed four to a byte in vector order. + payload[kHighOffset + slot / 4] |= + static_cast(((entry >> 8) & 0x3) << (2 * (slot % 4))); + } + __syncthreads(); + } + } + + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); +} + +} // namespace + +at::Tensor iq2_s_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("IQ2_S", 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(); + // The 1024-entry codebook and its norms exceed the static shared limit, so they are dynamic. + const size_t smem = static_cast(kEntries) * (kVectorSize + 1) * sizeof(float); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_s_pack", [&] { + auto kernel = encode; + C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem))); + kernel<<(num_blocks), kThreads, smem, 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/iq2_xxs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu new file mode 100644 index 00000000000..c501273a669 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu @@ -0,0 +1,240 @@ +/* + * 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 IQ2_XXS 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 +constexpr int kEntries = kIq2xxsEntries; +constexpr int kGroups = 8; // one 4-bit local scale per 32 values +constexpr int kVectorsPerGroup = 4; // four 8-value codebook vectors per group +constexpr int kLocalScales = 16; +constexpr int kRecordBytes = 8; // four index bytes then one little-endian uint32 +constexpr int kCodeOffset = kScaleBytes; +constexpr int kPayloadBytes = kCodeOffset + kGroups * kRecordBytes; +constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. + +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"); +static_assert(kPayloadBytes == 66, "IQ2_XXS blocks are 66 bytes"); + +// Dot product of |x| against one codebook vector, under the format's even-parity sign rule. +// IQ2_XXS stores seven sign bits per vector and recovers the eighth from their parity, exactly as +// IQ2_XS does, so an odd sign pattern must flip the coordinate with the smallest |x| * q penalty. +__device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) { + float dot = 0.0f; + float weakest = FLT_MAX; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + dot += term; + weakest = fminf(weakest, term); + } + return odd_parity ? dot - 2.0f * weakest : dot; +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + __shared__ float shared_grid[kEntries * kVectorSize]; + __shared__ float grid_norm[kEntries]; + __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]; + __shared__ uint8_t entry_bytes[kGroups * kVectorsPerGroup]; + __shared__ uint8_t sign_bits[kGroups * kVectorsPerGroup]; + + const int tid = threadIdx.x; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x) + shared_grid[i] = grid[i]; + __syncthreads(); + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float norm = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float q = shared_grid[entry * kVectorSize + j]; + norm = fmaf(q, q, norm); + } + grid_norm[entry] = norm; + } + __syncthreads(); + + 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); + if (!store_block_scale(payload, d_bits)) + return; + +#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 = 0.0f; + int negative_count = 0; + 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); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + 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) { + const float *q = shared_grid + entry * kVectorSize; + const float dot = even_parity_dot(x, q, odd_parity); +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + const float scale = d * (2 * local + 1) * kLocalScaleStep; + local_best[local] = + fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } + 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) * kLocalScaleStep; + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + 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); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = clamped_quant_error( + xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), + grid_norm[entry], selected_scale); + const unsigned long long candidate = error_key(error, entry); + key = candidate < key ? candidate : key; + } + key = block_min_key(key, warp_keys); + if (tid == 0) { + const int entry = static_cast(key & (kEntries - 1)); + const float *q = shared_grid + entry * kVectorSize; + int flip_index = 0; + float weakest = fabsf(x[0]) * q[0]; +#pragma unroll + for (int j = 1; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + if (term < weakest) { + weakest = term; + flip_index = j; + } + } + int sign_mask = 0; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + bool is_negative = x[j] < 0.0f; + if (odd_parity && j == flip_index) + is_negative = !is_negative; + sign_mask |= static_cast(is_negative) << j; + } + const int slot = group * kVectorsPerGroup + vector; + entry_bytes[slot] = static_cast(entry); + sign_bits[slot] = static_cast(sign_mask & 0x7f); + } + __syncthreads(); + } + } + + // One 8-byte record per group: four index bytes, then a uint32 holding four 7-bit sign + // indices in bits 0..27 and the 4-bit local scale in bits 28..31. + if (tid < kGroups) { + uint8_t *record = payload + kCodeOffset + tid * kRecordBytes; + const int base = tid * kVectorsPerGroup; +#pragma unroll + for (int j = 0; j < kVectorsPerGroup; ++j) + record[j] = entry_bytes[base + j]; + const uint32_t aux = static_cast(sign_bits[base]) | + (static_cast(sign_bits[base + 1]) << 7) | + (static_cast(sign_bits[base + 2]) << 14) | + (static_cast(sign_bits[base + 3]) << 21) | + (static_cast(locals[tid]) << 28); +#pragma unroll + for (int j = 0; j < 4; ++j) + record[kVectorsPerGroup + j] = static_cast(aux >> (8 * j)); + } +} + +} // namespace + +at::Tensor iq2_xxs_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("IQ2_XXS", 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(), "iq2_xxs_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/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 08ca0dbb719..e6f7f5961db 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -95,6 +95,9 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False): kernels_ggml / "ggml.cpp", kernels_ggml / "iq1_s.cu", kernels_ggml / "iq2_xs.cu", + kernels_ggml / "iq2_xxs.cu", + kernels_ggml / "iq2_s.cu", + kernels_ggml / "iq1_m.cu", ], cuda_version_specifiers=">=11.8", fail_msg="GGML IQ CUDA packing extension is unavailable.", diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index 6bdcd31af74..86b7ff5072d 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -39,6 +39,7 @@ import torch +from ..extensions import get_cuda_ext_ggml from .common import ( GGML_BLOCK_SIZE, fake_quantize_with_cache, @@ -75,6 +76,7 @@ _IQ1_M_SUBBLOCKS = 8 _DEFAULT_BLOCK_CHUNK_SIZE = 1024 _DEFAULT_DECODE_CHUNK_SIZE = 4096 +_SCALE_BLOCK_CHUNK_SIZE = 4096 def _predict_iq1_m_scales(blocks: torch.Tensor) -> torch.Tensor: @@ -183,8 +185,7 @@ def quantize_iq1_m( """Pack a floating-point weight into GGML-compatible IQ1_M blocks. Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 56]`` - and ``[weight.ndim]``. There is no CUDA encoder for this format yet, so - packing always runs the torch search. + and ``[weight.ndim]``. """ validate_weight(weight, "IQ1_M") validate_block_chunk_size(block_chunk_size) @@ -192,11 +193,21 @@ def quantize_iq1_m( logical_shape = torch.tensor(weight.shape, dtype=torch.int64) 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) ] - packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ1_M_BLOCK_SIZE, IQ1_M_BLOCK_BYTES) return torch.cat(chunks).reshape(packed_shape), logical_shape diff --git a/modelopt/torch/quantization/ggml/iq2_s.py b/modelopt/torch/quantization/ggml/iq2_s.py index 25aff6232d0..c4c130511f1 100644 --- a/modelopt/torch/quantization/ggml/iq2_s.py +++ b/modelopt/torch/quantization/ggml/iq2_s.py @@ -37,6 +37,7 @@ import torch +from ..extensions import get_cuda_ext_ggml from .codebooks import iq2_s_grid_bytes from .common import ( GGML_BLOCK_SIZE, @@ -71,6 +72,7 @@ # The grid is twice IQ2_XS's, so the same search tile costs twice the memory. _DEFAULT_BLOCK_CHUNK_SIZE = 128 _DEFAULT_DECODE_CHUNK_SIZE = 4096 +_SCALE_BLOCK_CHUNK_SIZE = 4096 _GRID_CACHE: dict[torch.device, torch.Tensor] = {} @@ -162,8 +164,7 @@ def quantize_iq2_s( """Pack a floating-point weight into GGML-compatible IQ2_S blocks. Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 82]`` - and ``[weight.ndim]``. There is no CUDA encoder for this format yet, so - packing always runs the torch search. + and ``[weight.ndim]``. """ validate_weight(weight, "IQ2_S") validate_block_chunk_size(block_chunk_size) @@ -171,11 +172,21 @@ def quantize_iq2_s( logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ2_S_BLOCK_SIZE) grid = iq2_s_grid(weight.device) + packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ2_S_BLOCK_SIZE, IQ2_S_BLOCK_BYTES) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + scale_chunks = [ + _predict_iq2_s_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) + for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) + ] + packed = extension.iq2_s_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) ] - packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ2_S_BLOCK_SIZE, IQ2_S_BLOCK_BYTES) return torch.cat(chunks).reshape(packed_shape), logical_shape diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py index 106a1c2c9af..01d8bce0a76 100644 --- a/modelopt/torch/quantization/ggml/iq2_xxs.py +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -37,6 +37,7 @@ import torch +from ..extensions import get_cuda_ext_ggml from .codebooks import iq2_xxs_grid_bytes from .common import ( GGML_BLOCK_SIZE, @@ -75,6 +76,7 @@ _DEFAULT_BLOCK_CHUNK_SIZE = 512 # The decode runs on every forward and is launch-bound, so it takes a much larger chunk. _DEFAULT_DECODE_CHUNK_SIZE = 4096 +_SCALE_BLOCK_CHUNK_SIZE = 4096 _GRID_CACHE: dict[torch.device, torch.Tensor] = {} @@ -190,8 +192,7 @@ def quantize_iq2_xxs( Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 66]`` 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. There is no CUDA encoder for this format - yet, so packing always runs the torch search. + treated as zero during packing. """ validate_weight(weight, "IQ2_XXS") validate_block_chunk_size(block_chunk_size) @@ -199,15 +200,25 @@ def quantize_iq2_xxs( logical_shape = torch.tensor(weight.shape, dtype=torch.int64) blocks = weight.contiguous().reshape(-1, IQ2_XXS_BLOCK_SIZE) grid = iq2_xxs_grid(weight.device) - chunks = [ - _encode_blocks(blocks[start : start + block_chunk_size], grid) - for start in range(0, blocks.shape[0], block_chunk_size) - ] packed_shape = ( *weight.shape[:-1], weight.shape[-1] // IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, ) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + scale_chunks = [ + _predict_iq2_xxs_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) + for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) + ] + packed = extension.iq2_xxs_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) + ] return torch.cat(chunks).reshape(packed_shape), logical_shape From 0892bda9032994fdd9adc3c0f376633c3f61b76c Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 21:52:56 +0000 Subject: [PATCH 5/6] Bring every GGML IQ format to the same tests and surface The five IQ formats differ only in codebook size, payload layout and bits per weight, so anything asserted for one should hold for all. They were not tested that way: IQ1_S and IQ2_XS had a per-format file each while the three added formats shared a smaller one, and each side tested things the other did not. Replace the split with one parametrized module per layer -- unit and CUDA -- so a format is a row in a table and a new format inherits the whole contract. That closes gaps in both directions: the new formats gain the grid, metadata, payload-field, saturation, underflow, invalid-shape, scalar-payload and default-dtype checks; IQ1_S and IQ2_XS gain the determinism and at-scale reconstruction checks the newer ones had. Generalising the underflow test exposed a real asymmetry: the IQ1 formats divide by a native max of 16.875 against the IQ2 formats' 166.6, so the -1e-6 weight the IQ2_XS test used still yields an FP16 subnormal for IQ1 rather than a zero scale. That test had only ever existed for IQ2_XS, so the difference had never been visible. The shared test uses a magnitude that underflows all five. Two surface asymmetries go with it: IQ1_S now exposes _predict_iq1_s_scales like the other four instead of computing its anchor inline, and IQ1_M exposes iq1_m_grid aliasing the IQ1_S table it shares. Also add the recipes the three formats were missing -- numerics, model preset and general/ptq entry each -- so they are reachable by --recipe rather than only by num_bits, and add all five to the hf_ptq example matrix. general/ptq now holds 31 recipes; ptq.md is updated to match. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/iq1_m.py | 10 + modelopt/torch/quantization/ggml/iq1_s.py | 17 +- modelopt_recipes/configs/numerics/iq1_m.yaml | 26 ++ modelopt_recipes/configs/numerics/iq2_s.yaml | 26 ++ .../configs/numerics/iq2_xxs.yaml | 26 ++ .../configs/ptq/presets/model/iq1_m.yaml | 32 ++ .../configs/ptq/presets/model/iq2_s.yaml | 32 ++ .../configs/ptq/presets/model/iq2_xxs.yaml | 32 ++ modelopt_recipes/general/ptq/iq1_m.yaml | 27 ++ modelopt_recipes/general/ptq/iq2_s.yaml | 27 ++ modelopt_recipes/general/ptq/iq2_xxs.yaml | 27 ++ modelopt_recipes/ptq.md | 15 +- tests/examples/hf_ptq/test_llm_ptq.py | 14 +- .../quantization/test_iq_formats_cuda.py | 185 ++++++++++++ tests/unit/recipe/test_presets.py | 9 + .../torch/quantization/test_iq_formats.py | 280 ++++++++++++++++++ .../test_iq_llama_cpp_conformance.py | 149 ---------- 17 files changed, 771 insertions(+), 163 deletions(-) create mode 100644 modelopt_recipes/configs/numerics/iq1_m.yaml create mode 100644 modelopt_recipes/configs/numerics/iq2_s.yaml create mode 100644 modelopt_recipes/configs/numerics/iq2_xxs.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq1_m.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml create mode 100644 modelopt_recipes/general/ptq/iq1_m.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_s.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_xxs.yaml create mode 100644 tests/gpu/torch/quantization/test_iq_formats_cuda.py create mode 100644 tests/unit/torch/quantization/test_iq_formats.py delete mode 100644 tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index 86b7ff5072d..12ee59b4eb1 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -56,6 +56,7 @@ "IQ1_M_EFFECTIVE_BITS", "dequantize_iq1_m", "iq1_m_fake_quant", + "iq1_m_grid", "quantize_iq1_m", ] @@ -79,6 +80,15 @@ _SCALE_BLOCK_CHUNK_SIZE = 4096 +def iq1_m_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ1_M ternary grid as float32. + + IQ1_M indexes the same 2048-entry table as IQ1_S; this alias exists so every format + exposes a grid accessor under its own name. + """ + return iq1_s_grid(device) + + def _predict_iq1_m_scales(blocks: torch.Tensor) -> torch.Tensor: """Predict one FP16 super-block scale for each flattened block.""" x = narrow_to_float32(blocks) diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index a43e1699db1..75c62306204 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -87,6 +87,18 @@ def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: return _GRID_CACHE[resolved_device] +def _predict_iq1_s_scales(blocks: torch.Tensor) -> torch.Tensor: + """Predict one FP16 super-block scale for each flattened block. + + The fixed-scale search favors a compressed super-block scale, so this empirical anchor + puts d below the full-range value. The CUDA encoder computes the same quantity in its + own scale kernel; this is the reference the torch path uses. + """ + x = narrow_to_float32(blocks) + amax = x.abs().amax(dim=1) + return ((amax / _IQ1_S_NATIVE_MAX) * _IQ1_S_SCALE_ANCHOR).clamp(max=65504.0).to(torch.float16) + + 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) @@ -95,10 +107,7 @@ def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: 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 = _predict_iq1_s_scales(x) d_float = d.float() best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device) 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/numerics/iq2_s.yaml b/modelopt_recipes/configs/numerics/iq2_s.yaml new file mode 100644 index 00000000000..861179f8462 --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_s.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. + +# IQ2_S weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_s +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq2_s", so the generic estimator cannot derive the +# storage cost: 82 packed bytes * 8 / 256 weights. Keep in sync with IQ2_S_BLOCK_BYTES. +effective_bits: 2.5625 +block_sizes: + -1: 256 +backend: ggml diff --git a/modelopt_recipes/configs/numerics/iq2_xxs.yaml b/modelopt_recipes/configs/numerics/iq2_xxs.yaml new file mode 100644 index 00000000000..2c7a61a08ab --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_xxs.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. + +# IQ2_XXS weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_xxs +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq2_xxs", so the generic estimator cannot derive the +# storage cost: 66 packed bytes * 8 / 256 weights. Keep in sync with IQ2_XXS_BLOCK_BYTES. +effective_bits: 2.0625 +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/configs/ptq/presets/model/iq2_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml new file mode 100644 index 00000000000..0b3ad6154e1 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_s.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 IQ2_S 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 + iq2_s: configs/numerics/iq2_s + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_s + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml new file mode 100644 index 00000000000..b02b17b417b --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.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 IQ2_XXS 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 + iq2_xxs: configs/numerics/iq2_xxs + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_xxs + - 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/general/ptq/iq2_s.yaml b/modelopt_recipes/general/ptq/iq2_s.yaml new file mode 100644 index 00000000000..733388de6df --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_s.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. + +# IQ2_S weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_s + +metadata: + description: >- + Applies uniform GGML-compatible IQ2_S 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/general/ptq/iq2_xxs.yaml b/modelopt_recipes/general/ptq/iq2_xxs.yaml new file mode 100644 index 00000000000..3ef8e26dd99 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xxs.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. + +# IQ2_XXS weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_xxs + +metadata: + description: >- + Applies uniform GGML-compatible IQ2_XXS 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 fa493a5c945..5a8062ec184 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 28 general/ptq/ recipes (click to expand) +All 31 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -59,7 +59,10 @@ 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) | -| `iq2_xs` | IQ2_XS W2A16, 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) |
@@ -138,9 +141,11 @@ 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_xs`** — GGML-compatible IQ1_S or IQ2_XS 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. No calibration data is +- **`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 five + formats 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 6ae697b3c09..2282e55d81a 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -76,13 +76,17 @@ 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. These encoders require every weight's input - # dimension to be a multiple of 256; TinyLlama's 2048 and 5632 both are. Neither - # recipe calibrates -- both set algorithm: null -- so the only IQ-specific cost is - # packing each weight once and decoding it on each forward. 95s and 103s on 2xH100, - # inside the 300s tests/examples default. + # 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"), PTQCommand(quant="nvfp4"), PTQCommand(quant="nvfp4_awq_lite"), # autoquant (recipe-driven) diff --git a/tests/gpu/torch/quantization/test_iq_formats_cuda.py b/tests/gpu/torch/quantization/test_iq_formats_cuda.py new file mode 100644 index 00000000000..f6a7e419f75 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq_formats_cuda.py @@ -0,0 +1,185 @@ +# 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. + +"""CUDA encoders for every GGML IQ format. + +Parametrized over all five rather than per format, so a behaviour asserted for one is +asserted for all. IQ1_S is the odd one out only in its entry point: it derives the block +scale in its own kernel instead of taking a precomputed one. +""" + +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, + IQ2_XXS_BLOCK_BYTES, +) + +# 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), +} + + +def _extension(): + extension = get_cuda_ext_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +def _pack(name, weight): + module, packer, _, takes_scales = FORMATS[name] + grid = getattr(module, f"{name}_grid")("cuda") + if not takes_scales: + return getattr(_extension(), packer)(weight, grid) + blocks = weight.contiguous().reshape(-1, 256) + scales = getattr(module, f"_predict_{name}_scales")(blocks) + return getattr(_extension(), packer)(weight, grid, scales) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch, name): + """Byte parity with the reference encoder on a fixed small weight. + + Exact parity is asserted on this input, not in general: the two encoders evaluate the same + squared error with different floating-point fusion, so where two local scales fall within a + float32 ULP they can round to different sides. That happens in roughly one block in several + thousand, costs under 1e-8 of relative reconstruction error, and favours neither encoder -- + see ``test_cuda_pack_reconstruction_matches_pytorch_at_scale``. + """ + module, _, block_bytes, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _pack(name, weight).reshape(8, 2, block_bytes) + packed_again = _pack(name, weight).reshape(8, 2, block_bytes) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, shape = getattr(module, f"quantize_{name}")(weight) + reconstructed = getattr(module, f"dequantize_{name}")(packed, shape) + + assert packed.shape == (8, 2, block_bytes) + 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 + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_is_deterministic_on_one_device(name): + """The contract reproducibility actually needs: same machine, same bytes.""" + generator = torch.Generator(device="cuda").manual_seed(7) + weight = torch.randn((64, 1024), generator=generator, device="cuda", dtype=torch.bfloat16) + first = _pack(name, weight) + for _ in range(3): + assert torch.equal(_pack(name, weight), first) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_reconstruction_matches_pytorch_at_scale(monkeypatch, name): + """Over many blocks the encoders may disagree on a near-tied scale, but not on quality.""" + module, _, block_bytes, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(11) + weight = torch.randn((128, 2048), generator=generator, device="cuda", dtype=torch.bfloat16) + blocks = weight.shape[0] * weight.shape[1] // 256 + + packed = _pack(name, weight).reshape(weight.shape[0], weight.shape[1] // 256, block_bytes) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, shape = getattr(module, f"quantize_{name}")(weight) + + dequantize = getattr(module, f"dequantize_{name}") + target = weight.float() + denominator = target.square().sum() + cuda_error = ( + dequantize(packed, shape, dtype=torch.float32) - target + ).square().sum() / denominator + torch_error = ( + dequantize(reference, shape, dtype=torch.float32) - target + ).square().sum() / denominator + + differing = int((packed != reference).any(dim=-1).sum()) + assert differing <= blocks // 1000, f"{differing} of {blocks} blocks differ" + assert torch.isclose(cuda_error, torch_error, rtol=1e-5) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_zero_encoding_matches_ggml_block_layout(name): + module, _, block_bytes, _ = FORMATS[name] + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed = _pack(name, weight).reshape(1, 1, block_bytes) + shape = torch.tensor(weight.shape, device="cuda") + + assert not packed.any() + assert torch.equal(getattr(module, f"dequantize_{name}")(packed, shape), weight) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch, name): + module, _, _, _ = FORMATS[name] + weight = torch.zeros((1, 256), device="cuda", dtype=torch.float32) + weight[0, 0] = float("nan") + weight[0, 1] = float("inf") + weight[0, 2] = float("-inf") + + packed, _ = getattr(module, f"quantize_{name}")(weight) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, _ = getattr(module, f"quantize_{name}")(weight) + assert torch.equal(packed, reference) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_falls_back_to_pytorch_encoder(monkeypatch, name): + """Without the extension the format still packs, through the torch search.""" + module, _, block_bytes, _ = FORMATS[name] + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(5) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = getattr(module, f"quantize_{name}")(weight) + reconstructed = getattr(module, f"dequantize_{name}")(packed, shape) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + + assert packed.shape == (2, 1, block_bytes) + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_float64_matches_pytorch_encoder(monkeypatch, name): + module, _, _, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(99) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.float64) + + packed, _ = getattr(module, f"quantize_{name}")(weight) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, _ = getattr(module, f"quantize_{name}")(weight) + assert torch.equal(reference, packed) diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 2bf899cf650..f002bcc3c59 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -34,10 +34,16 @@ 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, + IQ2_S_EFFECTIVE_BITS, IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS, + IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_EFFECTIVE_BITS, ) @@ -135,7 +141,10 @@ 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), ], ) def test_iq_recipe_matches_packing_contract(qformat, block_size, effective_bits): diff --git a/tests/unit/torch/quantization/test_iq_formats.py b/tests/unit/torch/quantization/test_iq_formats.py new file mode 100644 index 00000000000..fec908dc0e2 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq_formats.py @@ -0,0 +1,280 @@ +# 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. + +"""Behaviour every GGML IQ format shares, exercised identically for all five. + +The formats differ only in codebook size, payload layout and bits per weight. Anything that +should hold for one should hold for all, so the contract lives here once and is parametrized +rather than duplicated per format -- a new format is a row in ``FORMATS``. +""" + +import numpy as np +import pytest +import torch +from _test_utils.torch.quantization.iq_llama_cpp_vectors import ( + expected_values, + formats, + packed_blocks, +) + +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.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer + +# name -> (module, packed bytes per block, codebook entries, bits per weight) +FORMATS = { + "iq1_s": (iq1_s_module, 50, 2048, 1.5625), + "iq1_m": (iq1_m_module, 56, 2048, 1.75), + "iq2_xxs": (iq2_xxs_module, 66, 256, 2.0625), + "iq2_xs": (iq2_xs_module, 74, 512, 2.3125), + "iq2_s": (iq2_s_module, 82, 1024, 2.5625), +} +NAMES = sorted(FORMATS) +# IQ1 grids are ternary; IQ2 grids hold the magnitudes 8, 25 and 43. +TERNARY = {"iq1_s", "iq1_m"} + + +def _parts(name): + module, block_bytes, entries, bits = FORMATS[name] + return ( + module, + getattr(module, f"quantize_{name}"), + getattr(module, f"dequantize_{name}"), + getattr(module, f"{name}_grid"), + block_bytes, + entries, + bits, + ) + + +@pytest.mark.parametrize("name", NAMES) +def test_canonical_grid(name): + _, _, _, grid_fn, _, entries, _ = _parts(name) + grid = grid_fn() + + assert grid.shape == (entries, 8) + assert grid.dtype == torch.float32 + if name in TERNARY: + assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0} + assert grid[0].tolist() == [-1.0] * 8 + else: + assert set(grid.unique().tolist()) <= {8.0, 25.0, 43.0} + assert grid[0].tolist() == [8.0] * 8 + + +@pytest.mark.parametrize("name", NAMES) +def test_grid_normalizes_unindexed_cuda_device(monkeypatch, name): + module, _, _, grid_fn, _, _, _ = _parts(name) + # IQ1_M shares the IQ1_S cache, being the same table. + cache_owner = iq1_s_module if name == "iq1_m" else module + cached = torch.empty(0) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 7) + monkeypatch.setitem(cache_owner._GRID_CACHE, torch.device("cuda", 7), cached) + + assert grid_fn("cuda") is cached + + +@pytest.mark.parametrize("name", NAMES) +def test_effective_bits_matches_the_payload(name): + module, _, _, _, block_bytes, _, bits = _parts(name) + assert getattr(module, f"{name.upper()}_BLOCK_BYTES") == block_bytes + assert getattr(module, f"{name.upper()}_EFFECTIVE_BITS") == pytest.approx(bits) + assert getattr(module, f"{name.upper()}_BLOCK_SIZE") == 256 + + +@pytest.mark.parametrize("name", NAMES) +def test_round_trip_and_payload_fields(name): + _, quantize, dequantize, _, block_bytes, _, _ = _parts(name) + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize(weight) + reconstructed = dequantize(packed, shape) + chunked = dequantize(packed, shape, block_chunk_size=1) + + assert packed.shape == (2, 2, block_bytes) + assert packed.dtype == torch.uint8 + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + assert torch.equal(reconstructed, chunked) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", NAMES) +def test_decode_is_invariant_to_chunk_size(name): + """Chunking the decode is a memory bound, not a numerical choice.""" + _, quantize, dequantize, _, _, _, _ = _parts(name) + torch.manual_seed(0) + weight = torch.randn(3, 1024, dtype=torch.bfloat16) + packed, shape = quantize(weight) + + reference = dequantize(packed, shape, dtype=weight.dtype, block_chunk_size=1) + for chunk in (2, 7, 4096): + assert torch.equal( + dequantize(packed, shape, dtype=weight.dtype, block_chunk_size=chunk), reference + ) + + +@pytest.mark.parametrize("name", NAMES) +def test_zero_block_has_canonical_zero_encoding(name): + _, quantize, dequantize, _, _, _, _ = _parts(name) + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + packed, shape = quantize(weight) + + assert not packed.any() + assert torch.equal(dequantize(packed, shape), weight) + + +@pytest.mark.parametrize("name", NAMES) +def test_underflowed_scale_has_canonical_zero_encoding(name): + """A block whose scale rounds to zero in FP16 packs as all zero bytes. + + The magnitude has to clear every format's threshold at once: the IQ1 formats divide by a + native max of 16.875 against the IQ2 formats' 166.6, so a weight that underflows an IQ2 + scale still lands on an FP16 subnormal for IQ1. + """ + _, quantize, dequantize, _, _, _, _ = _parts(name) + weight = torch.full((1, 256), -1e-8, dtype=torch.bfloat16) + packed, shape = quantize(weight) + + assert not packed.any() + assert torch.equal(dequantize(packed, shape), torch.zeros_like(weight)) + + +@pytest.mark.parametrize("name", NAMES) +def test_requires_complete_last_dimension_blocks(name): + _, quantize, _, _, _, _, _ = _parts(name) + with pytest.raises(ValueError, match="last weight dimension"): + quantize(torch.ones(2, 257)) + + +@pytest.mark.parametrize("name", NAMES) +def test_treats_nonfinite_values_as_zero(name): + _, quantize, _, _, _, _, _ = _parts(name) + weight = torch.zeros(1, 256) + weight[0, 0] = float("nan") + weight[0, 1] = float("inf") + weight[0, 2] = float("-inf") + + packed, _ = quantize(weight) + assert not packed.any() + + +@pytest.mark.parametrize("name", NAMES) +def test_saturates_finite_values_above_the_float32_range(name): + """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. + """ + _, quantize, _, _, _, _, _ = _parts(name) + torch.manual_seed(0) + 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(weight) + assert torch.equal(packed, quantize(saturated)[0]) + assert not torch.equal(packed, quantize(zeroed)[0]) + + +@pytest.mark.parametrize("name", NAMES) +def test_rejects_invalid_shape_metadata(name): + _, _, dequantize, _, block_bytes, _, _ = _parts(name) + packed = torch.zeros((1, 1, block_bytes), dtype=torch.uint8) + for weight_shape in (torch.tensor([[1, 256]]), torch.tensor([1.0, 256.0]), torch.tensor([257])): + with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"): + dequantize(packed, weight_shape) + + +@pytest.mark.parametrize("name", NAMES) +def test_rejects_scalar_packed_payload(name): + _, _, dequantize, _, _, _, _ = _parts(name) + with pytest.raises(ValueError, match="packed_weights"): + dequantize(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256])) + + +@pytest.mark.parametrize("name", NAMES) +def test_fake_quant_has_pass_through_gradient(name): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + weight = torch.randn(2, 256, requires_grad=True) + quantizer(weight).sum().backward() + assert torch.equal(weight.grad, torch.ones_like(weight)) + + +@pytest.mark.parametrize("name", NAMES) +def test_search_is_independent_of_default_dtype(name): + """The encoder must not inherit a global default dtype; it works in float32 throughout.""" + _, quantize, _, _, _, _, _ = _parts(name) + torch.manual_seed(0) + weight = torch.randn(2, 256) + expected, _ = quantize(weight) + try: + torch.set_default_dtype(torch.float64) + assert torch.equal(quantize(weight)[0], expected) + finally: + torch.set_default_dtype(torch.float32) + + +@pytest.mark.parametrize("name", formats()) +def test_decoder_matches_llama_cpp_on_captured_blocks(name): + """Decode bytes we did not produce and match llama.cpp's own output exactly. + + A round-trip against our own encoder cannot catch a layout error that the encoder makes + symmetrically; these blocks come from a real checkpoint, so they can. + """ + _, _, dequantize, _, block_bytes, _, _ = _parts(name) + blocks = packed_blocks(name) + assert blocks.shape[1] == block_bytes + + count = blocks.shape[0] + decoded = dequantize( + torch.from_numpy(blocks).reshape(count, 1, block_bytes), + torch.tensor([count, 256]), + dtype=torch.float32, + ).reshape(count, 256) + + assert np.array_equal(decoded.numpy(), expected_values(name)) + + +def test_every_format_has_conformance_vectors(): + """A new format must arrive with blocks captured from a real llama.cpp checkpoint.""" + assert sorted(formats()) == NAMES + + +def test_error_decreases_with_bit_width(): + """More bits must buy less error, or a format's scale handling is wrong.""" + generator = torch.Generator().manual_seed(7) + weight = torch.randn((4, 1024), generator=generator) + errors = [] + for name in sorted(NAMES, key=lambda n: FORMATS[n][3]): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + errors.append(float((quantizer(weight) - weight).square().mean())) + + assert errors == sorted(errors, reverse=True), dict(zip(sorted(NAMES), errors)) diff --git a/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py b/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py deleted file mode 100644 index eba6ba655e7..00000000000 --- a/tests/unit/torch/quantization/test_iq_llama_cpp_conformance.py +++ /dev/null @@ -1,149 +0,0 @@ -# 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. - -"""Conformance of the GGML IQ decoders against real llama.cpp-produced blocks.""" - -import numpy as np -import pytest -import torch -from _test_utils.torch.quantization.iq_llama_cpp_vectors import ( - expected_values, - formats, - packed_blocks, -) - -from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.ggml import ( - IQ1_M_BLOCK_BYTES, - IQ1_S_BLOCK_BYTES, - IQ2_S_BLOCK_BYTES, - IQ2_XS_BLOCK_BYTES, - IQ2_XXS_BLOCK_BYTES, - dequantize_iq1_m, - dequantize_iq1_s, - dequantize_iq2_s, - dequantize_iq2_xs, - dequantize_iq2_xxs, - quantize_iq1_m, - quantize_iq2_s, - quantize_iq2_xxs, -) -from modelopt.torch.quantization.nn import TensorQuantizer - -DECODERS = { - "iq1_s": (dequantize_iq1_s, IQ1_S_BLOCK_BYTES), - "iq1_m": (dequantize_iq1_m, IQ1_M_BLOCK_BYTES), - "iq2_xxs": (dequantize_iq2_xxs, IQ2_XXS_BLOCK_BYTES), - "iq2_xs": (dequantize_iq2_xs, IQ2_XS_BLOCK_BYTES), - "iq2_s": (dequantize_iq2_s, IQ2_S_BLOCK_BYTES), -} -NEW_FORMATS = { - "iq1_m": (quantize_iq1_m, dequantize_iq1_m, IQ1_M_BLOCK_BYTES), - "iq2_xxs": (quantize_iq2_xxs, dequantize_iq2_xxs, IQ2_XXS_BLOCK_BYTES), - "iq2_s": (quantize_iq2_s, dequantize_iq2_s, IQ2_S_BLOCK_BYTES), -} - - -@pytest.mark.parametrize("name", formats()) -def test_decoder_matches_llama_cpp_on_captured_blocks(name): - """Decode bytes we did not produce and match llama.cpp's own output exactly. - - A round-trip against our own encoder cannot catch a layout error that the encoder - makes symmetrically; these blocks come from a real checkpoint, so they can. - """ - decode, block_bytes = DECODERS[name] - blocks = packed_blocks(name) - assert blocks.shape[1] == block_bytes - - count = blocks.shape[0] - decoded = decode( - torch.from_numpy(blocks).reshape(count, 1, block_bytes), - torch.tensor([count, 256]), - dtype=torch.float32, - ).reshape(count, 256) - - assert np.array_equal(decoded.numpy(), expected_values(name)) - - -@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) -def test_round_trip_is_close_and_shaped(name): - quantize, dequantize, block_bytes = NEW_FORMATS[name] - generator = torch.Generator().manual_seed(1234) - weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) - - packed, shape = quantize(weight) - reconstructed = dequantize(packed, shape) - chunked = dequantize(packed, shape, block_chunk_size=1) - - assert packed.shape == (2, 2, block_bytes) - assert reconstructed.shape == weight.shape - assert reconstructed.dtype == torch.bfloat16 - # Chunking the decode is a memory bound, not a numerical choice. - assert torch.equal(reconstructed, chunked) - normalized_mse = ( - reconstructed.float() - weight.float() - ).square().mean() / weight.float().square().mean() - assert normalized_mse < 0.25 - - -@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) -def test_zero_block_has_canonical_zero_encoding(name): - quantize, dequantize, block_bytes = NEW_FORMATS[name] - packed, shape = quantize(torch.zeros(1, 256)) - assert torch.equal(packed, torch.zeros_like(packed)) - assert torch.equal(dequantize(packed, shape), torch.zeros(1, 256, dtype=torch.bfloat16)) - - -@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) -def test_requires_complete_last_dimension_blocks(name): - quantize, _, _ = NEW_FORMATS[name] - with pytest.raises(ValueError, match="last weight dimension"): - quantize(torch.ones(2, 257)) - - -@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) -def test_nonfinite_values_are_treated_as_zero(name): - quantize, _, _ = NEW_FORMATS[name] - weight = torch.zeros(1, 256) - weight[0, 0] = float("nan") - weight[0, 1] = float("inf") - packed, _ = quantize(weight) - assert torch.equal(packed, torch.zeros_like(packed)) - - -@pytest.mark.parametrize("name", sorted(NEW_FORMATS)) -def test_fake_quant_has_pass_through_gradient(name): - quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") - ) - weight = torch.randn(2, 256, requires_grad=True) - quantizer(weight).sum().backward() - assert torch.equal(weight.grad, torch.ones_like(weight)) - - -def test_error_decreases_with_bit_width(): - """More bits must buy less error, or a format's scale handling is wrong.""" - generator = torch.Generator().manual_seed(7) - weight = torch.randn((4, 1024), generator=generator) - errors = {} - for name in ("iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s"): - quantizer = TensorQuantizer( - QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") - ) - reconstructed = quantizer(weight) - errors[name] = float((reconstructed - weight).square().mean()) - - ordered = [errors[n] for n in ("iq1_s", "iq1_m", "iq2_xxs", "iq2_xs", "iq2_s")] - assert ordered == sorted(ordered, reverse=True), errors From 029821fef417f5bd8f31e3d44397a6f1d9a878ea Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 21:54:08 +0000 Subject: [PATCH 6/6] Drop the torch-only caveat from the IQ changelog entry The three formats now ship CUDA encoders and recipes, so the note about slow PyTorch packing no longer describes them. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 67744b6c3ea..1f2ad30b05e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,7 +14,7 @@ Changelog *Quantization* - 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 ``iq1_m``, ``iq2_xxs`` and ``iq2_s`` weight-only quantization, completing the GGML IQ formats at one and two bits. The same 256-value block constraint applies; these three pack with the PyTorch encoder rather than a CUDA kernel, so quantizing a large model with them is slower than with ``iq1_s`` or ``iq2_xs``. +- Add ``iq1_m``, ``iq2_xxs`` and ``iq2_s`` weight-only quantization with CUDA encoders and ``general/ptq`` recipes, completing the GGML IQ formats at one and two bits. The same 256-value block constraint applies as for ``iq1_s`` and ``iq2_xs``. - 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.