[4/5] Add the IQ2_S CUDA encoder and register the format (#2565)

### What does this PR do?

Type of change: new feature

**Second of two PRs adding IQ2_S** (2.5625 bits per weight). #2512
landed the PyTorch codec; this PR adds its **CUDA encoder** and makes
the format reachable:

- the CUDA encoder, its binding and extension build wiring, plus the
CUDA path in `quantize_iq2_s`
- an `IQFormat` record and **one `IQ_FORMAT_REGISTRY` entry**, so
backend dispatch, both exporters and `convert_hf_config` take it from
there
- the `ggml` package export
- the `general/ptq/iq2_s` recipe, its presets, `ptq.md` and a CHANGELOG
entry

The kernel lands with the registration so every registered format keeps
a CUDA encoder.

On the mixed-precision checkpoint #2511 measured
(`unsloth/Qwen3.8-27B-GGUF`), IQ2_S covers **9 tensors and 0.6 B
parameters**.

### The kernel

IQ2_S's **1024-entry codebook is twice IQ2_XS's**, which makes its
search the most expensive in the family. The codebook and its norms take
36 KiB of shared memory, the most of any IQ kernel but inside the 48 KiB
static limit, so they are declared statically like the IQ2_XS and
IQ2_XXS kernels.

That cost is why the kernel matters more here than anywhere else:

| | torch | CUDA | |
|---|---|---|---|
| IQ2_S, 5632×2048 weight | 0.8 M elem/s | **725.7 M elem/s** | **907×**
|
| extrapolated to a 27B model | ~9.8 hours | **~37 s** | |

### Usage

```bash
python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path <model> --recipe general/ptq/iq2_s
```

### Testing

Registering the format brings it under every registry-driven test with
no IQ2_S-specific test code: backend dispatch and weight caching, the
`num_bits` guard, `convert_hf_config` metadata (uniform and mixed
precision), all 9 Megatron export tests, and the two `TensorQuantizer`
tests in the shared battery. The shared CUDA battery gains one row.

- `tests/unit/torch/quantization/test_ggml_backend.py`,
`test_iq_formats.py`,
`tests/unit/torch/export/test_convert_hf_config.py`,
`tests/unit/recipe/test_presets.py`: **134 passed**
- broader unit sweep (`-k 'ggml or iq or gguf or registry'` over
quantization, export and recipe tests): **192 passed**. The one failure,
`test_export_registry.py::test_builtin_dispatch_covers_all_handler_shapes`,
is a `torchvision` import error in my environment, unrelated to IQ.
- `tests/gpu/torch/quantization/test_iq_formats_cuda.py`,
`test_iq1_s_cuda.py`, `test_iq2_xs_cuda.py`: **42 passed** on RTX PRO
6000 Blackwell (sm_120). 7 of them are IQ2_S: CUDA-vs-PyTorch encoder
parity, determinism, reconstruction at scale, zero and non-finite
policy, float64 input and the fallback path.
- `tests/gpu_megatron/torch/export/test_unified_export_megatron.py -k
'iq or ggml'`: **36 passed** (9 tests × 4 formats) in
`nvcr.io/nvidia/nemo:26.08`
- `tests/examples/hf_ptq/test_llm_ptq.py -k iq2_s`: **passed**.
TinyLlama PTQ through unified HF export writes `quant_algo: IQ2_S`,
`block_payload_bytes: 82`, and `down_proj` packed as `(2048, 22, 82)`
uint8.
- `general/ptq` now holds 30 recipes.
- The shared-memory change in `b7739d5d0` leaves the packed bytes
identical (same hash on a 5632×2048 weight), and packing runs at 849.1 M
elem/s against 825.7 before on RTX PRO 6000. The GPU battery was rerun:
42 passed.

All of the above was rerun after rebasing onto `main` at `c2aaa44f6`.
That base adds a Q8_0 packer to the same GGML extension (#2515), and
changes the hf_ptq example and the export code this format goes through.
The packed IQ2_S bytes still hash the same. On this RTX PRO 6000
(sm_120), two of #2515's own Q8_0 tests in
`tests/gpu/_extensions/test_torch_extensions.py` fail:
`test_cuda_ext_q8_0_zero_and_roundf_layout` and
`test_cuda_ext_q8_0_dequantizes_with_small_error`. They fail identically
on a clean `main` checkout, so they are not from this PR.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ No new code
sources or dependencies.
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ✅
- Did you get Claude approval on this PR?: ❌ Not yet run.

### Additional Information

Merge order: #2511 (IQ2_XXS, merged) → #2525 (format registry, merged) →
#2512 (IQ2_S codec, merged) → **this** → #2513 (IQ1_M).

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added IQ2_S weight-only quantization for eligible linear layers, at
2.5625 bits per weight.
* Added a PTQ recipe that requires no calibration data. Weights must
meet the existing 256-value block-size constraint.
* Added CUDA-accelerated packing for CUDA weights, with a Python
fallback when the CUDA extension is unavailable.
* **Documentation**
  * Updated the PTQ recipe catalog and IQ-format size tradeoffs.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

CHANGELOG.rst

@@ -19,6 +19,7 @@ Changelog
   until a runtime consumes that metadata combination.
 - Add IQ1_S and IQ2_XS weight-only quantization with GGML-compatible 256-value block encoders, built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes, and unified HF and Megatron export of the packed blocks. Quantized weights must have a final dimension divisible by 256, and Megatron export requires tensor and pipeline parallel sizes of 1.
 - Add ``iq2_xxs`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 2.0625 bits per weight between ``iq1_s`` and ``iq2_xs``. The same 256-value block constraint applies.
+- Add ``iq2_s`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 2.5625 bits per weight above ``iq2_xs``. The same 256-value block constraint applies.
 - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it.
 - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes.
 - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.

modelopt/torch/export/quant_format.py

@@ -41,6 +41,7 @@ QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
 QUANTIZATION_IQ1_S = "iq1_s"
 QUANTIZATION_IQ2_XXS = "iq2_xxs"
 QUANTIZATION_IQ2_XS = "iq2_xs"
+QUANTIZATION_IQ2_S = "iq2_s"
 
 # 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,

modelopt/torch/kernels/quantization/ggml/common.cuh

@@ -50,6 +50,7 @@ constexpr int kScaleBytes = 2;
 constexpr int kIq1sEntries = 2048;
 constexpr int kIq2xsEntries = 512;
 constexpr int kIq2xxsEntries = 256;
+constexpr int kIq2sEntries = 1024;
 
 // One CUDA block encodes one GGML block. The reductions below fold over exactly this many warps,
 // and each kernel static_asserts that its codebook divides evenly among the threads.

modelopt/torch/kernels/quantization/ggml/ggml.cpp

@@ -24,6 +24,7 @@
 at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid);
 at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales);
 at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales);
