[2/4] Register each GGML IQ format once for dispatch and export

modelopt/torch/export/convert_hf_config.py

@@ -19,7 +19,9 @@
 from collections import defaultdict
 from typing import Any
 
-from .quant_format import IQ_BLOCK_METADATA, IQ_FORMATS
+from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY
+
+from .quant_format import IQ_FORMATS
 
 
 def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) -> dict[str, Any]:
@@ -121,7 +123,9 @@ 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.lower() in IQ_FORMATS:
-        block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[quant_algo.lower()]
+        iq_format = IQ_FORMAT_REGISTRY[quant_algo.lower()]
+        block_size, payload_bytes = iq_format.block_size, iq_format.block_bytes
+        effective_bits = iq_format.effective_bits
         if group_size not in (None, block_size):
             raise ValueError(f"{quant_algo} requires group size {block_size}, got {group_size}")
         # IQ payloads are self-contained blocks, not compressed-tensors integer groups.

modelopt/torch/export/quant_format.py

@@ -19,20 +19,7 @@
 constants, for example, are in :mod:`modelopt.torch.export.trtllm.model_config`.
 """
 
-from modelopt.torch.quantization.ggml import (
-    IQ1_S_BLOCK_BYTES,
-    IQ1_S_BLOCK_SIZE,
-    IQ1_S_EFFECTIVE_BITS,
-    IQ2_XS_BLOCK_BYTES,
-    IQ2_XS_BLOCK_SIZE,
-    IQ2_XS_EFFECTIVE_BITS,
-    IQ2_XXS_BLOCK_BYTES,
-    IQ2_XXS_BLOCK_SIZE,
-    IQ2_XXS_EFFECTIVE_BITS,
-    quantize_iq1_s,
-    quantize_iq2_xs,
-    quantize_iq2_xxs,
-)
+from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY
 
 QUANTIZATION_NONE = None
 QUANTIZATION_FP8 = "fp8"
@@ -55,34 +42,15 @@
 QUANTIZATION_IQ2_XXS = "iq2_xxs"
 QUANTIZATION_IQ2_XS = "iq2_xs"
 
-# Every GGML IQ format. They share the weight-only, 256-value-block, per-module-scale
-# shape, so export treats them as one family; adding a format means adding it here
-# rather than extending a tuple at each use site.
-IQ_FORMATS = frozenset(
-    {
-        QUANTIZATION_IQ1_S,
-        QUANTIZATION_IQ2_XXS,
-        QUANTIZATION_IQ2_XS,
-    }
-)
-
-# Block geometry per IQ format: (block size, packed bytes per block, bits per weight). Checkpoint
-# metadata spells the algorithm in upper case, so consumers look up
-# ``IQ_BLOCK_METADATA[algo.lower()]`` rather than carrying a second spelling of the family.
-IQ_BLOCK_METADATA = {
-    QUANTIZATION_IQ1_S: (IQ1_S_BLOCK_SIZE, IQ1_S_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS),
-    QUANTIZATION_IQ2_XXS: (IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, IQ2_XXS_EFFECTIVE_BITS),
-    QUANTIZATION_IQ2_XS: (IQ2_XS_BLOCK_SIZE, IQ2_XS_BLOCK_BYTES, IQ2_XS_EFFECTIVE_BITS),
-}
-
-
-# The packer each format's checkpoint weights are written with. Both exporters resolve through
-# this one mapping so they cannot drift apart.
-IQ_PACKERS = {
-    QUANTIZATION_IQ1_S: quantize_iq1_s,
-    QUANTIZATION_IQ2_XXS: quantize_iq2_xxs,
-    QUANTIZATION_IQ2_XS: quantize_iq2_xs,
-}
+# Every GGML IQ format, derived from the registry the quantization backend dispatches through, so
+# export and dispatch cannot disagree about which formats exist. They share the weight-only,
+# 256-value-block, per-module-scale shape, so export treats them as one family. A format's block
+# geometry and packer are read from IQ_FORMAT_REGISTRY directly.
+#
+# Registering a format therefore declares it exportable, and that is intended rather than a side
+# effect: fake quant is dequantize(quantize(w)), so a format cannot be dispatched without the
+# packer and block geometry that are all export reads.
+IQ_FORMATS = frozenset(IQ_FORMAT_REGISTRY)
 
 
 # Formats whose scales are purely per-module, so export never merges them across the q/k/v

modelopt/torch/export/quant_utils.py

@@ -27,6 +27,7 @@
 
 from modelopt import __version__
 from modelopt.torch.models import get_spec, list_all_possible
+from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY
 from modelopt.torch.quantization.model_calib import (
     enable_stats_collection,
     finish_stats_collection,
@@ -51,7 +52,6 @@
 from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer
 from .model_utils import TiedWeightMap, get_language_model_from_vl
 from .quant_format import (
-    IQ_BLOCK_METADATA,
     IQ_FORMATS,
     KV_CACHE_FP8,
     KV_CACHE_FP8_K_NVFP4_V,
@@ -773,7 +773,9 @@ def process_layer_quant_config(layer_config_dict):
                 "group_size": block_size_value,
             }
         elif v in IQ_FORMATS:
-            block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[v]
+            iq_format = IQ_FORMAT_REGISTRY[v]
+            block_size, payload_bytes = iq_format.block_size, iq_format.block_bytes
+            effective_bits = iq_format.effective_bits
             if block_size_value != block_size:
                 raise ValueError(
                     f"{v.upper()} requires block size {block_size}, got {block_size_value}"

modelopt/torch/export/unified_export_hf.py

@@ -67,6 +67,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 IQ_FORMAT_REGISTRY
 from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
 from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper
@@ -101,7 +102,6 @@
 from .quant_format import (
     FUSION_FREE_FORMATS,
     IQ_FORMATS,
-    IQ_PACKERS,
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
     QUANTIZATION_FP8_PC_PT,
@@ -635,7 +635,7 @@ def _export_quantized_weight(
                 "IQ unified export currently supports modules with a standard 'weight' "
                 f"attribute, got {weight_name!r} on {type(sub_module).__name__}"
             )
-        quantize_iq = IQ_PACKERS[quantization_format]
+        quantize_iq = IQ_FORMAT_REGISTRY[quantization_format].quantize
         packed_weight, _ = quantize_iq(weight.to(dtype))
         setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False))
         maybe_clear_cuda_cache()

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 IQ_FORMAT_REGISTRY
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer
 from modelopt.torch.utils import import_plugin, warn_rank_0
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
@@ -57,7 +58,6 @@
 from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d
 from .quant_format import (
     IQ_FORMATS,
-    IQ_PACKERS,
     KV_CACHE_FP8,
     KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
@@ -1185,7 +1185,7 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
     @staticmethod
     def _pack_iq_weight(weight: torch.Tensor, qformat: str) -> torch.Tensor:
         """Pack one ``[out, in]`` weight and return its CPU payload."""
-        quantize_iq = IQ_PACKERS[qformat]
+        quantize_iq = IQ_FORMAT_REGISTRY[qformat].quantize
         packed_weight, _ = quantize_iq(weight)
         return packed_weight.detach().cpu()
 

modelopt/torch/quantization/ggml/__init__.py

@@ -23,9 +23,12 @@
 from .iq2_xs import __all__ as _iq2_xs_all
 from .iq2_xxs import *
 from .iq2_xxs import __all__ as _iq2_xxs_all
+from .registry import IQ_FORMAT_REGISTRY, IQFormat
 
 __all__ = [  # noqa: PLE0604
     *_iq1_s_all,
     *_iq2_xs_all,
     *_iq2_xxs_all,
+    "IQ_FORMAT_REGISTRY",
+    "IQFormat",
 ]

modelopt/torch/quantization/ggml/backend.py

@@ -18,16 +18,7 @@
 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
-from .iq2_xxs import iq2_xxs_fake_quant
-
-# One entry per GGML IQ format; adding a format is adding a row here.
-_FAKE_QUANTS = {
-    "iq1_s": iq1_s_fake_quant,
-    "iq2_xs": iq2_xs_fake_quant,
-    "iq2_xxs": iq2_xxs_fake_quant,
-}
+from .registry import IQ_FORMAT_REGISTRY
 
 
 def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor:
@@ -39,11 +30,11 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor:
         raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}")
     # num_bits arrives untyped from the quantizer and is a tuple for scalar formats,
     # so narrow before the lookup rather than relying on the dict to reject it.
-    fake_quant = _FAKE_QUANTS.get(num_bits) if isinstance(num_bits, str) else None
-    if fake_quant is None:
-        supported = ", ".join(repr(name) for name in sorted(_FAKE_QUANTS))
+    iq_format = IQ_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None
+    if iq_format is None:
+        supported = ", ".join(repr(name) for name in sorted(IQ_FORMAT_REGISTRY))
         raise ValueError(f"The ggml backend requires num_bits in ({supported})")
-    return fake_quant(inputs, quantizer, **extra_args)
+    return iq_format.fake_quant(inputs, quantizer, **extra_args)
 
 
 register_quant_backend("ggml", ggml_fake_quant)

modelopt/torch/quantization/ggml/common.py

@@ -121,6 +121,59 @@ def fake_quantize_with_cache(
     return inputs + (reconstructed - inputs).detach()
 
 
+@dataclass(frozen=True)
+class IQFormat:
+    """Everything backend dispatch and export need to know about one IQ format.
+
+    Each format module declares one of these beside its encoder and decoder, and
+    :data:`~modelopt.torch.quantization.ggml.registry.IQ_FORMAT_REGISTRY` lists them. The
+    per-format pieces -- codebook, search, payload layout -- stay in the format's module; what
+    lives here is the part every format does the same way.
+    """
+
+    name: str
+    block_size: int
+    block_bytes: int
+    quantize: Callable[..., tuple[torch.Tensor, torch.Tensor]]
+    dequantize: Callable[..., torch.Tensor]
+    # Encode and decode are chunked separately: packing runs once per weight and is bounded by
+    # its search temporaries, decoding runs every forward and is bounded by kernel launches.
+    block_chunk_size: int
+    decode_chunk_size: int
+
+    @property
+    def effective_bits(self) -> float:
+        """Packed storage cost per weight."""
+        return self.block_bytes * 8 / self.block_size
+
+    def fake_quant(
+        self,
+        inputs: torch.Tensor,
+        quantizer,
+        *,
+        block_chunk_size: int | None = None,
+        decode_chunk_size: int | None = None,
+    ) -> torch.Tensor:
+        """TensorQuantizer backend for this format, with pass-through backward."""
+        if getattr(quantizer, "num_bits", None) != self.name:
+            raise ValueError(
+                f"The ggml {self.name.upper()} backend requires num_bits={self.name!r}"
+            )
+        return fake_quantize_with_cache(
+            inputs,
+            quantizer,
+            format_name=self.name,
+            block_chunk_size=(
+                self.block_chunk_size if block_chunk_size is None else block_chunk_size
+            ),
+            decode_chunk_size=(
+                self.decode_chunk_size if decode_chunk_size is None else decode_chunk_size
+            ),
+            quantize=self.quantize,
+            dequantize=self.dequantize,
+        )
+
+
 def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor:
     """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does.
 

