[OMNIML-5899] Add Q8_0 quantization codec and backend

modelopt/torch/kernels/quantization/ggml/q8_0.cu

@@ -53,7 +53,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, uint8_t *outpu
     payload[1] = static_cast<uint8_t>(d_bits >> 8);
   }
   d = __shfl_sync(0xffffffff, d, 0);
-  const float rounded = d > 0.0f ? roundf(value / d) : 0.0f;
+  const float inverse = d > 0.0f ? 1.0f / d : 0.0f;
+  const float rounded = roundf(value * inverse);
   const int quant = static_cast<int>(fminf(fmaxf(rounded, -kMaxQuant), kMaxQuant));
   payload[kScaleBytes + lane] = static_cast<uint8_t>(static_cast<int8_t>(quant));
 }

modelopt/torch/quantization/ggml/__init__.py

@@ -27,14 +27,19 @@
 from .iq2_xs import __all__ as _iq2_xs_all
 from .iq2_xxs import *
 from .iq2_xxs import __all__ as _iq2_xxs_all
-from .registry import IQ_FORMAT_REGISTRY, IQFormat
+from .q8_0 import *
+from .q8_0 import __all__ as _q8_0_all
+from .registry import GGML_FORMAT_REGISTRY, IQ_FORMAT_REGISTRY, GGMLFormat, IQFormat
 
 __all__ = [  # noqa: PLE0604
     *_iq1_m_all,
     *_iq1_s_all,
     *_iq2_s_all,
     *_iq2_xs_all,
     *_iq2_xxs_all,
+    *_q8_0_all,
+    "GGML_FORMAT_REGISTRY",
+    "GGMLFormat",
     "IQ_FORMAT_REGISTRY",
     "IQFormat",
 ]

modelopt/torch/quantization/ggml/backend.py

@@ -13,28 +13,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""TensorQuantizer backend dispatch for GGML-compatible weight-only IQ formats."""
+"""TensorQuantizer backend dispatch for GGML-compatible weight-only formats."""
 
 import torch
 
 from ..nn.modules.tensor_quantizer import register_quant_backend
-from .registry import IQ_FORMAT_REGISTRY
+from .registry import GGML_FORMAT_REGISTRY
 
 
 def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor:
-    """Dispatch an IQ quantizer to its format-specific implementation."""
+    """Dispatch a GGML quantizer to its format-specific implementation."""
     num_bits = getattr(quantizer, "num_bits", None)
     extra_args = getattr(quantizer, "backend_extra_args", None) or {}
     unknown_args = set(extra_args) - {"block_chunk_size", "decode_chunk_size"}
     if unknown_args:
         raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}")
     # num_bits arrives untyped from the quantizer and is a tuple for scalar formats,
     # so narrow before the lookup rather than relying on the dict to reject it.
-    iq_format = IQ_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None
-    if iq_format is None:
-        supported = ", ".join(repr(name) for name in sorted(IQ_FORMAT_REGISTRY))
+    ggml_format = GGML_FORMAT_REGISTRY.get(num_bits) if isinstance(num_bits, str) else None
+    if ggml_format is None:
+        supported = ", ".join(repr(name) for name in sorted(GGML_FORMAT_REGISTRY))
         raise ValueError(f"The ggml backend requires num_bits in ({supported})")
-    return iq_format.fake_quant(inputs, quantizer, **extra_args)
+    return ggml_format.fake_quant(inputs, quantizer, **extra_args)
 
 
 register_quant_backend("ggml", ggml_fake_quant)

modelopt/torch/quantization/ggml/common.py

@@ -173,11 +173,11 @@ def fake_quantize_with_cache(
 
 
 @dataclass(frozen=True)
-class IQFormat:
-    """Everything backend dispatch and export need to know about one IQ format.
+class GGMLFormat:
+    """Everything backend dispatch and export need to know about one GGML format.
 
     Each format module declares one of these beside its encoder and decoder, and
-    :data:`~modelopt.torch.quantization.ggml.registry.IQ_FORMAT_REGISTRY` lists them. The
+    :data:`~modelopt.torch.quantization.ggml.registry.GGML_FORMAT_REGISTRY` lists them. The
     per-format pieces -- codebook, search, payload layout -- stay in the format's module; what
     lives here is the part every format does the same way.
     """