+at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales);
 at::Tensor q8_0_pack_cuda(at::Tensor input);
 
 namespace {
@@ -72,6 +73,21 @@ at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
   return iq2_xxs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
 }
 
+at::Tensor iq2_s_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
+  TORCH_CHECK(input.is_cuda(), "IQ2_S packing requires a CUDA input");
+  TORCH_CHECK(grid.is_cuda(), "IQ2_S packing requires a CUDA grid");
+  TORCH_CHECK(scales.is_cuda(), "IQ2_S packing requires CUDA scales");
+  modelopt::ggml::check_pack_inputs("IQ2_S", input, grid, modelopt::ggml::kIq2sEntries);
+  const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize;
+  TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
+                  scales.numel() == num_blocks,
+              "scales must be float16 [numel / 256]");
+  TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item<bool>(),
+              "scales must be finite and non-negative");
+  TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
+  return iq2_s_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
+}
+
 at::Tensor q8_0_pack(at::Tensor input) {
   TORCH_CHECK(input.is_cuda(), "Q8_0 packing requires a CUDA input");
   modelopt::ggml::check_scalar_pack_input("Q8_0", input, 32);
@@ -102,6 +118,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
              "Returns uint8 [numel / 256, 66] on the input device. Non-finite input elements are "
              "treated as zero during packing, and finite elements outside the float32 range "
              "saturate.");
