[4/5] Add the IQ2_S CUDA encoder and register the format

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_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 @@
 
 __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 @@
 
 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 @@
     "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 @@
 # 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 @@
 # 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)
 }