@@ -243,6 +243,11 @@ def pack(self, weight: torch.Tensor, quantizer=None) -> torch.Tensor:
         )
 
 
+# Compatibility alias for callers that imported the record type before the registry was
+# generalized from the IQ family to every supported GGML block format.
+IQFormat = GGMLFormat
+
+
 def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor:
     """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does.
 
@@ -261,14 +266,16 @@ def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor:
     return finite.float()
 
 
-def validate_weight(weight: torch.Tensor, format_name: str) -> None:
+def validate_weight(
+    weight: torch.Tensor, format_name: str, *, block_size: int = GGML_BLOCK_SIZE
+) -> None:
     """Validate weight metadata accepted by the current GGML block encoders."""
     if weight.numel() == 0:
         raise ValueError(f"{format_name} requires a non-empty weight")
-    if weight.dim() == 0 or weight.shape[-1] % GGML_BLOCK_SIZE:
+    if weight.dim() == 0 or weight.shape[-1] % block_size:
         raise ValueError(
             f"{format_name} requires the last weight dimension to be divisible by "
-            f"{GGML_BLOCK_SIZE}, got shape {tuple(weight.shape)}"
+            f"{block_size}, got shape {tuple(weight.shape)}"
         )
     if not weight.is_floating_point():
         raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}")
@@ -287,6 +294,7 @@ def validate_packed_weights(
     weight_shape: torch.Tensor,
     *,
     block_bytes: int,
+    block_size: int = GGML_BLOCK_SIZE,
     format_name: str,
 ) -> tuple[int, ...]:
     """Validate a packed payload and return its logical shape."""
@@ -303,9 +311,9 @@ def validate_packed_weights(
     if weight_shape.dim() != 1 or weight_shape.dtype not in integral_dtypes:
         raise ValueError("weight_shape must be a one-dimensional integral tensor")
     shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist())
-    if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % GGML_BLOCK_SIZE:
+    if not shape or any(dimension <= 0 for dimension in shape) or shape[-1] % block_size:
         raise ValueError(f"invalid {format_name} logical weight shape: {shape}")
-    expected_payload_values = math.prod(shape) // GGML_BLOCK_SIZE * block_bytes
+    expected_payload_values = math.prod(shape) // block_size * block_bytes
     if packed_weights.numel() != expected_payload_values:
         raise ValueError("packed_weights size does not match weight_shape")
     return shape

modelopt/torch/quantization/ggml/iq1_s.py

@@ -274,5 +274,5 @@ def dequantize_iq1_s(
 
 # Kept for callers of the per-format entry point. The record captured quantize_iq1_s and
 # dequantize_iq1_s when it was built, so patching those module functions changes neither backend
-# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY.
+# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY.
 iq1_s_fake_quant = IQ1_S_FORMAT.fake_quant

modelopt/torch/quantization/ggml/iq2_xs.py

@@ -272,5 +272,5 @@ def dequantize_iq2_xs(
 
 # Kept for callers of the per-format entry point. The record captured quantize_iq2_xs and
 # dequantize_iq2_xs when it was built, so patching those module functions changes neither backend
-# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY.
+# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY.
 iq2_xs_fake_quant = IQ2_XS_FORMAT.fake_quant

modelopt/torch/quantization/ggml/iq2_xxs.py

@@ -282,5 +282,5 @@ def dequantize_iq2_xxs(
 
 # Kept for callers of the per-format entry point. The record captured quantize_iq2_xxs and
 # dequantize_iq2_xxs when it was built, so patching those module functions changes neither backend
-# dispatch nor this alias; substitute a format's encoder or decoder in IQ_FORMAT_REGISTRY.
+# dispatch nor this alias; substitute a format's encoder or decoder in GGML_FORMAT_REGISTRY.
 iq2_xxs_fake_quant = IQ2_XXS_FORMAT.fake_quant

modelopt/torch/quantization/ggml/q8_0.py

@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Q8_0 fake quantization and GGML-compatible block packing.
+
+Every 32 logical values become one 34-byte ``block_q8_0`` payload: a
+little-endian FP16 block scale followed by 32 signed int8 quants.
+"""
+
+import torch
+
+from ..extensions import get_cuda_ext_ggml
+from .common import (
+    GGMLFormat,
+    narrow_to_float32,
+    validate_block_chunk_size,
+    validate_packed_weights,
+    validate_weight,
+)
+
+__all__ = [
+    "Q8_0_BLOCK_BYTES",
+    "Q8_0_BLOCK_SIZE",
+    "Q8_0_EFFECTIVE_BITS",
+    "Q8_0_FORMAT",
+    "dequantize_q8_0",
+    "q8_0_fake_quant",
+    "quantize_q8_0",
+]
+
+Q8_0_BLOCK_SIZE = 32
+Q8_0_BLOCK_BYTES = 34
+Q8_0_EFFECTIVE_BITS = Q8_0_BLOCK_BYTES * 8 / Q8_0_BLOCK_SIZE
+_Q8_0_MAX_QUANT = 127
+_DEFAULT_BLOCK_CHUNK_SIZE = 4096
+_DEFAULT_DECODE_CHUNK_SIZE = 16384
+
+
+def _encode_blocks(blocks: torch.Tensor) -> torch.Tensor:
+    """Encode a moderate-size batch of flattened 32-value blocks."""
+    x = narrow_to_float32(blocks)
+    block_count = x.shape[0]
+    amax = x.abs().amax(dim=1)
+    # Keep the serialized fp16 scale finite. This extends the existing GGML backend policy for
+    # finite float64 values outside the float32/fp16 scale range.
+    d_float = (amax / _Q8_0_MAX_QUANT).clamp(max=65504.0)
+    d = d_float.to(torch.float16)
+    # Select quants with the unrounded float32 scale, not the serialized fp16 scale.
+    inverse = torch.where(d_float > 0, d_float.reciprocal(), torch.zeros_like(d_float))
+    normalized = x * inverse.unsqueeze(1)
+    # C roundf, used by the canonical Q8_0 encoder, rounds half-way cases away from zero.
+    magnitude = normalized.abs()
+    rounded = normalized.sign() * torch.where(
+        magnitude.frac() == 0.5,
+        magnitude.floor() + 1,
+        magnitude.round(),
+    )
+    quants = rounded.clamp(-_Q8_0_MAX_QUANT, _Q8_0_MAX_QUANT).to(torch.int8)
+
+    packed = torch.empty((block_count, Q8_0_BLOCK_BYTES), dtype=torch.uint8, device=x.device)
+    packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2)
+    packed[:, 2:] = quants.view(torch.uint8)
+    return packed
+
+
+@torch.no_grad()
+def quantize_q8_0(
+    weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a floating-point weight into GGML-compatible Q8_0 blocks."""
+    validate_weight(weight, "Q8_0", block_size=Q8_0_BLOCK_SIZE)
+    validate_block_chunk_size(block_chunk_size)
+
+    logical_shape = torch.tensor(weight.shape, dtype=torch.int64)
+    blocks = weight.contiguous().reshape(-1, Q8_0_BLOCK_SIZE)
+    if weight.is_cuda:
+        extension = get_cuda_ext_ggml()
+        if extension is not None:
+            packed = extension.q8_0_pack(blocks)
+            packed_shape = (
+                *weight.shape[:-1],
+                weight.shape[-1] // Q8_0_BLOCK_SIZE,
+                Q8_0_BLOCK_BYTES,
+            )
+            return packed.reshape(packed_shape), logical_shape
+
+    chunks = [
+        _encode_blocks(blocks[start : start + block_chunk_size])
+        for start in range(0, blocks.shape[0], block_chunk_size)
+    ]
+    packed_shape = (
+        *weight.shape[:-1],
+        weight.shape[-1] // Q8_0_BLOCK_SIZE,
+        Q8_0_BLOCK_BYTES,
+    )
+    return torch.cat(chunks).reshape(packed_shape), logical_shape
+
+
+@torch.no_grad()
+def dequantize_q8_0(
+    packed_weights: torch.Tensor,
+    weight_shape: torch.Tensor,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
+) -> torch.Tensor:
+    """Decode GGML-compatible Q8_0 payload bytes."""
+    shape = validate_packed_weights(
+        packed_weights,
+        weight_shape,
+        block_bytes=Q8_0_BLOCK_BYTES,
+        block_size=Q8_0_BLOCK_SIZE,
+        format_name="Q8_0",
+    )
+    validate_block_chunk_size(block_chunk_size)
+
+    blocks = packed_weights.contiguous().reshape(-1, Q8_0_BLOCK_BYTES)
+    decoded = torch.empty((blocks.shape[0], Q8_0_BLOCK_SIZE), dtype=dtype, device=blocks.device)
+    for start in range(0, blocks.shape[0], block_chunk_size):
+        stop = min(start + block_chunk_size, blocks.shape[0])
+        block_chunk = blocks[start:stop]
+        d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float()
+        quants = block_chunk[:, 2:].contiguous().view(torch.int8).float()
+        decoded[start:stop] = (d.unsqueeze(1) * quants).to(dtype)
+    return decoded.reshape(shape)
+
+
+Q8_0_FORMAT = GGMLFormat(
+    name="q8_0",
+    block_size=Q8_0_BLOCK_SIZE,
+    block_bytes=Q8_0_BLOCK_BYTES,
+    quantize=quantize_q8_0,
+    dequantize=dequantize_q8_0,
+    block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE,
+    decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE,
+)
+
+# Keep the public per-format entry point on the same record used by backend dispatch.
+q8_0_fake_quant = Q8_0_FORMAT.fake_quant

modelopt/torch/quantization/ggml/registry.py

@@ -13,24 +13,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""The GGML IQ formats, listed once for backend dispatch and export."""
+"""GGML block formats, listed once for backend dispatch and export."""
 
-from .common import IQFormat
+from .common import GGMLFormat, IQFormat
 from .iq1_m import IQ1_M_FORMAT
 from .iq1_s import IQ1_S_FORMAT
 from .iq2_s import IQ2_S_FORMAT
 from .iq2_xs import IQ2_XS_FORMAT
 from .iq2_xxs import IQ2_XXS_FORMAT
+from .q8_0 import Q8_0_FORMAT
 
-__all__ = ["IQ_FORMAT_REGISTRY", "IQFormat"]
+__all__ = ["GGML_FORMAT_REGISTRY", "IQ_FORMAT_REGISTRY", "GGMLFormat", "IQFormat"]
 
-# Every IQ format, keyed by the name a quantizer's num_bits carries, in increasing bits per
-# weight. Backend dispatch and both exporters read this mapping and export's IQ_FORMATS is derived
-# from it, so adding a format is one entry here rather than a row in several parallel tables.
+# Every GGML block format, keyed by the name a quantizer's num_bits carries, in increasing bits
+# per weight. Backend dispatch reads this mapping, and export will derive its supported formats
+# from it rather than maintaining parallel metadata and packer tables.
 #
 # It is an explicit list rather than formats registering themselves on import, so its contents
 # never depend on which modules happen to have been imported first.
-IQ_FORMAT_REGISTRY: dict[str, IQFormat] = {
+GGML_FORMAT_REGISTRY: dict[str, GGMLFormat] = {
     fmt.name: fmt
-    for fmt in (IQ1_S_FORMAT, IQ1_M_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, IQ2_S_FORMAT)
+    for fmt in (
+        IQ1_S_FORMAT,
+        IQ1_M_FORMAT,
+        IQ2_XXS_FORMAT,
+        IQ2_XS_FORMAT,
+        IQ2_S_FORMAT,
+        Q8_0_FORMAT,
+    )
+}
+
+# Compatibility view for the already-merged IQ export path. Export migrates to the general
+# registry in the Q8_0 export PR; IQ-specific conformance tests continue to use this subset.
+IQ_FORMAT_REGISTRY: dict[str, GGMLFormat] = {
+    name: fmt for name, fmt in GGML_FORMAT_REGISTRY.items() if name.startswith("iq")
 }