+  module.def("iq2_s_pack", &iq2_s_pack,
+             "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
+             "dimension is a multiple of 256. The grid must be float32 [1024, 8] holding "
+             "non-negative codebook magnitudes, and scales must be finite non-negative float16 "
+             "[numel / 256]. "
+             "Returns uint8 [numel / 256, 82] on the input device. Non-finite input elements are "
+             "treated as zero during packing, and finite elements outside the float32 range "
+             "saturate.");
   module.def("q8_0_pack", &q8_0_pack,
              "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
              "dimension is a multiple of 32. Returns uint8 [numel / 32, 34] on the input device. "

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

@@ -0,0 +1,210 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common.cuh"
+
+namespace {
+
+using namespace modelopt::ggml;
+
+// The IQ2_S packed payload layout and format constants below follow the GGML
+// definition at:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
+constexpr int kEntries = kIq2sEntries;
+constexpr int kGroups = 16; // one 4-bit local scale per 16 values
+constexpr int kVectorsPerGroup = 2;
+constexpr int kVectors = kGroups * kVectorsPerGroup; // 32 codebook vectors per block
+constexpr int kLocalScales = 16;
+constexpr int kSubBlocks = 8;
+constexpr int kLowOffset = kScaleBytes;             // 32 low index bytes
+constexpr int kSignOffset = kLowOffset + kVectors;  // 32 sign masks
+constexpr int kHighOffset = kSignOffset + kVectors; // 8 bytes, two high bits per vector
+constexpr int kLocalScaleOffset = kHighOffset + kSubBlocks;
+constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2;
+constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8.
+
+static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries");
+static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two");
+static_assert(kPayloadBytes == 82, "IQ2_S blocks are 82 bytes");
+
+// IQ2_S stores all eight sign bits per vector rather than seven plus a parity bit, so unlike
+// IQ2_XS and IQ2_XXS the encoder never has to flip a coordinate: the best signs are simply the
+// input's, and the search compares magnitudes directly.
+__device__ __forceinline__ float magnitude_dot(const float *x, const float *q) {
+  float dot = 0.0f;
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j)
+    dot = fmaf(fabsf(x[j]), q[j], dot);
+  return dot;
+}
+
+template <typename scalar_t>
+__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid,
+                       const __half *scales, uint8_t *output) {
+  // The codebook and its norms take 36 KiB, the most in the family but inside the 48 KiB
+  // static shared-memory limit.
+  __shared__ float shared_grid[kEntries * kVectorSize];
+  __shared__ float grid_norm[kEntries];
+  __shared__ float warp_best[kWarps * kLocalScales];
+  __shared__ float group_error[kLocalScales];
+  __shared__ unsigned long long warp_keys[kWarps];
+  __shared__ int selected_local;
+  __shared__ uint8_t locals[kGroups];
+
+  const int tid = threadIdx.x;
+  const int64_t block = blockIdx.x;
+  if (block >= num_blocks)
+    return;
+
+  for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x)
+    shared_grid[i] = grid[i];
+  __syncthreads();
+  for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+    float norm = 0.0f;
+#pragma unroll
+    for (int j = 0; j < kVectorSize; ++j) {
+      const float q = shared_grid[entry * kVectorSize + j];
+      norm = fmaf(q, q, norm);
+    }
+    grid_norm[entry] = norm;
+  }
+  __syncthreads();
+
+  const scalar_t *source = input + block * kBlockSize;
+  uint8_t *payload = output + block * kPayloadBytes;
+  const __half d_half = scales[block];
+  const uint16_t d_bits = __half_as_ushort(d_half);
+  const float d = __half2float(d_half);
+  if (!store_block_scale<kPayloadBytes>(payload, d_bits))
+    return;
+  // The high-bit bytes are OR-ed into below, so they must start clear.
+  if (tid < kSubBlocks)
+    payload[kHighOffset + tid] = 0;
+  __syncthreads();
+
+#pragma unroll 1
+  for (int group = 0; group < kGroups; ++group) {
+    if (tid < kLocalScales)
+      group_error[tid] = 0.0f;
+    __syncthreads();
+
+#pragma unroll
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+      }
+      float local_best[kLocalScales];
+#pragma unroll
+      for (int local = 0; local < kLocalScales; ++local)
+        local_best[local] = FLT_MAX;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float dot = magnitude_dot(x, shared_grid + entry * kVectorSize);
+#pragma unroll
+        for (int local = 0; local < kLocalScales; ++local) {
+          const float scale = d * (2 * local + 1) * kLocalScaleStep;
+          local_best[local] =
+              fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale));
+        }
+      }
+      block_min_accumulate<kLocalScales>(local_best, warp_best, group_error);
+    }
+
+    if (tid == 0) {
+      selected_local = 0;
+      float best = group_error[0];
+#pragma unroll
+      for (int local = 1; local < kLocalScales; ++local) {
+        if (group_error[local] < best) {
+          best = group_error[local];
+          selected_local = local;
+        }
+      }
+      locals[group] = static_cast<uint8_t>(selected_local);
+    }
+    __syncthreads();
+    const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep;
+
+#pragma unroll
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+      }
+      unsigned long long key = ~0ULL;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float error =
+            clamped_quant_error(xnorm, magnitude_dot(x, shared_grid + entry * kVectorSize),
+                                grid_norm[entry], selected_scale);
+        const unsigned long long candidate = error_key(error, entry);
+        key = candidate < key ? candidate : key;
+      }
+      key = block_min_key(key, warp_keys);
+      if (tid == 0) {
+        const int entry = static_cast<int>(key & (kEntries - 1));
+        int sign_mask = 0;
+#pragma unroll
+        for (int j = 0; j < kVectorSize; ++j)
+          sign_mask |= static_cast<int>(x[j] < 0.0f) << j;
+        const int slot = group * kVectorsPerGroup + vector; // 0..31, GGML's qs index
+        payload[kLowOffset + slot] = static_cast<uint8_t>(entry & 0xFF);
+        payload[kSignOffset + slot] = static_cast<uint8_t>(sign_mask);
+        // Two high index bits per vector, packed four to a byte in vector order.
+        payload[kHighOffset + slot / 4] |=
+            static_cast<uint8_t>(((entry >> 8) & 0x3) << (2 * (slot % 4)));
+      }
+      __syncthreads();
+    }
+  }
+
+  if (tid < kGroups / 2)
+    payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4);
+}
+
+} // namespace
+
+at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) {
+  TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(),
+              "inputs must be contiguous");
+  check_pack_inputs("IQ2_S", input, grid, kEntries);
+  const int64_t num_blocks = input.numel() / kBlockSize;
+  TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
+                  scales.numel() == num_blocks,
+              "scales must be float16 [numel / 256]");
+  TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
+  c10::cuda::CUDAGuard guard(input.device());
+  auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte));
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_s_pack", [&] {
+        encode<scalar_t><<<static_cast<int>(num_blocks), kThreads, 0, stream>>>(
+            input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
+            reinterpret_cast<const __half *>(scales.data_ptr<at::Half>()),
+            output.data_ptr<uint8_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}

