[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)  ## 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.  --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-02 03:14:52 +08:00 · 2026-09-29 20:35:38 +00:00
parent 1576f7d4ad
commit 3091b8ff69
16 changed files with 377 additions and 12 deletions
@@ -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,
@@ -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.
@@ -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. "
@@ -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;
+}
@@ -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",
@@ -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",
@@ -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
@@ -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)
 }