[OMNIML-5899] Add IQ quantization codecs and backend (#2446)

## Summary - add IQ1_S and IQ2_XS reference codecs and a weight-only fake-quant backend - register and export both formats from the quantization package - cache compact packed weights across unchanged forwards and invalidate on tensor or config changes - use one Python-side IQ2_XS FP16 scale predictor for both reference and CUDA packing - validate packed payload metadata, normalize CUDA cache keys, and define a shared non-finite policy ## PR split This work is split into four focused PRs. Each PR targets `main` and owns a disjoint file set: 1. **Kernel** — [#2448: Add CUDA kernels for IQ packing](https://github.com/NVIDIA/Model-Optimizer/pull/2448) 2. **Quantization** — [#2446: Add IQ quantization codecs and backend](https://github.com/NVIDIA/Model-Optimizer/pull/2446) 3. **Export** — [#2447: Export IQ checkpoints from HF and Megatron](https://github.com/NVIDIA/Model-Optimizer/pull/2447) 4. **Recipes** — [#2449: Add IQ post-training quantization recipes](https://github.com/NVIDIA/Model-Optimizer/pull/2449) The required merge order is #2448, #2446, #2447, then #2449. ## Scope This PR owns the Python codecs, backend dispatch, package registration, license attribution, CPU codec/backend tests, and CUDA numerical/reference-path tests. The native CUDA layer and direct extension tests remain in #2448; export and recipes remain in their own PRs. ## Why the codecs are separate from `qtensor` The new `ggml/` package contains stateless reference codecs and fake-quant backend functions. They transform ordinary tensors into packed format payloads and reconstruct tensors for fake quantization; they do not define persistent runtime quantized-tensor objects. `BaseQuantizedTensor` subclasses under `qtensor/` own runtime tensor objects and execution dispatch. Keeping the codecs separate avoids claiming a runtime tensor contract that these formats do not yet provide. A `qtensor` type can be added later if a runtime execution path requires one. ## Compatibility boundary The Python encoders intentionally use fixed-scale, unweighted searches. They are not intended to reproduce another encoder's bytes for every input when that encoder performs iterative scale refinement or importance weighting. Compatibility is defined by the canonical codebooks, 50/74-byte payload layouts, and pinned dequantization formulas. IQ2_XS computes the FP16 superblock scale once in the Python predictor and passes it to the CUDA packer. This removes a duplicate floating-point reduction and makes native/reference byte parity use the same scale. Non-finite input elements are treated as zero during packing in both implementations. The unit tests construct nonzero payload fields independently and validate metadata, signs, local scales, and global scales. The CUDA tests compare native packed bytes with this Python reference encoder. ## Test coverage - [IQ1_S CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq1_s.py) - [IQ2_XS CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq2_xs.py) - [registered backend and cache tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_ggml_backend.py) - [IQ1_S CUDA byte-parity, numerical, non-finite, zero-payload, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq1_s_cuda.py) - [IQ2_XS CUDA byte-parity, numerical, non-finite, zero-payload, underflow, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq2_xs_cuda.py) ## Licensing The embedded codebook data cites the pinned upstream MIT source, carries its license notice, and uses the repository's third-party license mechanism. Human OSRB/code-owner confirmation is still required; this PR does not claim that approval. ## Validation - focused lint, format, and type checks pass for all changed Python files - 36 focused CPU codec and backend tests pass locally - all 20 direct-extension and CUDA integration test cases collect locally; runtime CUDA execution remains delegated to GPU CI - restricted-term scan passes  ## Summary by CodeRabbit - **New Features** - Added GGML quantization support for IQ1_S and IQ2_XS formats. - Added quantization, dequantization, and fake-quantization workflows with pass-through gradients. - Added CPU fallback when CUDA acceleration is unavailable. - Added validation for packed weights, tensor shapes, formats, and backend options. - Added configurable chunk processing and caching for repeated quantization. - **Tests** - Added comprehensive CPU and CUDA coverage for accuracy, validation, caching, fallback behavior, and edge cases.  --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-10-02 03:14:52 +08:00 · 2026-09-18 06:27:58 +00:00
parent 9895d6f129
commit 9e3d555aa1
14 changed files with 1650 additions and 0 deletions
@@ -0,0 +1,112 @@
+# 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.
+
+import torch
+
+import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
+from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
+from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, iq1_s_grid, quantize_iq1_s
+
+
+def _extension():
+    extension = get_cuda_ext_ggml(raise_if_failed=True)
+    assert extension is not None
+    return extension
+
+
+def test_iq1_s_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch):
+    generator = torch.Generator(device="cuda").manual_seed(1234)
+    weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
+
+    packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50)
+    packed_again = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(8, 1, 50)
+    monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
+    reference, shape = quantize_iq1_s(weight)
+    reconstructed = dequantize_iq1_s(packed, shape)
+
+    assert packed.shape == (8, 1, 50)
+    assert torch.equal(packed, packed_again)
+    assert torch.equal(packed, reference)
+    assert shape.device.type == "cpu"
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+    assert normalized_mse < 0.25
+
+
+def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout():
+    weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16)
+    packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50)
+    shape = torch.tensor(weight.shape, device="cuda")
+
+    assert not packed.any()
+    assert torch.equal(dequantize_iq1_s(packed, shape), weight)
+
+
+def test_iq1_s_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch):
+    weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16)
+    weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda")
+
+    packed = _extension().iq1_s_pack(weight, iq1_s_grid("cuda")).reshape(1, 1, 50)
+    monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
+    reference, _ = quantize_iq1_s(weight)
+
+    assert torch.equal(packed, reference)
+
+
+def test_iq1_s_cuda_falls_back_to_pytorch_encoder(monkeypatch):
+    monkeypatch.setattr(iq1_s_module, "get_cuda_ext_ggml", lambda: None)
+    generator = torch.Generator(device="cuda").manual_seed(1234)
+    weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq1_s(weight)
+    reconstructed = dequantize_iq1_s(packed, shape)
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+
+    assert packed.shape == (2, 1, 50)
+    assert normalized_mse < 0.25
+
+
+def test_iq1_s_cuda_float64_matches_pytorch_encoder():
+    """float64 weights inside the float32 range must pack identically on both paths."""
+    weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7))
+
+    reference, _ = quantize_iq1_s(weight)
+    packed, _ = quantize_iq1_s(weight.cuda())
+
+    assert torch.equal(reference, packed.cpu())
+
+
+def test_iq1_s_cuda_saturates_finite_values_above_the_float32_range():
+    """The extension saturates such values rather than dropping them to zero.
+
+    Byte parity with the reference encoder is not asserted here: at these magnitudes the
+    squared-error objective overflows to infinity in float32, so every codebook candidate ties
+    and the two search implementations break that tie differently. The saturation policy is
+    what both paths must agree on.
+    """
+    weight = torch.randn(1, 256, dtype=torch.float64, device="cuda")
+    weight[0, 7] = 1e100
+    saturated = weight.clone()
+    saturated[0, 7] = torch.finfo(torch.float32).max
+    zeroed = weight.clone()
+    zeroed[0, 7] = 0.0
+
+    packed, _ = quantize_iq1_s(weight)
+
+    assert torch.equal(packed, quantize_iq1_s(saturated)[0])
+    assert not torch.equal(packed, quantize_iq1_s(zeroed)[0])
@@ -0,0 +1,125 @@
+# 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.
+
+import torch
+
+import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
+from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
+from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, iq2_xs_grid, quantize_iq2_xs
+
+
+def _extension():
+    extension = get_cuda_ext_ggml(raise_if_failed=True)
+    assert extension is not None
+    return extension
+
+
+def _pack(weight):
+    blocks = weight.contiguous().reshape(-1, 256)
+    scales = iq2_xs_module._predict_iq2_xs_scales(blocks)
+    return _extension().iq2_xs_pack(weight, iq2_xs_grid("cuda"), scales)
+
+
+def test_iq2_xs_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch):
+    generator = torch.Generator(device="cuda").manual_seed(1234)
+    weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16)
+
+    packed = _pack(weight).reshape(8, 2, 74)
+    packed_again = _pack(weight).reshape(8, 2, 74)
+    monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
+    reference, shape = quantize_iq2_xs(weight)
+    reconstructed = dequantize_iq2_xs(packed, shape)
+
+    assert packed.shape == (8, 2, 74)
+    assert torch.equal(packed, packed_again)
+    assert torch.equal(packed, reference)
+    assert shape.device.type == "cpu"
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+    assert normalized_mse < 0.1
+
+
+def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout():
+    weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16)
+    packed = _pack(weight).reshape(1, 1, 74)
+    shape = torch.tensor(weight.shape, device="cuda")
+
+    assert not packed.any()
+    assert torch.equal(dequantize_iq2_xs(packed, shape), weight)
+
+
+def test_iq2_xs_cuda_underflowed_scale_has_canonical_zero_encoding():
+    weight = torch.full((1, 256), -1e-6, device="cuda", dtype=torch.bfloat16)
+    packed = _pack(weight).reshape(1, 1, 74)
+
+    assert not packed.any()
+
+
+def test_iq2_xs_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch):
+    weight = torch.randn((1, 256), device="cuda", dtype=torch.bfloat16)
+    weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf], device="cuda")
+
+    packed = _pack(weight).reshape(1, 1, 74)
+    monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
+    reference, _ = quantize_iq2_xs(weight)
+
+    assert torch.equal(packed, reference)
+
+
+def test_iq2_xs_cuda_falls_back_to_pytorch_encoder(monkeypatch):
+    monkeypatch.setattr(iq2_xs_module, "get_cuda_ext_ggml", lambda: None)
+    generator = torch.Generator(device="cuda").manual_seed(1234)
+    weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq2_xs(weight)
+    reconstructed = dequantize_iq2_xs(packed, shape)
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+
+    assert packed.shape == (2, 1, 74)
+    assert normalized_mse < 0.1
+
+
+def test_iq2_xs_cuda_float64_matches_pytorch_encoder():
+    """float64 weights inside the float32 range must pack identically on both paths."""
+    weight = torch.randn(4, 256, dtype=torch.float64, generator=torch.Generator().manual_seed(7))
+
+    reference, _ = quantize_iq2_xs(weight)
+    packed, _ = quantize_iq2_xs(weight.cuda())
+
+    assert torch.equal(reference, packed.cpu())
+
+
+def test_iq2_xs_cuda_saturates_finite_values_above_the_float32_range():
+    """The extension saturates such values rather than dropping them to zero.
+
+    Byte parity with the reference encoder is not asserted here: at these magnitudes the
+    squared-error objective overflows to infinity in float32, so every codebook candidate ties
+    and the two search implementations break that tie differently. The saturation policy is
+    what both paths must agree on.
+    """
+    weight = torch.randn(1, 256, dtype=torch.float64, device="cuda")
+    weight[0, 7] = 1e100
+    saturated = weight.clone()
+    saturated[0, 7] = torch.finfo(torch.float32).max
+    zeroed = weight.clone()
+    zeroed[0, 7] = 0.0
+
+    packed, _ = quantize_iq2_xs(weight)
+
+    assert torch.equal(packed, quantize_iq2_xs(saturated)[0])
+    assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0])
@@ -0,0 +1,137 @@
+# 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.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.ggml.backend as backend_module
+import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
+import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
+from modelopt.torch.quantization.ggml.backend import ggml_fake_quant
+from modelopt.torch.quantization.ggml.common import narrow_to_float32
+
+
+@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
+def test_ggml_backend_via_quantize(num_bits):
+    torch.manual_seed(1234)
+    model = torch.nn.Linear(256, 2, bias=False)
+    inputs = torch.randn(2, 256)
+    unquantized_output = model(inputs).detach()
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": num_bits, "backend": "ggml"},
+                "enable": True,
+            },
+        ],
+        "algorithm": "max",
+    }
+
+    mtq.quantize(model, config, forward_loop=lambda module: module(inputs))
+    output = model(inputs)
+
+    assert model.weight_quantizer.backend == "ggml"
+    assert model.weight_quantizer.num_bits == num_bits
+    assert output.shape == (2, 2)
+    assert torch.isfinite(output).all()
+    assert not torch.equal(output, unquantized_output)
+
+
+def test_ggml_backend_rejects_unknown_format():
+    with pytest.raises(ValueError, match="requires num_bits"):
+        ggml_fake_quant(torch.ones(1, 256), SimpleNamespace(num_bits="unknown"))
+
+
+def test_ggml_codecs_are_exported_from_quantization_package():
+    assert mtq.quantize_iq1_s is iq1_s_module.quantize_iq1_s
+    assert mtq.quantize_iq2_xs is iq2_xs_module.quantize_iq2_xs
+
+
+def test_ggml_backend_forwards_block_chunk_size(monkeypatch):
+    received = {}
+
+    def fake_quant(inputs, _quantizer, *, block_chunk_size):
+        received["block_chunk_size"] = block_chunk_size
+        return inputs
+
+    monkeypatch.setattr(backend_module, "iq1_s_fake_quant", fake_quant)
+    inputs = torch.ones(1, 256)
+    quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"block_chunk_size": 17})
+
+    assert ggml_fake_quant(inputs, quantizer) is inputs
+    assert received == {"block_chunk_size": 17}
+
+
+def test_ggml_backend_rejects_unknown_extra_arg():
+    quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args={"unknown": 1})
+
+    with pytest.raises(ValueError, match="Unsupported ggml backend_extra_args"):
+        ggml_fake_quant(torch.ones(1, 256), quantizer)
+
+
+@pytest.mark.parametrize(
+    ("num_bits", "module", "fake_quant_name", "quantize_name"),
+    [
+        ("iq1_s", iq1_s_module, "iq1_s_fake_quant", "quantize_iq1_s"),
+        ("iq2_xs", iq2_xs_module, "iq2_xs_fake_quant", "quantize_iq2_xs"),
+    ],
+)
+def test_ggml_backend_caches_packed_weight_and_invalidates_on_change(
+    monkeypatch, num_bits, module, fake_quant_name, quantize_name
+):
+    weight = torch.randn(1, 256)
+    quantizer = SimpleNamespace(num_bits=num_bits, _quantizer_cache=None)
+    original_quantize = getattr(module, quantize_name)
+    call_count = 0
+
+    def counted_quantize(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        return original_quantize(*args, **kwargs)
+
+    monkeypatch.setattr(module, quantize_name, counted_quantize)
+    fake_quant = getattr(module, fake_quant_name)
+
+    fake_quant(weight, quantizer, block_chunk_size=1)
+    fake_quant(weight, quantizer, block_chunk_size=1)
+    assert call_count == 1
+
+    fake_quant(weight, quantizer, block_chunk_size=2)
+    assert call_count == 2
+
+    with torch.no_grad():
+        weight.add_(0.01)
+    fake_quant(weight, quantizer, block_chunk_size=2)
+    assert call_count == 3
+
+
+def test_narrow_to_float32_matches_the_cuda_load_float_policy():
+    """Non-finite elements become zero; finite out-of-range elements saturate."""
+    largest = torch.finfo(torch.float32).max
+    values = torch.tensor(
+        [torch.nan, torch.inf, -torch.inf, 1e100, -1e100, 1.5], dtype=torch.float64
+    )
+
+    narrowed = narrow_to_float32(values)
+
+    assert narrowed.dtype is torch.float32
+    assert torch.equal(
+        narrowed, torch.tensor([0.0, 0.0, 0.0, largest, -largest, 1.5], dtype=torch.float32)
+    )
@@ -0,0 +1,183 @@
+# 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.
+
+import pytest
+import torch
+
+import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module
+from modelopt.torch.quantization.ggml.iq1_s import (
+    IQ1_S_BLOCK_BYTES,
+    dequantize_iq1_s,
+    iq1_s_fake_quant,
+    iq1_s_grid,
+    quantize_iq1_s,
+)
+
+
+def test_iq1_s_canonical_grid():
+    grid = iq1_s_grid()
+
+    assert grid.shape == (2048, 8)
+    assert grid.dtype == torch.float32
+    assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0}
+    assert grid[0].tolist() == [-1.0] * 8
+
+
+def test_iq1_s_grid_normalizes_unindexed_cuda_device(monkeypatch):
+    cached = torch.empty(0)
+    indexed_device = torch.device("cuda", 7)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 7)
+    monkeypatch.setitem(iq1_s_module._GRID_CACHE, indexed_device, cached)
+
+    assert iq1_s_grid("cuda") is cached
+
+
+def test_iq1_s_zero_block_has_canonical_zero_encoding():
+    weight = torch.zeros((2, 256), dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq1_s(weight)
+
+    assert packed.shape == (2, 1, IQ1_S_BLOCK_BYTES)
+    assert packed.dtype == torch.uint8
+    assert not packed.any()
+    assert shape.tolist() == [2, 256]
+    assert torch.equal(dequantize_iq1_s(packed, shape), weight)
+
+
+def test_iq1_s_dequantizes_ggml_metadata_bit_fields():
+    packed = torch.zeros((1, 1, 50), dtype=torch.uint8)
+    d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8)
+    packed[0, 0, :2] = d
+    entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64)
+    packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8)
+    qh = (
+        ((entries[0] >> 8) & 7)
+        | (((entries[1] >> 8) & 7) << 3)
+        | (((entries[2] >> 8) & 7) << 6)
+        | (((entries[3] >> 8) & 7) << 9)
+        | (3 << 12)
+        | (1 << 15)
+    )
+    packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8)
+    packed[0, 0, 35] = (qh >> 8).to(torch.uint8)
+
+    decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32)
+    expected = (iq1_s_grid()[entries] - 0.125) * 14.0
+
+    assert torch.equal(decoded[0, :32].reshape(4, 8), expected)
+
+
+def test_iq1_s_round_trip_and_payload_fields():
+    generator = torch.Generator().manual_seed(1234)
+    weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq1_s(weight, block_chunk_size=1)
+    reconstructed = dequantize_iq1_s(packed, shape, block_chunk_size=1)
+    default_reconstructed = dequantize_iq1_s(packed, shape)
+
+    assert packed.shape == (2, 1, 50)
+    assert reconstructed.shape == weight.shape
+    assert reconstructed.dtype == torch.bfloat16
+    assert torch.equal(reconstructed, default_reconstructed)
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+    assert normalized_mse < 0.25
+
+    blocks = packed.reshape(-1, 50)
+    qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8)
+    assert torch.all(((qh >> 12) & 0x7) < 8)
+    assert torch.all((qh & 0xFFF) < 0x1000)
+
+
+def test_iq1_s_search_is_independent_of_default_dtype():
+    generator = torch.Generator().manual_seed(0)
+    weight = torch.randn((8, 256), generator=generator, dtype=torch.float32)
+    expected, _ = quantize_iq1_s(weight)
+
+    default_dtype = torch.get_default_dtype()
+    try:
+        torch.set_default_dtype(torch.bfloat16)
+        actual, _ = quantize_iq1_s(weight)
+    finally:
+        torch.set_default_dtype(default_dtype)
+
+    assert torch.equal(actual, expected)
+
+
+def test_iq1_s_requires_complete_last_dimension_blocks():
+    with pytest.raises(ValueError, match="last weight dimension"):
+        quantize_iq1_s(torch.ones(2, 257))
+
+
+def test_iq1_s_treats_nonfinite_values_as_zero():
+    weight = torch.randn(1, 256)
+    weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf])
+
+    packed, _ = quantize_iq1_s(weight)
+    expected, _ = quantize_iq1_s(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0))
+
+    assert torch.equal(packed, expected)
+
+
+@pytest.mark.parametrize(
+    "weight_shape",
+    [
+        torch.tensor(256),
+        torch.tensor([[1, 256]]),
+        torch.tensor([1.0, 256.0]),
+        torch.tensor([0, 256]),
+    ],
+)
+def test_iq1_s_rejects_invalid_shape_metadata(weight_shape):
+    packed = torch.zeros((1, 1, 50), dtype=torch.uint8)
+
+    with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"):
+        dequantize_iq1_s(packed, weight_shape)
+
+
+def test_iq1_s_rejects_scalar_packed_payload():
+    with pytest.raises(ValueError, match="packed_weights"):
+        dequantize_iq1_s(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256]))
+
+
+def test_iq1_s_fake_quant_has_pass_through_gradient():
+    class Quantizer:
+        num_bits = "iq1_s"
+
+    weight = torch.randn(1, 256, requires_grad=True)
+    output = iq1_s_fake_quant(weight, Quantizer())
+    output.sum().backward()
+
+    assert torch.equal(weight.grad, torch.ones_like(weight))
+
+
+def test_iq1_s_saturates_finite_values_above_the_float32_range():
+    """float64 weights are accepted, so a finite value too large for float32 must saturate.
+
+    Converting before sanitizing would turn it into infinity and then zero, which silently
+    encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy.
+    """
+    weight = torch.randn(1, 256, dtype=torch.float64)
+    weight[0, 7] = 1e100
+    saturated = weight.clone()
+    saturated[0, 7] = torch.finfo(torch.float32).max
+    zeroed = weight.clone()
+    zeroed[0, 7] = 0.0
+
+    packed, _ = quantize_iq1_s(weight)
+
+    assert torch.equal(packed, quantize_iq1_s(saturated)[0])
+    assert not torch.equal(packed, quantize_iq1_s(zeroed)[0])
@@ -0,0 +1,168 @@
+# 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.
+
+import pytest
+import torch
+
+import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module
+from modelopt.torch.quantization.ggml.iq2_xs import (
+    IQ2_XS_BLOCK_BYTES,
+    dequantize_iq2_xs,
+    iq2_xs_fake_quant,
+    iq2_xs_grid,
+    quantize_iq2_xs,
+)
+
+
+def test_iq2_xs_canonical_grid():
+    grid = iq2_xs_grid()
+
+    assert grid.shape == (512, 8)
+    assert grid.dtype == torch.float32
+    assert set(grid.unique().tolist()) == {8.0, 25.0, 43.0}
+    assert grid[0].tolist() == [8.0] * 8
+    assert grid[-1].tolist() == [43.0] * 8
+
+
+def test_iq2_xs_grid_normalizes_unindexed_cuda_device(monkeypatch):
+    cached = torch.empty(0)
+    indexed_device = torch.device("cuda", 7)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 7)
+    monkeypatch.setitem(iq2_xs_module._GRID_CACHE, indexed_device, cached)
+
+    assert iq2_xs_grid("cuda") is cached
+
+
+def test_iq2_xs_zero_block_has_canonical_zero_encoding():
+    weight = torch.zeros((2, 256), dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq2_xs(weight)
+
+    assert packed.shape == (2, 1, IQ2_XS_BLOCK_BYTES)
+    assert packed.dtype == torch.uint8
+    assert not packed.any()
+    assert shape.tolist() == [2, 256]
+    assert torch.equal(dequantize_iq2_xs(packed, shape), weight)
+
+
+def test_iq2_xs_underflowed_scale_has_canonical_zero_encoding():
+    weight = torch.full((1, 256), -1e-6, dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq2_xs(weight)
+
+    assert not packed.any()
+    assert torch.equal(dequantize_iq2_xs(packed, shape), torch.zeros_like(weight))
+
+
+def test_iq2_xs_round_trip_and_payload_fields():
+    generator = torch.Generator().manual_seed(1234)
+    weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16)
+
+    packed, shape = quantize_iq2_xs(weight, block_chunk_size=2)
+    reconstructed = dequantize_iq2_xs(packed, shape, block_chunk_size=1)
+    default_reconstructed = dequantize_iq2_xs(packed, shape)
+
+    assert packed.shape == (2, 2, 74)
+    assert reconstructed.shape == weight.shape
+    assert reconstructed.dtype == torch.bfloat16
+    assert torch.equal(reconstructed, default_reconstructed)
+    normalized_mse = (
+        reconstructed.float() - weight.float()
+    ).square().mean() / weight.float().square().mean()
+    assert normalized_mse < 0.1
+
+    blocks = packed.reshape(-1, 74)
+    codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8)
+    assert torch.all((codes & 0x1FF) < 512)
+    assert torch.all((codes >> 9) < 128)
+
+
+def test_iq2_xs_dequantizes_pinned_scale_factor():
+    packed = torch.zeros((1, 1, 74), dtype=torch.uint8)
+    packed[0, 0, :2] = torch.tensor([1.0], dtype=torch.float16).view(torch.uint8)
+    packed[0, 0, 2:66:2] = 0xFF
+    packed[0, 0, 3:66:2] = 0x01
+    packed[0, 0, 66:] = 0xFF
+
+    decoded = dequantize_iq2_xs(packed, torch.tensor([1, 256]), dtype=torch.float32)
+
+    # Entry 511 contains eight 43s and local code 15 gives (2 * 15 + 1) / 8.
+    assert torch.equal(decoded, torch.full((1, 256), 43 * 31 / 8, dtype=torch.float32))
+
+
+def test_iq2_xs_requires_complete_last_dimension_blocks():
+    with pytest.raises(ValueError, match="last weight dimension"):
+        quantize_iq2_xs(torch.ones(2, 257))
+
+
+def test_iq2_xs_treats_nonfinite_values_as_zero():
+    weight = torch.randn(1, 256)
+    weight[0, :3] = torch.tensor([torch.nan, torch.inf, -torch.inf])
+
+    packed, _ = quantize_iq2_xs(weight)
+    expected, _ = quantize_iq2_xs(torch.nan_to_num(weight, nan=0.0, posinf=0.0, neginf=0.0))
+
+    assert torch.equal(packed, expected)
+
+
+@pytest.mark.parametrize(
+    "weight_shape",
+    [
+        torch.tensor(256),
+        torch.tensor([[1, 256]]),
+        torch.tensor([1.0, 256.0]),
+        torch.tensor([0, 256]),
+    ],
+)
+def test_iq2_xs_rejects_invalid_shape_metadata(weight_shape):
+    packed = torch.zeros((1, 1, 74), dtype=torch.uint8)
+
+    with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"):
+        dequantize_iq2_xs(packed, weight_shape)
+
+
+def test_iq2_xs_rejects_scalar_packed_payload():
+    with pytest.raises(ValueError, match="packed_weights"):
+        dequantize_iq2_xs(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256]))
+
+
+def test_iq2_xs_fake_quant_has_pass_through_gradient():
+    class Quantizer:
+        num_bits = "iq2_xs"
+
+    weight = torch.randn(1, 256, requires_grad=True)
+    output = iq2_xs_fake_quant(weight, Quantizer())
+    output.sum().backward()
+
+    assert torch.equal(weight.grad, torch.ones_like(weight))
+
+
+def test_iq2_xs_saturates_finite_values_above_the_float32_range():
+    """float64 weights are accepted, so a finite value too large for float32 must saturate.
+
+    Converting before sanitizing would turn it into infinity and then zero, which silently
+    encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy.
+    """
+    weight = torch.randn(1, 256, dtype=torch.float64)
+    weight[0, 7] = 1e100
+    saturated = weight.clone()
+    saturated[0, 7] = torch.finfo(torch.float32).max
+    zeroed = weight.clone()
+    zeroed[0, 7] = 0.0
+
+    packed, _ = quantize_iq2_xs(weight)
+
+    assert torch.equal(packed, quantize_iq2_xs(saturated)[0])
+    assert not torch.equal(packed, quantize_iq2_xs(zeroed)[0])