modelopt/torch/quantization/extensions.py

@@ -84,8 +84,8 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False):
 
     The formats share their packing helpers, CUDA version requirement, and build flags, so they
     build as a single extension: ``iq1_s_pack(input, grid)``,
-    ``iq2_xs_pack(input, grid, scales)``, ``iq2_xxs_pack(input, grid, scales)``, and
-    ``q8_0_pack(input)``.
+    ``iq2_xs_pack(input, grid, scales)``, ``iq2_xxs_pack(input, grid, scales)``,
+    ``iq2_s_pack(input, grid, scales)``, and ``q8_0_pack(input)``.
     """
     if not hasattr(get_cuda_ext_ggml, "extension") or (
         raise_if_failed and get_cuda_ext_ggml.extension is None
@@ -97,6 +97,7 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False):
                 kernels_ggml / "iq1_s.cu",
                 kernels_ggml / "iq2_xs.cu",
                 kernels_ggml / "iq2_xxs.cu",
+                kernels_ggml / "iq2_s.cu",
                 kernels_ggml / "q8_0.cu",
             ],
             cuda_version_specifiers=">=11.8",

modelopt/torch/quantization/ggml/__init__.py

@@ -19,6 +19,8 @@
 from . import backend as _backend
 from .iq1_s import *
 from .iq1_s import __all__ as _iq1_s_all
+from .iq2_s import *
+from .iq2_s import __all__ as _iq2_s_all
 from .iq2_xs import *
 from .iq2_xs import __all__ as _iq2_xs_all
 from .iq2_xxs import *
@@ -27,6 +29,7 @@ from .registry import IQ_FORMAT_REGISTRY, IQFormat
 
 __all__ = [  # noqa: PLE0604
     *_iq1_s_all,
+    *_iq2_s_all,
     *_iq2_xs_all,
     *_iq2_xxs_all,
     "IQ_FORMAT_REGISTRY",

modelopt/torch/quantization/ggml/iq2_s.py

@@ -37,9 +37,11 @@ https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015
 
 import torch
 
+from ..extensions import get_cuda_ext_ggml
 from .codebooks import iq2_s_grid_bytes
 from .common import (
     GGML_BLOCK_SIZE,
+    IQFormat,
     narrow_to_float32,
     validate_block_chunk_size,
     validate_packed_weights,
@@ -51,6 +53,7 @@ __all__ = [
     "IQ2_S_BLOCK_SIZE",
     "IQ2_S_EFFECTIVE_BITS",
     "dequantize_iq2_s",
+    "iq2_s_fake_quant",
     "iq2_s_grid",
     "quantize_iq2_s",
 ]
@@ -69,6 +72,7 @@ _IQ2_S_PEAK_TO_RMS_TAPER = 0.035
 # The grid is twice IQ2_XS's, so the same search tile costs twice the memory.
 _DEFAULT_BLOCK_CHUNK_SIZE = 128
 _DEFAULT_DECODE_CHUNK_SIZE = 4096
+_SCALE_BLOCK_CHUNK_SIZE = 4096
 
 _GRID_CACHE: dict[torch.device, torch.Tensor] = {}
 
@@ -169,6 +173,16 @@ def quantize_iq2_s(
     blocks = weight.contiguous().reshape(-1, IQ2_S_BLOCK_SIZE)
     grid = iq2_s_grid(weight.device)
     packed_shape = (*weight.shape[:-1], weight.shape[-1] // IQ2_S_BLOCK_SIZE, IQ2_S_BLOCK_BYTES)
+    if weight.is_cuda:
+        extension = get_cuda_ext_ggml()
+        if extension is not None:
+            scale_chunks = [
+                _predict_iq2_s_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE])
+                for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE)
+            ]
+            packed = extension.iq2_s_pack(blocks, grid, torch.cat(scale_chunks))
+            return packed.reshape(packed_shape), logical_shape
+
     chunks = [
         _encode_blocks(blocks[start : start + block_chunk_size], grid)
         for start in range(0, blocks.shape[0], block_chunk_size)
@@ -215,3 +229,19 @@ def dequantize_iq2_s(
         chunk_decoded = values.reshape(count, 16, 2, 8) * scales.unsqueeze(-1).unsqueeze(-1)
         decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_S_BLOCK_SIZE)
     return decoded.reshape(shape)
+
+
+IQ2_S_FORMAT = IQFormat(
+    name="iq2_s",
+    block_size=IQ2_S_BLOCK_SIZE,
+    block_bytes=IQ2_S_BLOCK_BYTES,
+    quantize=quantize_iq2_s,
+    dequantize=dequantize_iq2_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_iq2_s and
+# dequantize_iq2_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.
+iq2_s_fake_quant = IQ2_S_FORMAT.fake_quant

modelopt/torch/quantization/ggml/registry.py

@@ -17,6 +17,7 @@
 
 from .common import IQFormat
 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
 
@@ -29,5 +30,5 @@ __all__ = ["IQ_FORMAT_REGISTRY", "IQFormat"]
 # 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] = {
-    fmt.name: fmt for fmt in (IQ1_S_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT)
+    fmt.name: fmt for fmt in (IQ1_S_FORMAT, IQ2_XXS_FORMAT, IQ2_XS_FORMAT, IQ2_S_FORMAT)
 }

modelopt_recipes/configs/numerics/iq2_s.yaml

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# IQ2_S weight quantizer using the built-in fixed-scale codebook search.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig
+num_bits: iq2_s
+# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or
+# numerics. num_bits is the string "iq2_s", so the generic estimator cannot derive the
+# storage cost: 82 packed bytes * 8 / 256 weights. Keep in sync with IQ2_S_BLOCK_BYTES.
+effective_bits: 2.5625
+block_sizes:
+  -1: 256
+backend: ggml

modelopt_recipes/configs/ptq/presets/model/iq2_s.yaml

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# QuantizeConfig preset for IQ2_S weight-only quantization.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  iq2_s: configs/numerics/iq2_s
+
+algorithm:
+quant_cfg:
+  - $import: base_disable_all
+  - quantizer_name: '*weight_quantizer'
+    cfg:
+      $import: iq2_s
+  - quantizer_name: '*input_quantizer'
+    enable: false
+  - $import: default_disabled_quantizers

modelopt_recipes/general/ptq/iq2_s.yaml

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# IQ2_S weight-only PTQ.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+  preset: configs/ptq/presets/model/iq2_s
+
+metadata:
+  description: >-
+    Applies uniform GGML-compatible IQ2_S weight-only quantization to eligible linear layers.
+    This is not a mixed per-tensor precision preset. No calibration data is required.
+quantize:
+  $import: preset

modelopt_recipes/ptq.md

@@ -28,7 +28,7 @@ supported combinations.
 ### The shipped recipes
 
 <details>
-<summary>All 29 <code>general/ptq/</code> recipes (click to expand)</summary>
+<summary>All 30 <code>general/ptq/</code> recipes (click to expand)</summary>
 
 | Recipe | Model body | KV cache | Calibration |
 |--------|-----------|----------|-------------|
@@ -61,6 +61,7 @@ supported combinations.
 | `iq1_s` | IQ1_S W1A16, eligible linears | none | none (no calibration) |
 | `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), eligible linears | none | none (no calibration) |
 | `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), eligible linears | none | none (no calibration) |
+| `iq2_s` | IQ2_S W2A16 (2.56 bpw), eligible linears | none | none (no calibration) |
 
 </details>
 
@@ -139,10 +140,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_xxs` / `iq2_xs`** — GGML-compatible IQ weights
+- **`iq1_s` / `iq2_xxs` / `iq2_xs` / `iq2_s`** — GGML-compatible IQ weights
   on the eligible linear layers, with BF16 activations; `lm_head`, MoE routers,
   `conv1d` and the vision branch stay in BF16 like every other preset. The formats
-  trade size against accuracy in order: 1.56, 2.06 and 2.31 bits per weight. No calibration data is
+  trade size against accuracy in order: 1.56, 2.06, 2.31 and 2.56 bits per weight. No calibration data is
   required. Quantized weights must have a final dimension divisible by 256.
   Unified HF export writes the packed GGML blocks; Megatron export additionally
   requires tensor and pipeline parallel sizes of 1, and does not support

tests/examples/hf_ptq/test_llm_ptq.py

@@ -76,15 +76,16 @@ def test_ptq_whisper(command):
         PTQCommand(quant="int8_weight_only", kv_cache_quant="none"),
         PTQCommand(quant="int4_awq", kv_cache_quant="none"),
         PTQCommand(quant="w4a8_awq_beta", kv_cache_quant="none"),
-        # GGML IQ weight-only, recipe-driven: three formats between 1.56 and 2.31 bits
-        # per weight. These encoders require every weight's input dimension to be a
-        # multiple of 256; TinyLlama's 2048 and 5632 both are. None of them calibrates --
-        # every recipe sets algorithm: null -- so the only IQ-specific cost is packing each
-        # weight once on a CUDA encoder and decoding it on each forward, which fits the
-        # 300s tests/examples default.
+        # GGML IQ weight-only, recipe-driven: four formats between 1.56 and 2.56 bits
+        # per weight. These encoders require every weight's input dimension to be a multiple of
+        # 256; TinyLlama's 2048 and 5632 both are. None of them calibrates -- every recipe
+        # sets algorithm: null -- so the only IQ-specific cost is packing each weight once on
+        # a CUDA encoder and decoding it on each forward, which fits the 300s
+        # tests/examples default.
         PTQCommand(recipe="general/ptq/iq1_s", kv_cache_quant="none"),
         PTQCommand(recipe="general/ptq/iq2_xxs", kv_cache_quant="none"),
         PTQCommand(recipe="general/ptq/iq2_xs", kv_cache_quant="none"),
+        PTQCommand(recipe="general/ptq/iq2_s", kv_cache_quant="none"),
         PTQCommand(quant="nvfp4"),
         PTQCommand(quant="nvfp4_awq_lite"),
         # autoquant (recipe-driven)

tests/gpu/torch/quantization/test_iq_formats_cuda.py

@@ -24,11 +24,13 @@ import pytest
 import torch
 
 import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
+import modelopt.torch.quantization.ggml.iq2_s as iq2_s_module
 import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
 import modelopt.torch.quantization.ggml.iq2_xxs as iq2_xxs_module
 from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
 from modelopt.torch.quantization.ggml import (
     IQ1_S_BLOCK_BYTES,
+    IQ2_S_BLOCK_BYTES,
     IQ2_XS_BLOCK_BYTES,
     IQ2_XXS_BLOCK_BYTES,
     IQ_FORMAT_REGISTRY,
@@ -39,6 +41,7 @@ FORMATS = {
     "iq1_s": (iq1_s_module, "iq1_s_pack", IQ1_S_BLOCK_BYTES, False),
     "iq2_xxs": (iq2_xxs_module, "iq2_xxs_pack", IQ2_XXS_BLOCK_BYTES, True),
     "iq2_xs": (iq2_xs_module, "iq2_xs_pack", IQ2_XS_BLOCK_BYTES, True),
+    "iq2_s": (iq2_s_module, "iq2_s_pack", IQ2_S_BLOCK_BYTES, True),
 }
 
 

tests/unit/recipe/test_presets.py

@@ -36,6 +36,8 @@ from modelopt.torch.quantization.config import LocalHessianCalibConfig, Quantize
 from modelopt.torch.quantization.ggml import (
     IQ1_S_BLOCK_SIZE,
     IQ1_S_EFFECTIVE_BITS,
+    IQ2_S_BLOCK_SIZE,
+    IQ2_S_EFFECTIVE_BITS,
     IQ2_XS_BLOCK_SIZE,
     IQ2_XS_EFFECTIVE_BITS,
     IQ2_XXS_BLOCK_SIZE,
@@ -139,6 +141,7 @@ def test_mlp_weight_only_recipe_matches_its_mtq_cfg(recipe_name, cfg_name):
         ("iq1_s", IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS),
         ("iq2_xxs", IQ2_XXS_BLOCK_SIZE, IQ2_XXS_EFFECTIVE_BITS),
         ("iq2_xs", IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS),
+        ("iq2_s", IQ2_S_BLOCK_SIZE, IQ2_S_EFFECTIVE_BITS),
     ],
 )
 def test_iq_recipe_matches_packing_contract(qformat, block_size, effective_bits):