diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0a6afcb0017..1f2ad30b05e 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 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. 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/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/__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..12ee59b4eb1 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -0,0 +1,309 @@ +# 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 ..extensions import get_cuda_ext_ggml +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", + "iq1_m_grid", + "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 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 +_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) + 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) + 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) + + 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. + 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]``. + """ + 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) + 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) + ] + 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/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/torch/quantization/ggml/iq2_s.py b/modelopt/torch/quantization/ggml/iq2_s.py new file mode 100644 index 00000000000..c4c130511f1 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_s.py @@ -0,0 +1,252 @@ +# 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 ..extensions import get_cuda_ext_ggml +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 +_SCALE_BLOCK_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]``. + """ + 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) + 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) + ] + 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..01d8bce0a76 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -0,0 +1,286 @@ +# 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 ..extensions import get_cuda_ext_ggml +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 +_SCALE_BLOCK_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. + """ + 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) + 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 + + +@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/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/_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/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_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_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))