modelopt/torch/quantization/ggml/iq1_s.py

@@ -38,7 +38,7 @@
 from .codebooks import iq1_s_grid_bytes
 from .common import (
     GGML_BLOCK_SIZE,
-    fake_quantize_with_cache,
+    IQFormat,
     narrow_to_float32,
     validate_block_chunk_size,
     validate_packed_weights,
@@ -235,22 +235,17 @@ def dequantize_iq1_s(
     return decoded.reshape(shape)
 
 
-def iq1_s_fake_quant(
-    inputs: torch.Tensor,
-    quantizer,
-    *,
-    block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
-    decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
-) -> torch.Tensor:
-    """IQ1_S weight backend for TensorQuantizer, with pass-through backward."""
-    if getattr(quantizer, "num_bits", None) != "iq1_s":
-        raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'")
-    return fake_quantize_with_cache(
-        inputs,
-        quantizer,
-        format_name="iq1_s",
-        block_chunk_size=block_chunk_size,
-        decode_chunk_size=decode_chunk_size,
-        quantize=quantize_iq1_s,
-        dequantize=dequantize_iq1_s,
-    )
+IQ1_S_FORMAT = IQFormat(
+    name="iq1_s",
+    block_size=IQ1_S_BLOCK_SIZE,
+    block_bytes=IQ1_S_BLOCK_BYTES,
+    quantize=quantize_iq1_s,
+    dequantize=dequantize_iq1_s,
+    block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE,
+    decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE,
+)
+
+# Kept for callers of the per-format entry point. The record captured quantize_iq1_s and
+# dequantize_iq1_s when it was built, so patching those module functions changes neither backend
+# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY.
+iq1_s_fake_quant = IQ1_S_FORMAT.fake_quant

modelopt/torch/quantization/ggml/iq2_xs.py

@@ -36,7 +36,7 @@
 from .codebooks import iq2_xs_grid_bytes
 from .common import (
     GGML_BLOCK_SIZE,
-    fake_quantize_with_cache,
+    IQFormat,
     narrow_to_float32,
     validate_block_chunk_size,
     validate_packed_weights,
@@ -255,22 +255,17 @@ def dequantize_iq2_xs(
     return decoded.reshape(shape)
 
 
-def iq2_xs_fake_quant(
-    inputs: torch.Tensor,
-    quantizer,
-    *,
-    block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
-    decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
-) -> torch.Tensor:
-    """IQ2_XS weight backend for TensorQuantizer, with pass-through backward."""
-    if getattr(quantizer, "num_bits", None) != "iq2_xs":
-        raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'")
-    return fake_quantize_with_cache(
-        inputs,
-        quantizer,
-        format_name="iq2_xs",
-        block_chunk_size=block_chunk_size,
-        decode_chunk_size=decode_chunk_size,
-        quantize=quantize_iq2_xs,
-        dequantize=dequantize_iq2_xs,
-    )
+IQ2_XS_FORMAT = IQFormat(
+    name="iq2_xs",
+    block_size=IQ2_XS_BLOCK_SIZE,
+    block_bytes=IQ2_XS_BLOCK_BYTES,
+    quantize=quantize_iq2_xs,
+    dequantize=dequantize_iq2_xs,
+    block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE,
+    decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE,
+)
+
+# Kept for callers of the per-format entry point. The record captured quantize_iq2_xs and
+# dequantize_iq2_xs when it was built, so patching those module functions changes neither backend
+# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY.
+iq2_xs_fake_quant = IQ2_XS_FORMAT.fake_quant