diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1ab4d7e7b85..48e973c6dee 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,6 +103,8 @@ repos: exclude: > (?x)^( modelopt/torch/quantization/utils/calib_utils.py| + modelopt/torch/quantization/ggml/iq1_s.py| + modelopt/torch/quantization/ggml/iq2_xs.py| modelopt/onnx/quantization/operators.py| modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 59e999ea2ee..10a7f814601 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add IQ1_S and IQ2_XS weight quantization and unified checkpoint export from Hugging Face and TP=1 Megatron models. Supported weights are encoded in GGML-compatible 256-value blocks and stored as packed ``uint8`` weight tensors in safetensors. - 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. *Misc* diff --git a/LICENSE b/LICENSE index c58bddda878..57c1104ae67 100644 --- a/LICENSE +++ b/LICENSE @@ -250,6 +250,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2023 DeepSeek Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors + Copyright (c) 2023-2026 The ggml authors Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index ccef639d00e..2923c227a79 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -50,6 +50,36 @@ The unified HF export API supports the following quantization formats: 4. NVFP4_AWQ - NVIDIA 4-bit floating point with AWQ optimization 5. INT4_AWQ - 4-bit integer with AWQ optimization 6. W4A8_AWQ - 4-bit weights and 8-bit activations with AWQ optimization +7. IQ1_S - 1-bit importance-aware quantization using the GGML block layout +8. IQ2_XS - 2-bit importance-aware quantization using the GGML block layout + +.. note:: + GGML has no equivalent for ModelOpt's per-tensor FP8 weight-and-activation format. In particular, + GGML does not define a first-class FP8 tensor type with the corresponding per-tensor weight and + activation scale semantics. Converting a ModelOpt FP8 checkpoint to GGUF therefore requires + conversion to another GGML-supported tensor type rather than a lossless FP8 encoding. + +IQ weight representation +~~~~~~~~~~~~~~~~~~~~~~~~ + +For IQ1_S and IQ2_XS, unified export replaces each floating-point ``.weight`` with a +``uint8`` tensor containing byte-exact GGML blocks. Its shape is +``[*logical_shape[:-1], logical_shape[-1] // 256, payload_bytes]``, where ``payload_bytes`` is 50 +for IQ1_S and 74 for IQ2_XS. No separate shape tensor is stored: a loader recovers the logical +shape as ``[*weight.shape[:-2], weight.shape[-2] * 256]``. This is unambiguous because IQ export +requires the logical last dimension to be divisible by 256. + +Each 74-byte IQ2_XS block represents 256 logical weights: + +* bytes 0--1 are the little-endian FP16 super-block scale ``d``; +* bytes 2--65 are 32 little-endian ``uint16`` codes, one per group of eight weights. Each code + contains a 9-bit codebook index and seven stored sign bits; the eighth sign bit is derived from + parity; and +* bytes 66--73 contain sixteen 4-bit local-scale codes, packed two per byte. Each local scale is + shared by two adjacent eight-weight groups. + +The canonical 512-by-8 IQ2_XS codebook is part of the implementation rather than the checkpoint. +The complete block therefore costs ``74 * 8 / 256 = 2.3125`` bits per logical weight. Minimum Framework Versions -------------------------- diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 45fa0c30f3b..a3084ab7e0c 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -117,6 +117,19 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo in ("IQ1_S", "IQ2_XS"): + effective_bits, payload_bytes = (1.5625, 50) if quant_algo == "IQ1_S" else (2.3125, 74) + return { + "weights": { + "dynamic": False, + "num_bits": 1 if quant_algo == "IQ1_S" else 2, + "effective_bits": effective_bits, + "type": "int", + "group_size": 256, + "packing": "ggml", + "block_payload_bytes": payload_bytes, + } + } else: warnings.warn( f"Unsupported quantization algorithm '{quant_algo}' in " @@ -209,6 +222,10 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "targets": ["Linear"], } new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value in ("IQ1_S", "IQ2_XS"): + config_group_details = _quant_algo_to_group_config(quant_algo_value, 256) + config_group_details["targets"] = ["Linear"] + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "NVFP4_SVD": # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 734302690f1..e8ae02090d2 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -218,7 +218,10 @@ def _export_fused_experts( _export_quantized_weight(wrapper, dtype) proj = nn.Module() - proj.weight = wrapper.weight + if isinstance(wrapper.weight, nn.Parameter): + proj.weight = wrapper.weight + else: + proj.register_buffer("weight", wrapper.weight) for attr in ("weight_scale", "weight_scale_2", "input_scale"): if hasattr(wrapper, attr): proj.register_buffer(attr, getattr(wrapper, attr)) diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index b270aec1a68..b4475fa1f7c 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -36,11 +36,21 @@ QUANTIZATION_FP8_PB_REAL = "fp8_pb_real" QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" +QUANTIZATION_IQ1_S = "iq1_s" +QUANTIZATION_IQ2_XS = "iq2_xs" # 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({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL}) +FUSION_FREE_FORMATS = frozenset( + { + QUANTIZATION_FP8, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, + QUANTIZATION_NONE, + QUANTIZATION_FP8_PB_REAL, + } +) KV_CACHE_FP8 = "FP8" KV_CACHE_INT8 = "INT8" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 5320f012a94..3bd75432a51 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -61,6 +61,8 @@ QUANTIZATION_INT4_AWQ, QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -434,6 +436,11 @@ 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.backend != "psx_luts": + raise ValueError("IQ formats require the built-in 'psx_luts' quantization backend") + return weight_quantizer.num_bits + if weight_quantizer.num_bits == 4: assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( "Invalid block_sizes for INT4 quantizer" @@ -682,6 +689,14 @@ 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): + payload_bytes = 50 if v == QUANTIZATION_IQ1_S else 74 + layer_config = { + "quant_algo": v.upper(), + "group_size": 256, + "block_payload_bytes": payload_bytes, + "packing": "ggml", + } else: layer_config = {"quant_algo": v} @@ -1075,6 +1090,7 @@ def _export_key(key: str) -> str: # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.) weight_suffixes = ( "weight", + "weight_shape", "weight_scale", "weight_scale_2", "input_scale", diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 354dda65f40..a4cfc388dbf 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -59,6 +59,7 @@ 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.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper @@ -100,6 +101,8 @@ QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_MXFP8, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -622,6 +625,21 @@ 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 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 + ) + packed_weight, _ = quantize_iq(weight.to(dtype)) + delattr(sub_module, weight_name) + sub_module.register_buffer("weight", packed_weight) + maybe_clear_cuda_cache() + return + weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( sub_module, quantizer_attrs.weight_quantizer ) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 37c896fd609..87eadeb12d1 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,6 +35,7 @@ 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.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 @@ -61,6 +62,8 @@ QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PB_WO, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, @@ -94,6 +97,7 @@ get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, ) from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp @@ -309,10 +313,19 @@ def save_pretrained( is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0 is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank) + quantization_format = self._get_quantization_format(self.model) + if ( + quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + and get_tensor_model_parallel_world_size() != 1 + ): + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + # Main export process layer_state_dicts = self.layer_state_dicts - quantization_format = self._get_quantization_format(self.model) quantization = None if quantization_format in ( QUANTIZATION_FP8_PB_REAL, @@ -325,6 +338,8 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" + elif quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + quantization = quantization_format.upper() if is_last_stage_main_rank: if is_writer_rank: @@ -1027,6 +1042,7 @@ def _get_weight_bias( module: torch.nn.Module, dtype: torch.dtype = torch.float16, name_to_value: dict[str, torch.Tensor] | None = None, + keep_weight_device: bool = False, ) -> dict[str, torch.Tensor]: """Get the weight and bias of the module. @@ -1035,6 +1051,7 @@ def _get_weight_bias( dtype: The data type of the weight and bias. name_to_value: The dictionary to store the weight and bias. A new dict is created if not provided. + keep_weight_device: Keep the weight on its current device instead of moving it to CPU. Returns: The dictionary containing the weight and bias. @@ -1045,7 +1062,9 @@ def _get_weight_bias( # layers whose weight is a placeholder) so callers can use "weight" in name_to_value # as a reliable guard without re-inspecting module.weight. if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0: - weight = module.weight.to(dtype).cpu() + weight = module.weight.to(dtype) + if not keep_weight_device: + weight = weight.cpu() name_to_value["weight"] = weight if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0: @@ -1082,13 +1101,21 @@ def _get_quantized_state( self._record_excluded_module(prefix) block_size = get_weight_block_size(module) - name_to_value = self._get_weight_bias(module, dtype, name_to_value) + is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + name_to_value = self._get_weight_bias( + module, dtype, name_to_value, keep_weight_device=is_iq + ) if "weight" not in name_to_value: return name_to_value, qformat, block_size if qformat == QUANTIZATION_NONE: return name_to_value, qformat, block_size + # IQ formats derive all block metadata directly from the weight and do not use amax or + # separately exported scaling tensors. Keep the weight on-device until its final HF layout + # has been produced, so the CUDA packer can be used. + if is_iq: + return name_to_value, qformat, block_size # Getting the weight scales weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) @@ -1124,6 +1151,15 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str): return weight_scale, weight_scale_2 + @staticmethod + def _get_iq_weight_state( + weight_key: str, weight: torch.Tensor, qformat: str + ) -> dict[str, torch.Tensor]: + """Pack one final-layout weight into the IQ unified-checkpoint representation.""" + quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs + packed_weight, _ = quantize_iq(weight) + return {weight_key: packed_weight.detach().cpu()} + def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None): """Record per-HF-layer quantization metadata for mixed precision exports.""" if qformat in (None, QUANTIZATION_NONE): @@ -1188,7 +1224,9 @@ def _name_remapping( weight = weight + 1.0 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat)) + elif weight_scale is None: self._state_dict[prefix + "weight"] = weight else: self._state_dict[prefix + "weight"] = to_quantized_weight( @@ -1233,7 +1271,14 @@ def _gated_mlp_slicing( gate_proj_weight = weight[:ffn_hidden_size, :] up_proj_weight = weight[ffn_hidden_size:, :] - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update( + self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat) + ) + self._state_dict.update( + self._get_iq_weight_state(up_proj_prefix + "weight", up_proj_weight, qformat) + ) + elif weight_scale is None: self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight self._state_dict[up_proj_prefix + "weight"] = up_proj_weight else: @@ -1399,7 +1444,9 @@ def _grouped_mlp_slicing( name_to_value.pop("weight", None) seen_qformat, seen_block_size = qformat, block_size - weight = state_dict[weight_key].to(self.dtype).cpu() + weight = state_dict[weight_key].to(self.dtype) + if qformat not in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + weight = weight.cpu() weight_scale_cpu = ( weight_scale.detach().cpu().clone() if weight_scale is not None else None ) @@ -1430,7 +1477,13 @@ def _grouped_mlp_slicing( ] for shard_prefix, shard_weight, shard_scale in shards: - if shard_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + local_expert_state.update( + self._get_iq_weight_state( + shard_prefix + "weight", shard_weight, qformat + ) + ) + elif shard_scale is None: local_expert_state[shard_prefix + "weight"] = shard_weight else: local_expert_state[shard_prefix + "weight"] = to_quantized_weight( @@ -1593,7 +1646,10 @@ 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 weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + 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: for key, weight in zip(proj_keys, proj_weights): self._state_dict[key] = weight else: @@ -1708,7 +1764,15 @@ 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 weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + 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() + else: + self._state_dict.update( + self._get_iq_weight_state(proj_prefix + "weight", proj_weight, qformat) + ) + elif weight_scale is None: for key, proj_weight in zip(proj_keys, proj_weights): self._state_dict[key] = proj_weight else: @@ -1842,7 +1906,9 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr merged_input_scale = None # Save the merged weights - if merged_weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) + elif merged_weight_scale is None: self._state_dict[prefix] = merged_weight else: self._state_dict[prefix] = to_quantized_weight( @@ -1953,7 +2019,9 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F merged_input_scale = None # Save the merged weights - if merged_weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) + elif merged_weight_scale is None: # TODO: May need to modify the key name later. self._state_dict[prefix] = merged_weight else: diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp new file mode 100644 index 00000000000..bc8a8fa325d --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp @@ -0,0 +1,31 @@ +/* + * 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 +#include + +at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid); + +at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_cuda(), "IQ1_S packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ1_S packing requires a CUDA grid"); + return iq1_s_pack_cuda(input.contiguous(), grid.contiguous()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("pack", &iq1_s_pack, "Pack a tensor into GGML IQ1_S blocks"); +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu new file mode 100644 index 00000000000..ddf1f762e16 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu @@ -0,0 +1,270 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace { + +constexpr int kBlockSize = 256; +constexpr int kVectorSize = 8; +constexpr int kEntries = 2048; +constexpr int kGroups = 8; +constexpr int kLocalScales = 8; +constexpr int kChoices = 16; +constexpr int kPayloadBytes = 50; +constexpr float kDelta = 0.125f; +constexpr float kNativeMax = 16.875f; + +template __device__ __forceinline__ float load_float(const scalar_t *input) { + return static_cast(*input); +} + +__device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x, + const float *q, float scale, float delta) { + float dot = 0.0f; + float qnorm = 0.0f; + float qsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + dot = fmaf(x[j], q[j], dot); + qnorm = fmaf(q[j], q[j], qnorm); + qsum += q[j]; + } + const float shifted_dot = dot + delta * xsum; + const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; + return fmaxf(fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f); +} + +template +__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { + const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) + return; + + float amax = 0.0f; + const scalar_t *values = input + block * kBlockSize; +#pragma unroll 1 + for (int i = 0; i < kBlockSize; ++i) + amax = fmaxf(amax, fabsf(load_float(values + i))); + const __half scale = __float2half_rn(fminf((amax / kNativeMax) * 0.61f, 65504.0f)); + scale_bits[block] = static_cast(__half_as_ushort(scale)); +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const int64_t *scale_bits, uint8_t *output) { + __shared__ float warp_best[8 * kChoices]; + __shared__ float group_error[kChoices]; + __shared__ unsigned long long warp_keys[8]; + __shared__ int selected_choice; + __shared__ uint16_t selected_entries[4]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + 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 uint16_t d_bits = static_cast(scale_bits[block]); + const float d = __half2float(__ushort_as_half(d_bits)); + if (d_bits == 0) { + if (tid < kPayloadBytes) + payload[tid] = 0; + return; + } + if (tid == 0) { + payload[0] = static_cast(d_bits); + payload[1] = static_cast(d_bits >> 8); + } + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kChoices) + group_error[tid] = 0.0f; + __syncthreads(); + +#pragma unroll + for (int vector = 0; vector < 4; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + float xsum = 0.0f; + const int offset = group * 32 + vector * 8; +#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[kChoices]; +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) + local_best[choice] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float *q = grid + entry * kVectorSize; + float dot = 0.0f; + float qnorm = 0.0f; + float qsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + dot = fmaf(x[j], q[j], dot); + qnorm = fmaf(q[j], q[j], qnorm); + qsum += q[j]; + } +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) { + const int local = choice & 7; + const float delta = choice < 8 ? kDelta : -kDelta; + const float scale = d * (2 * local + 1); + const float shifted_dot = dot + delta * xsum; + const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; + const float error = fmaxf( + fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f); + local_best[choice] = fminf(local_best[choice], error); + } + } +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) { + float value = local_best[choice]; +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + value = fminf(value, __shfl_down_sync(0xffffffff, value, delta)); + if (lane == 0) + warp_best[warp * kChoices + choice] = value; + } + __syncthreads(); + if (tid < kChoices) { + float value = warp_best[tid]; +#pragma unroll + for (int w = 1; w < 8; ++w) + value = fminf(value, warp_best[w * kChoices + tid]); + group_error[tid] += value; + } + __syncthreads(); + } + + if (tid == 0) { + selected_choice = 0; + float best = group_error[0]; +#pragma unroll + for (int choice = 1; choice < kChoices; ++choice) { + if (group_error[choice] < best) { + best = group_error[choice]; + selected_choice = choice; + } + } + } + __syncthreads(); + const int selected_local = selected_choice & 7; + const float selected_delta = selected_choice < 8 ? kDelta : -kDelta; + const float selected_scale = d * (2 * selected_local + 1); + +#pragma unroll + for (int vector = 0; vector < 4; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + float xsum = 0.0f; + const int offset = group * 32 + vector * 8; +#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) { + const float error = + quant_error(xnorm, xsum, x, grid + entry * kVectorSize, selected_scale, selected_delta); + const unsigned long long candidate = + (static_cast(__float_as_uint(error)) << 32) | + static_cast(entry); + key = candidate < key ? candidate : key; + } +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + const auto other = __shfl_down_sync(0xffffffff, key, delta); + key = other < key ? other : key; + } + if (lane == 0) + warp_keys[warp] = key; + __syncthreads(); + if (tid == 0) { + key = warp_keys[0]; +#pragma unroll + for (int w = 1; w < 8; ++w) + key = warp_keys[w] < key ? warp_keys[w] : key; + const uint16_t entry = static_cast(key & 0x7ff); + selected_entries[vector] = entry; + payload[2 + group * 4 + vector] = static_cast(entry); + } + __syncthreads(); + } + + if (tid == 0) { + const uint16_t qh = static_cast( + ((selected_entries[0] >> 8) & 7) | (((selected_entries[1] >> 8) & 7) << 3) | + (((selected_entries[2] >> 8) & 7) << 6) | (((selected_entries[3] >> 8) & 7) << 9) | + (selected_local << 12) | ((selected_choice >> 3) << 15)); + payload[34 + 2 * group] = static_cast(qh); + payload[35 + 2 * group] = static_cast(qh >> 8); + } + __syncthreads(); + } +} + +} // namespace + +at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0, + "input size must be a positive multiple of 256"); + TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.numel() == kEntries * kVectorSize, + "grid must be float32 [2048, 8]"); + TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(num_blocks <= std::numeric_limits::max(), "IQ1_S CUDA grid is too large"); + auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong)); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + const int scale_grid = static_cast((num_blocks + 255) / 256); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_s_pack", [&] { + find_scale<<>>(input.data_ptr(), num_blocks, + scales.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + encode<<(num_blocks), 256, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + scales.data_ptr(), output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp new file mode 100644 index 00000000000..ee90cae2650 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp @@ -0,0 +1,30 @@ +/* + * 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 + +at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid); + +at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_cuda(), "IQ2_XS packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ2_XS packing requires a CUDA grid"); + return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("pack", &iq2_xs_pack, "Pack a tensor into GGML IQ2_XS blocks"); +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu new file mode 100644 index 00000000000..cc798903aff --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu @@ -0,0 +1,292 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace { + +constexpr int kBlockSize = 256; +constexpr int kVectorSize = 8; +constexpr int kEntries = 512; +constexpr int kGroups = 16; +constexpr int kLocalScales = 16; +constexpr int kPayloadBytes = 74; +constexpr float kNativeMax = 166.625f; + +template __device__ __forceinline__ float load_float(const scalar_t *input) { + return static_cast(*input); +} + +__device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm, float scale) { + return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f); +} + +__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 find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { + const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) + return; + + float amax = 0.0f; + float sumsq = 0.0f; + const scalar_t *values = input + block * kBlockSize; +#pragma unroll 1 + for (int i = 0; i < kBlockSize; ++i) { + const float value = load_float(values + i); + amax = fmaxf(amax, fabsf(value)); + sumsq = fmaf(value, value, sumsq); + } + if (amax == 0.0f) { + scale_bits[block] = 0; + return; + } + const float rms = sqrtf(sumsq / kBlockSize); + const float peak_to_rms = rms > 0.0f ? amax / rms : 0.0f; + const float anchor = fminf(0.92f, fmaxf(0.65f, 1.0f - 0.035f * peak_to_rms)); + const __half scale = __float2half_rn(fminf((amax / kNativeMax) * anchor, 65504.0f)); + scale_bits[block] = static_cast(__half_as_ushort(scale)); +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const int64_t *scale_bits, uint8_t *output) { + __shared__ float shared_grid[kEntries * kVectorSize]; + __shared__ float grid_norm[kEntries]; + __shared__ float warp_best[8 * kLocalScales]; + __shared__ float group_error[kLocalScales]; + __shared__ unsigned long long warp_keys[8]; + __shared__ int selected_local; + __shared__ uint8_t locals[kGroups]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + 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 int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * kPayloadBytes; + const uint16_t d_bits = static_cast(scale_bits[block]); + const float d = __half2float(__ushort_as_half(d_bits)); + if (tid == 0) { + payload[0] = static_cast(d_bits); + payload[1] = static_cast(d_bits >> 8); + } + +#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 < 2; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * 16 + vector * 8; +#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) * 0.125f; + local_best[local] = + fminf(local_best[local], quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + float value = local_best[local]; +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + value = fminf(value, __shfl_down_sync(0xffffffff, value, delta)); + if (lane == 0) + warp_best[warp * kLocalScales + local] = value; + } + __syncthreads(); + if (tid < kLocalScales) { + float value = warp_best[tid]; +#pragma unroll + for (int w = 1; w < 8; ++w) + value = fminf(value, warp_best[w * kLocalScales + tid]); + group_error[tid] += value; + } + __syncthreads(); + } + + 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) * 0.125f; + +#pragma unroll + for (int vector = 0; vector < 2; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * 16 + vector * 8; +#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 = + quant_error(xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), + grid_norm[entry], selected_scale); + const unsigned long long candidate = + (static_cast(__float_as_uint(error)) << 32) | + static_cast(entry); + key = candidate < key ? candidate : key; + } +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + const auto other = __shfl_down_sync(0xffffffff, key, delta); + key = other < key ? other : key; + } + if (lane == 0) + warp_keys[warp] = key; + __syncthreads(); + if (tid == 0) { + key = warp_keys[0]; +#pragma unroll + for (int w = 1; w < 8; ++w) + key = warp_keys[w] < key ? warp_keys[w] : key; + const int entry = static_cast(key & 0x1ff); + 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 uint16_t code = static_cast(entry | ((sign_mask & 0x7f) << 9)); + const int code_offset = 2 + 2 * (group * 2 + vector); + payload[code_offset] = static_cast(code); + payload[code_offset + 1] = static_cast(code >> 8); + } + __syncthreads(); + } + } + + if (tid < 8) + payload[66 + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); +} + +} // namespace + +at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0, + "input size must be a positive multiple of 256"); + TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.numel() == kEntries * kVectorSize, + "grid must be float32 [512, 8]"); + TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(num_blocks <= std::numeric_limits::max(), "IQ2_XS CUDA grid is too large"); + auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong)); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + const int scale_grid = static_cast((num_blocks + 255) / 256); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] { + find_scale<<>>(input.data_ptr(), num_blocks, + scales.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + encode<<(num_blocks), 256, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + scales.data_ptr(), output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py index 87dbf30bb57..80dc83e7cb8 100644 --- a/modelopt/torch/quantization/__init__.py +++ b/modelopt/torch/quantization/__init__.py @@ -22,6 +22,7 @@ from .compress import * from .config import * from .conversion import * +from .ggml import * from .model_quant import * from .nn.modules.quant_module import QuantModuleRegistry from .utils import update_quant_cfg_with_kv_cache_quant diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index a65396d64ff..367f863ba4f 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -19,10 +19,18 @@ from modelopt.torch.utils import load_cpp_extension -__all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"] +__all__ = [ + "get_cuda_ext", + "get_cuda_ext_fp8", + "get_cuda_ext_iq1_s", + "get_cuda_ext_iq2_xs", + "get_cuda_ext_mx", + "precompile", +] path = Path(__file__).parent kernels_gemm = path.parent / "kernels" / "quantization" / "gemm" +kernels_ggml = path.parent / "kernels" / "quantization" / "ggml" def get_cuda_ext(raise_if_failed: bool = False): @@ -72,6 +80,34 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): return get_cuda_ext_mx.extension # type:ignore[attr-defined] +def get_cuda_ext_iq1_s(raise_if_failed: bool = False): + """Return the GGML-compatible IQ1_S packing extension.""" + if not hasattr(get_cuda_ext_iq1_s, "extension"): + get_cuda_ext_iq1_s.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_iq1_s", + sources=[kernels_ggml / "iq1_s.cpp", kernels_ggml / "iq1_s.cu"], + cuda_version_specifiers=">=11.8", + fail_msg="IQ1_S CUDA packing is unavailable; using the PyTorch reference encoder.", + extra_cuda_cflags=["-O3", "--use_fast_math"], + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_iq1_s.extension # type:ignore[attr-defined] + + +def get_cuda_ext_iq2_xs(raise_if_failed: bool = False): + """Return the GGML-compatible IQ2_XS packing extension.""" + if not hasattr(get_cuda_ext_iq2_xs, "extension"): + get_cuda_ext_iq2_xs.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_iq2_xs", + sources=[kernels_ggml / "iq2_xs.cpp", kernels_ggml / "iq2_xs.cu"], + cuda_version_specifiers=">=11.8", + fail_msg="IQ2_XS CUDA packing is unavailable; using the PyTorch reference encoder.", + extra_cuda_cflags=["-O3", "--use_fast_math"], + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_iq2_xs.extension # type:ignore[attr-defined] + + def __getattr__(name): if name == "cuda_ext": return get_cuda_ext() @@ -79,6 +115,10 @@ def __getattr__(name): return get_cuda_ext_fp8() elif name == "cuda_ext_mx": return get_cuda_ext_mx() + elif name == "cuda_ext_iq1_s": + return get_cuda_ext_iq1_s() + elif name == "cuda_ext_iq2_xs": + return get_cuda_ext_iq2_xs() else: raise AttributeError(f"module {__name__} has no attribute {name}") @@ -88,3 +128,5 @@ def precompile(): print(get_cuda_ext()) print(get_cuda_ext_fp8()) print(get_cuda_ext_mx()) + print(get_cuda_ext_iq1_s()) + print(get_cuda_ext_iq2_xs()) diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py new file mode 100644 index 00000000000..97d8863ce3f --- /dev/null +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GGML-compatible block quantization formats.""" + +# Importing the backend installs its TensorQuantizer dispatch entry. +from . import backend as _backend +from .iq1_s import * +from .iq2_xs import * diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py new file mode 100644 index 00000000000..bba15d64ccd --- /dev/null +++ b/modelopt/torch/quantization/ggml/backend.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TensorQuantizer backend dispatch for GGML-compatible IQ formats.""" + +import torch + +from ..nn.modules.tensor_quantizer import register_quant_backend +from .iq1_s import iq1_s_fake_quant +from .iq2_xs import iq2_xs_fake_quant + + +def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """Dispatch an IQ quantizer to its format-specific implementation.""" + num_bits = getattr(quantizer, "num_bits", None) + if num_bits == "iq1_s": + return iq1_s_fake_quant(inputs, quantizer) + if num_bits == "iq2_xs": + return iq2_xs_fake_quant(inputs, quantizer) + raise ValueError("The psx_luts backend requires num_bits='iq1_s' or 'iq2_xs'") + + +register_quant_backend("psx_luts", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py new file mode 100644 index 00000000000..ff977813ac1 --- /dev/null +++ b/modelopt/torch/quantization/ggml/common.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared validation for GGML-compatible block quantizers.""" + +import math + +import torch + +GGML_BLOCK_SIZE = 256 + + +def validate_weight(weight: torch.Tensor, format_name: str) -> None: + """Validate a weight accepted by the current GGML block encoders.""" + if weight.numel() == 0: + raise ValueError(f"{format_name} requires a non-empty weight") + if weight.dim() == 0 or weight.shape[-1] % GGML_BLOCK_SIZE: + raise ValueError( + f"{format_name} requires the last weight dimension to be divisible by " + f"{GGML_BLOCK_SIZE}, got shape {tuple(weight.shape)}" + ) + if not weight.is_floating_point(): + raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError(f"{format_name} requires finite weight values") + + +def validate_packed_weights( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + block_bytes: int, + format_name: str, +) -> tuple[int, ...]: + """Validate a packed payload and return its logical shape.""" + if packed_weights.dtype != torch.uint8 or packed_weights.shape[-1] != block_bytes: + raise ValueError( + f"packed_weights must be uint8 with last dimension {block_bytes}, " + f"got {packed_weights.dtype} {tuple(packed_weights.shape)}" + ) + shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist()) + if not shape or shape[-1] % GGML_BLOCK_SIZE: + raise ValueError(f"invalid {format_name} logical weight shape: {shape}") + expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes + if packed_weights.numel() != expected_payload_values: + raise ValueError("packed_weights size does not match weight_shape") + return shape diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py new file mode 100644 index 00000000000..65307f6c9b8 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -0,0 +1,304 @@ +# This file includes the IQ1_S codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""IQ1_S fake quantization and GGML-compatible block packing. + +The encoder follows the PSX-LUTS ``search_impl="auto"`` search. Every 256 +logical values become one 50-byte ``block_iq1_s`` payload: + +* bytes 0..1: little-endian FP16 super-block scale ``d`` +* bytes 2..33: low eight bits of 32 codebook indices +* bytes 34..49: eight little-endian uint16 metadata words + +Each metadata word describes four consecutive eight-value vectors. Bits 0..11 +hold the three high index bits, bits 12..14 select one of eight local scales, +and bit 15 selects the shared -0.125 rather than +0.125 delta. The canonical +2048 x 8 ternary grid below comes from llama.cpp ``ggml-common.h`` revision +9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +import zlib +from functools import cache + +import torch + +from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight + +__all__ = [ + "IQ1_S_BLOCK_BYTES", + "IQ1_S_BLOCK_SIZE", + "IQ1_S_EFFECTIVE_BITS", + "dequantize_iq1_s", + "iq1_s_fake_quant", + "iq1_s_grid", + "quantize_iq1_s", +] + +IQ1_S_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ1_S_BLOCK_BYTES = 50 +IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE +_IQ1_S_DELTA = 0.125 +_IQ1_S_NATIVE_MAX = 16.875 + +# zlib-compressed little-endian bytes of the canonical uint64_t table. The +# decoded int8 values are -1, 0, and 1. +_IQ1_S_GRID_ZLIB_B64 = ( + "eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv" + "c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY" + "4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX" + "bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8" + "Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W" + "YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH" + "9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL" + "2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E" + "UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg" + "MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0" + "jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+" + "LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs" + "fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2" + "a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g" + "n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq" + "ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ" + "q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/" + "1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2" + "kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM" + "T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8" + "0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT" + "oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy" + "i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44" + "rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg" + "O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO" + "k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI" + "Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8" + "73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn" + "gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6" + "hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+" + "EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA" + "0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH" + "OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD" + "FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn" + "h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6" + "DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB" + "MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI" + "4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh" + "eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ" + "zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep" + "V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5" + "/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C" + "qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T" + "rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD" + "hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr" + "x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO" + "umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0" + "Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA" + "MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ" + "jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8=" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64)) + + +def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ1_S ternary grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + raw = torch.tensor(list(_grid_bytes()), dtype=torch.uint8).view(torch.int8) + _GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to( + device=resolved_device, dtype=torch.float32 + ) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + xnorm = vectors.square().sum(dim=-1) + xsum = vectors.sum(dim=-1) + + amax = x.abs().amax(dim=1) + d = ((amax / _IQ1_S_NATIVE_MAX) * 0.61).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + best_error = torch.full((block_count, 32, 16), torch.inf, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + grid_norm = grid.square().sum(dim=-1) + grid_sum = grid.sum(dim=-1) + + # Tile the 2048-entry codebook to bound temporary memory. A strict update + # retains the lowest codebook index when two candidates have equal error. + for entry_start in range(0, 2048, 128): + grid_tile = grid[entry_start : entry_start + 128] + dot = torch.matmul(vectors, grid_tile.T) + tile_norm = grid_norm[entry_start : entry_start + 128].reshape(1, 1, -1) + tile_sum = grid_sum[entry_start : entry_start + 128].reshape(1, 1, -1) + + for shift in range(2): + delta = -_IQ1_S_DELTA if shift else _IQ1_S_DELTA + shifted_dot = dot + delta * xsum.unsqueeze(-1) + shifted_norm = tile_norm + 2 * delta * tile_sum + 8 * delta * delta + for local in range(8): + choice = shift * 8 + local + scale = d_float.reshape(-1, 1, 1) * (2 * local + 1) + error = ( + xnorm.unsqueeze(-1) - 2 * scale * shifted_dot + scale.square() * shifted_norm + ) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, choice] + best_error[:, :, choice] = torch.where( + replace, tile_error, best_error[:, :, choice] + ) + best_entry[:, :, choice] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, choice] + ) + + group_error = best_error.reshape(block_count, 8, 4, 16).sum(dim=2) + selected_choice = group_error.argmin(dim=-1) + vector_choice = selected_choice.repeat_interleave(4, dim=1) + selected_entry = best_entry.gather(2, vector_choice.unsqueeze(-1)).squeeze(-1) + selected_local = selected_choice & 0x7 + selected_shift = selected_choice >> 3 + + high = (selected_entry >> 8).reshape(block_count, 8, 4) + qh = ( + high[:, :, 0] + | (high[:, :, 1] << 3) + | (high[:, :, 2] << 6) + | (high[:, :, 3] << 9) + | (selected_local << 12) + | (selected_shift << 15) + ) + + packed = torch.empty((block_count, IQ1_S_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:34] = (selected_entry & 0xFF).to(torch.uint8) + packed[:, 34:50:2] = (qh & 0xFF).to(torch.uint8) + packed[:, 35:50:2] = (qh >> 8).to(torch.uint8) + packed[d_float == 0] = 0 + return packed + + +@torch.no_grad() +def quantize_iq1_s( + weight: torch.Tensor, *, block_chunk_size: int = 4 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ1_S blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 50]`` + and ``[weight.ndim]``. Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ1_S") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + blocks = weight.contiguous().reshape(-1, IQ1_S_BLOCK_SIZE) + grid = iq1_s_grid(weight.device) + if weight.is_cuda: + from ..extensions import get_cuda_ext_iq1_s + + extension = get_cuda_ext_iq1_s() + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ1_S_BLOCK_SIZE, + IQ1_S_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ1_S_BLOCK_SIZE, + IQ1_S_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq1_s( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ1_S payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ1_S_BLOCK_BYTES, format_name="IQ1_S" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = blocks[:, 2:34].to(torch.int64).reshape(-1, 8, 4) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) + high = (qh.unsqueeze(-1) >> shifts) & 0x7 + entries = low | (high << 8) + + local = (qh >> 12) & 0x7 + delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA) + values = iq1_s_grid(blocks.device)[entries] + delta.unsqueeze(-1).unsqueeze(-1) + scales = d.unsqueeze(-1) * (2 * local + 1).float() + decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + return decoded.reshape(shape).to(dtype) + + +def iq1_s_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ1_S backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq1_s": + raise ValueError("The psx_luts IQ1_S backend requires num_bits='iq1_s'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ1_S search_impl='auto' is currently supported") + packed, shape = quantize_iq1_s(inputs) + reconstructed = dequantize_iq1_s(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py new file mode 100644 index 00000000000..ec0816922cc --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -0,0 +1,301 @@ +# This file includes the IQ2_XS codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""IQ2_XS fake quantization and GGML-compatible block packing. + +The encoder follows the PSX-LUTS search_impl="auto" search. Every 256 +logical values become one 74-byte block_iq2_xs payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign) +* bytes 66..73: 16 four-bit local scales, two per byte + +The canonical 512 x 8 magnitude grid below comes from llama.cpp +ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +from functools import cache + +import torch + +from .common import GGML_BLOCK_SIZE, validate_packed_weights, validate_weight + +__all__ = [ + "IQ2_XS_BLOCK_BYTES", + "IQ2_XS_BLOCK_SIZE", + "IQ2_XS_EFFECTIVE_BITS", + "dequantize_iq2_xs", + "iq2_xs_fake_quant", + "iq2_xs_grid", + "quantize_iq2_xs", +] + +IQ2_XS_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ2_XS_BLOCK_BYTES = 74 +IQ2_XS_EFFECTIVE_BITS = IQ2_XS_BLOCK_BYTES * 8 / IQ2_XS_BLOCK_SIZE + +# Compact byte representation of the canonical [512, 8] grid. Values are only +# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids +# adding a pickle-backed torch.save artifact to the wheel. +_IQ2_XS_GRID_B64 = ( + "CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr" + "CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI" + "GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI" + "GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr" + "CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr" + "CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI" + "CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI" + "KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ" + "CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI" + "CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI" + "CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ" + "GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI" + "CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI" + "KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ" + "CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI" + "GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI" + "GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI" + "GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ" + "CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI" + "CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI" + "KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI" + "GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr" + "KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr" + "KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI" + "CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI" + "GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ" + "GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr" + "CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ" + "CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI" + "CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr" + "CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI" + "CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI" + "CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ" + "GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI" + "CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr" + "GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ" + "GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI" + "CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ" + "KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI" + "KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr" + "CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI" + "KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI" + "GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr" + "GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr" + "CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI" + "KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ" + "GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI" + "CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ" + "GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr" + "KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ" + "CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw==" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return base64.b64decode(_IQ2_XS_GRID_B64) + + +def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_XS magnitude grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + values = torch.tensor(list(_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + amax = x.abs().amax(dim=1) + rms = x.square().mean(dim=1).sqrt() + peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) + anchor_ratio = (1.0 - 0.035 * peak_to_rms).clamp(0.65, 0.92) + d = ((amax / 166.625) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + # Search the codebook in tiles to cap temporary memory. Strict comparison + # preserves the lowest grid index on equal error, matching the CUDA key. + for entry_start in range(0, 512, 64): + grid_tile = grid[entry_start : entry_start + 64] + products = magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8) + dot = products.sum(dim=-1) + dot = torch.where(odd_parity.unsqueeze(-1), dot - 2.0 * products.amin(dim=-1), dot) + tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1) + + for local in range(16): + scale = d_float.reshape(-1, 1, 1) * ((2 * local + 1) / 8.0) + error = ( + xnorm.unsqueeze(-1) - 2.0 * scale * dot + scale.square() * tile_qnorm + ).clamp_min_(0) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, local] + best_error[:, :, local] = torch.where(replace, tile_error, best_error[:, :, local]) + best_entry[:, :, local] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, local] + ) + + group_error = best_error.reshape(block_count, 16, 2, 16).sum(dim=2) + selected_local = group_error.argmin(dim=-1) + vector_local = selected_local.repeat_interleave(2, dim=1) + selected_entry = best_entry.gather(2, vector_local.unsqueeze(-1)).squeeze(-1) + + selected_grid = grid[selected_entry] + weakest_index = (magnitudes * selected_grid).argmin(dim=-1) + flip = torch.nn.functional.one_hot(weakest_index, num_classes=8).bool() + encoded_negative = negative ^ (flip & odd_parity.unsqueeze(-1)) + sign_bits = torch.arange(8, dtype=torch.int64, device=x.device) + sign_mask = (encoded_negative.to(torch.int64) << sign_bits).sum(dim=-1) + + codes = selected_entry | ((sign_mask & 0x7F) << 9) + packed = torch.empty((block_count, IQ2_XS_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:66:2] = (codes & 0xFF).to(torch.uint8) + packed[:, 3:66:2] = (codes >> 8).to(torch.uint8) + packed[:, 66:] = (selected_local[:, 0::2] | (selected_local[:, 1::2] << 4)).to(torch.uint8) + return packed + + +@torch.no_grad() +def quantize_iq2_xs( + weight: torch.Tensor, *, block_chunk_size: int = 64 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ2_XS blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 74]`` + and ``[weight.ndim]``. Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ2_XS") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + blocks = weight.contiguous().reshape(-1, IQ2_XS_BLOCK_SIZE) + grid = iq2_xs_grid(weight.device) + if weight.is_cuda: + from ..extensions import get_cuda_ext_iq2_xs + + extension = get_cuda_ext_iq2_xs() + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XS_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XS_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_xs( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ2_XS payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_XS_BLOCK_BYTES, format_name="IQ2_XS" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + entries = codes & 0x1FF + sign_index = codes >> 9 + + parity = torch.zeros_like(sign_index) + for bit in range(7): + parity ^= (sign_index >> bit) & 1 + sign_mask = sign_index | (parity << 7) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + + scale_bytes = blocks[:, 66:].to(torch.int64) + local = torch.empty((blocks.shape[0], 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 + values = iq2_xs_grid(blocks.device)[entries] * signs + decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) + return decoded.reshape(shape).to(dtype) + + +def iq2_xs_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ2_XS backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_xs": + raise ValueError("The psx_luts IQ2_XS backend requires num_bits='iq2_xs'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ2_XS search_impl='auto' is currently supported") + packed, shape = quantize_iq2_xs(inputs) + reconstructed = dequantize_iq2_xs(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() diff --git a/modelopt_recipes/configs/numerics/iq1_s.yaml b/modelopt_recipes/configs/numerics/iq1_s.yaml new file mode 100644 index 00000000000..507f4588f7e --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq1_s.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ1_S weight quantizer using the built-in PSX-LUTS-compatible search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq1_s +effective_bits: 1.5625 +block_sizes: + -1: 256 +backend: psx_luts +backend_extra_args: + search_impl: auto +pass_through_bwd: true diff --git a/modelopt_recipes/configs/numerics/iq2_xs.yaml b/modelopt_recipes/configs/numerics/iq2_xs.yaml new file mode 100644 index 00000000000..4fe6f954c5c --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_xs.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ2_XS weight quantizer using the built-in PSX-LUTS-compatible search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_xs +effective_bits: 2.3125 +block_sizes: + -1: 256 +backend: psx_luts +backend_extra_args: + search_impl: auto +pass_through_bwd: true diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml new file mode 100644 index 00000000000..31fab0f7a44 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# QuantizeConfig preset for IQ1_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 + iq1_s: configs/numerics/iq1_s + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq1_s + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml new file mode 100644 index 00000000000..96938c79f09 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# QuantizeConfig preset for IQ2_XS 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_xs: configs/numerics/iq2_xs + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml new file mode 100644 index 00000000000..06c328a8da8 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ1_S weight-only PTQ and unified-checkpoint export. + +imports: + preset: configs/ptq/presets/model/iq1_s + +metadata: + recipe_type: ptq + description: >- + Applies GGML-compatible IQ1_S weight-only quantization. No calibration data is required; + unified export writes the GGML block payload to weight for every quantized Linear weight. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml new file mode 100644 index 00000000000..1131a5a29c7 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ2_XS weight-only PTQ and unified-checkpoint export. + +imports: + preset: configs/ptq/presets/model/iq2_xs + +metadata: + recipe_type: ptq + description: >- + Applies GGML-compatible IQ2_XS weight-only quantization. No calibration data is required; + unified export writes the GGML block payload to weight for every quantized Linear weight. +quantize: + $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 2da6a0ff2b0..b334eb4ab2a 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 25 general/ptq/ recipes (click to expand) +All 27 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -57,6 +57,8 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `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, all linears | none | PSX-LUTS auto search (no calibration) | +| `iq2_xs` | IQ2_XS W2A16, all linears | none | PSX-LUTS auto search (no calibration) |
@@ -135,6 +137,10 @@ activations and tensor-core math are what deliver the throughput. - **`mxfp4_mlp_weight_only`** — MXFP4 weights on MLP/MoE layers only, BF16 activations. Needs no calibration forward pass; the QAT starting point for the GPT-OSS family (see `examples/gpt-oss`). +- **`iq1_s` / `iq2_xs`** — GGML-compatible IQ1_S or IQ2_XS weights on all linear + layers, with BF16 activations. Unified export stores each logical weight as a packed + 50- or 74-byte-per-256-values payload plus its original shape. No calibration data is + required. --- diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py new file mode 100644 index 00000000000..b6ebb2aada9 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, quantize_iq1_s + + +def test_iq1_s_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + packed_again, _ = quantize_iq1_s(weight) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (8, 1, 50) + assert torch.equal(packed, packed_again) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed, shape = quantize_iq1_s(weight) + + assert not packed.any() + assert torch.equal(dequantize_iq1_s(packed, shape), weight) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py new file mode 100644 index 00000000000..ae06c931b79 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, quantize_iq2_xs + + +def test_iq2_xs_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + packed_again, _ = quantize_iq2_xs(weight) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (8, 2, 74) + assert torch.equal(packed, packed_again) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + +def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed, shape = quantize_iq2_xs(weight) + + assert not packed.any() + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 2da25e90e15..daea55be1ea 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -46,6 +46,9 @@ import modelopt.torch.speculative as mtsp from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf from modelopt.torch.export.unified_export_megatron import GPTModelExporter +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_iq2_xs +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.speculative.eagle.default_config import default_eagle_config from modelopt.torch.speculative.plugins.megatron_eagle import _DynamicEagleGPTModel from modelopt.torch.speculative.plugins.megatron_medusa import _DynamicMedusaGPTModel @@ -86,6 +89,75 @@ def _verify_model_quant_config( assert quant_config_dict["kv_cache_quant_algo"] == KV_CACHE_FP8 +@pytest.mark.parametrize( + ("qformat", "payload_bytes", "dequantize"), + [("iq1_s", 50, dequantize_iq1_s), ("iq2_xs", 74, dequantize_iq2_xs)], +) +def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequantize): + """Megatron export writes the same scale-free IQ representation as HF export.""" + linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=qformat, + block_sizes={-1: 256}, + backend="psx_luts", + backend_extra_args={"search_impl": "auto"}, + ) + ) + exporter = object.__new__(GPTModelExporter) + exporter.dtype = torch.bfloat16 + exporter._state_dict = {} + exporter.exclude_modules = [] + exporter.layer_config_dict = {} + + exporter._name_remapping(linear, "model.layers.0.mlp.down_proj.") + + packed_key = "model.layers.0.mlp.down_proj.weight" + assert exporter._state_dict[packed_key].shape == (2, 1, payload_bytes) + assert exporter._state_dict[packed_key].dtype == torch.uint8 + logical_shape = torch.tensor( + [ + *exporter._state_dict[packed_key].shape[:-2], + exporter._state_dict[packed_key].shape[-2] * 256, + ] + ) + reconstructed = dequantize( + exporter._state_dict[packed_key], + logical_shape, + dtype=torch.bfloat16, + ) + torch.testing.assert_close(reconstructed, linear.weight_quantizer(linear.weight)) + assert exporter.layer_config_dict == { + "model.layers.0.mlp.down_proj.quantization": qformat, + "model.layers.0.mlp.down_proj.awq_block_size": 256, + } + + +def test_megatron_iq_export_rejects_tensor_parallelism(): + """IQ packing is intentionally limited to complete TP=1 weights.""" + linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="psx_luts", + backend_extra_args={"search_impl": "auto"}, + ) + ) + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Sequential(linear) + + with ( + patch.object(exporter, "_is_sidecar_writer_rank", return_value=False), + patch.object(uem, "get_pipeline_model_parallel_rank", return_value=0), + patch.object(uem, "get_pipeline_model_parallel_world_size", return_value=1), + patch.object(uem, "get_tensor_model_parallel_rank", return_value=0), + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2), + pytest.raises(NotImplementedError, match="tensor model parallel size 1"), + ): + exporter.save_pretrained("unused", "unused") + + def _test_unified_export_megatron( tmp_path, model_type, diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..cae13977db5 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -20,10 +20,13 @@ from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import postprocess_state_dict from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _process_quantized_modules, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -102,6 +105,27 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) +@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)]) +def test_export_iq_payload_as_weight(num_bits, payload_bytes): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=num_bits, + block_sizes={-1: 256}, + backend="psx_luts", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + _export_quantized_weight(linear, torch.bfloat16) + state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) + + assert state_dict["weight"].shape == (4, 1, payload_bytes) + assert state_dict["weight"].dtype == torch.uint8 + assert "packed_weights" not in state_dict + assert "weight_shape" not in state_dict + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 90740587bf7..8c977039223 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -26,8 +26,11 @@ import modelopt.torch.export.unified_export_megatron as unified_export_megatron import modelopt.torch.quantization as mtq +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.quant_format import ( QUANTIZATION_FP8, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_NVFP4, QUANTIZATION_W4A8_AWQ, ) @@ -39,6 +42,45 @@ from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +@pytest.mark.parametrize( + ("num_bits", "quantization_format", "payload_bytes", "effective_bits"), + [ + ("iq1_s", QUANTIZATION_IQ1_S, 50, 1.5625), + ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2.3125), + ], +) +def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, effective_bits): + model = torch.nn.Sequential(torch.nn.Linear(256, 256, bias=False)) + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": num_bits, + "block_sizes": {-1: 256}, + "backend": "psx_luts", + "backend_extra_args": {"search_impl": "auto"}, + }, + }, + ], + "algorithm": None, + }, + ) + + assert get_quantization_format(model) == quantization_format + config = get_quant_config(model) + assert config["quantization"]["quant_algo"] == num_bits.upper() + assert config["quantization"]["block_payload_bytes"] == payload_bytes + hf_config = convert_hf_quant_config_format(config) + weights = hf_config["config_groups"]["group_0"]["weights"] + assert weights["group_size"] == 256 + assert weights["effective_bits"] == effective_bits + assert weights["packing"] == "ggml" + + class _FakeKVCacheQuantizer(torch.nn.Module): """Minimal FP8 KV cache quantizer for scaling-factor tests.""" diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 49cd999a205..f6fe438a2d3 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -502,6 +502,37 @@ def forward_loop(m): self._cleanup_registry(expert_type) + def test_export_registers_packed_weight_buffers(self, monkeypatch): + """Packed expert weights must remain present in the exported state dict.""" + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + + try: + converted = QuantModuleRegistry.convert(model.moe.experts) + + def _pack_weight_as_buffer(wrapper, dtype): + packed = torch.zeros((*wrapper.weight.shape, 1), dtype=torch.uint8) + del wrapper.weight + wrapper.register_buffer("weight", packed) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _pack_weight_as_buffer, + ) + + _export_fused_experts(converted, torch.float16) + + state_dict = converted.state_dict() + for idx in range(NUM_EXPERTS): + for projection in ("gate_proj", "up_proj", "down_proj"): + key = f"{idx}.{projection}.weight" + assert key in state_dict + assert state_dict[key].dtype == torch.uint8 + finally: + self._cleanup_registry(expert_type) + def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch): """gate_proj and up_proj must share weight_scale_2 even when an expert was never routed during calibration. diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py new file mode 100644 index 00000000000..a623be608a4 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.ggml.iq1_s import ( + IQ1_S_BLOCK_BYTES, + dequantize_iq1_s, + iq1_s_fake_quant, + iq1_s_grid, + quantize_iq1_s, +) + + +def test_iq1_s_canonical_grid(): + grid = iq1_s_grid() + + assert grid.shape == (2048, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0} + assert grid[0].tolist() == [-1.0] * 8 + + +def test_iq1_s_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + + assert packed.shape == (2, 1, IQ1_S_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq1_s(packed, shape), weight) + + +def test_iq1_s_dequantizes_ggml_metadata_bit_fields(): + packed = torch.zeros((1, 1, 50), dtype=torch.uint8) + d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8) + packed[0, 0, :2] = d + entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64) + packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8) + qh = ( + ((entries[0] >> 8) & 7) + | (((entries[1] >> 8) & 7) << 3) + | (((entries[2] >> 8) & 7) << 6) + | (((entries[3] >> 8) & 7) << 9) + | (3 << 12) + | (1 << 15) + ) + packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8) + packed[0, 0, 35] = (qh >> 8).to(torch.uint8) + + decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32) + expected = (iq1_s_grid()[entries] - 0.125) * 14.0 + + assert torch.equal(decoded[0, :32].reshape(4, 8), expected) + + +def test_iq1_s_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight, block_chunk_size=1) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (2, 1, 50) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + blocks = packed.reshape(-1, 50) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + assert torch.all(((qh >> 12) & 0x7) < 8) + assert torch.all((qh & 0xFFF) < 0x1000) + + +def test_iq1_s_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="last weight dimension"): + quantize_iq1_s(torch.ones(2, 257)) + + +def test_iq1_s_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq1_s" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq1_s_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py new file mode 100644 index 00000000000..d3e491a2f24 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.quantization.ggml.iq2_xs import ( + IQ2_XS_BLOCK_BYTES, + dequantize_iq2_xs, + iq2_xs_fake_quant, + iq2_xs_grid, + quantize_iq2_xs, +) + + +def test_iq2_xs_canonical_grid(): + grid = iq2_xs_grid() + + assert grid.shape == (512, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {8.0, 25.0, 43.0} + assert grid[0].tolist() == [8.0] * 8 + assert grid[-1].tolist() == [43.0] * 8 + + +def test_iq2_xs_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + + assert packed.shape == (2, 1, IQ2_XS_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) + + +def test_iq2_xs_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight, block_chunk_size=2) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (2, 2, 74) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + blocks = packed.reshape(-1, 74) + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + assert torch.all((codes & 0x1FF) < 512) + assert torch.all((codes >> 9) < 128) + + +def test_iq2_xs_requires_complete_last_dimension_blocks(): + with pytest.raises(ValueError, match="last weight dimension"): + quantize_iq2_xs(torch.ones(2, 257)) + + +def test_iq2_xs_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq2_xs" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq2_xs_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight))