diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp similarity index 70% rename from modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp rename to modelopt/torch/kernels/quantization/ggml/ggml.cpp index 0929a961935..4af30719f2f 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -15,10 +15,24 @@ * limitations under the License. */ +// Every GGML IQ format shares common.cuh, the same CUDA version gate, and the same build flags, +// so they compile into one extension and bind here. Each format keeps its kernels in its own +// translation unit and exposes a single host entry point. + #include "common.cuh" +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); +namespace { + +at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_cuda(), "IQ1_S packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ1_S packing requires a CUDA grid"); + modelopt::ggml::check_pack_inputs("IQ1_S", input, grid, modelopt::ggml::kIq1sEntries); + return iq1_s_pack_cuda(input.contiguous(), grid.contiguous()); +} + at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { TORCH_CHECK(input.is_cuda(), "IQ2_XS packing requires a CUDA input"); TORCH_CHECK(grid.is_cuda(), "IQ2_XS packing requires a CUDA grid"); @@ -39,8 +53,15 @@ at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); } +} // namespace + PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { - module.def("pack", &iq2_xs_pack, + module.def("iq1_s_pack", &iq1_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 [2048, 8]. Returns uint8 " + "[numel / 256, 50] 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_xs_pack", &iq2_xs_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 [512, 8] holding " "non-negative codebook magnitudes, and scales must be finite non-negative float16 " diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp deleted file mode 100644 index dbc187d295d..00000000000 --- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp +++ /dev/null @@ -1,35 +0,0 @@ -/* - * 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" - -at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid); - -at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) { - TORCH_CHECK(input.is_cuda(), "IQ1_S packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ1_S packing requires a CUDA grid"); - modelopt::ggml::check_pack_inputs("IQ1_S", input, grid, modelopt::ggml::kIq1sEntries); - return iq1_s_pack_cuda(input.contiguous(), grid.contiguous()); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { - module.def("pack", &iq1_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 [2048, 8]. Returns uint8 " - "[numel / 256, 50] on the input device. Non-finite input elements are treated as " - "zero during packing, and finite elements outside the float32 range saturate."); -} diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 900bf666588..08ca0dbb719 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -22,8 +22,7 @@ __all__ = [ "get_cuda_ext", "get_cuda_ext_fp8", - "get_cuda_ext_iq1_s", - "get_cuda_ext_iq2_xs", + "get_cuda_ext_ggml", "get_cuda_ext_mx", "precompile", ] @@ -80,36 +79,29 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): return get_cuda_ext_mx.extension # type:ignore[attr-defined] -def get_cuda_ext_iq1_s(raise_if_failed: bool = False): - """Return the GGML-compatible IQ1_S packing extension.""" - if not hasattr(get_cuda_ext_iq1_s, "extension") or ( - raise_if_failed and get_cuda_ext_iq1_s.extension is None - ): - get_cuda_ext_iq1_s.extension = load_cpp_extension( # type:ignore[attr-defined] - name="modelopt_cuda_ext_iq1_s", - sources=[kernels_ggml / "iq1_s.cpp", kernels_ggml / "iq1_s.cu"], - cuda_version_specifiers=">=11.8", - fail_msg="IQ1_S CUDA packing extension is unavailable.", - extra_cuda_cflags=["-O3"], - raise_if_failed=raise_if_failed, - ) - return get_cuda_ext_iq1_s.extension # type:ignore[attr-defined] - +def get_cuda_ext_ggml(raise_if_failed: bool = False): + """Return the GGML-compatible IQ packing extension, exposing one packer per IQ format. -def get_cuda_ext_iq2_xs(raise_if_failed: bool = False): - """Return the GGML-compatible IQ2_XS packing extension.""" - if not hasattr(get_cuda_ext_iq2_xs, "extension") or ( - raise_if_failed and get_cuda_ext_iq2_xs.extension is None + The formats share their packing helpers, CUDA version requirement, and build flags, so they + build as a single extension: ``iq1_s_pack(input, grid)`` and + ``iq2_xs_pack(input, grid, scales)``. + """ + if not hasattr(get_cuda_ext_ggml, "extension") or ( + raise_if_failed and get_cuda_ext_ggml.extension is None ): - get_cuda_ext_iq2_xs.extension = load_cpp_extension( # type:ignore[attr-defined] - name="modelopt_cuda_ext_iq2_xs", - sources=[kernels_ggml / "iq2_xs.cpp", kernels_ggml / "iq2_xs.cu"], + get_cuda_ext_ggml.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_ggml", + sources=[ + kernels_ggml / "ggml.cpp", + kernels_ggml / "iq1_s.cu", + kernels_ggml / "iq2_xs.cu", + ], cuda_version_specifiers=">=11.8", - fail_msg="IQ2_XS CUDA packing extension is unavailable.", + fail_msg="GGML IQ CUDA packing extension is unavailable.", extra_cuda_cflags=["-O3"], raise_if_failed=raise_if_failed, ) - return get_cuda_ext_iq2_xs.extension # type:ignore[attr-defined] + return get_cuda_ext_ggml.extension # type:ignore[attr-defined] def __getattr__(name): @@ -119,10 +111,8 @@ def __getattr__(name): return get_cuda_ext_fp8() elif name == "cuda_ext_mx": return get_cuda_ext_mx() - elif name == "cuda_ext_iq1_s": - return get_cuda_ext_iq1_s() - elif name == "cuda_ext_iq2_xs": - return get_cuda_ext_iq2_xs() + elif name == "cuda_ext_ggml": + return get_cuda_ext_ggml() else: raise AttributeError(f"module {__name__} has no attribute {name}") @@ -132,5 +122,4 @@ def precompile(): print(get_cuda_ext()) print(get_cuda_ext_fp8()) print(get_cuda_ext_mx()) - print(get_cuda_ext_iq1_s()) - print(get_cuda_ext_iq2_xs()) + print(get_cuda_ext_ggml()) diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py index ea1e41a5b94..68ff3f2e7e5 100644 --- a/tests/gpu/_extensions/test_torch_extensions.py +++ b/tests/gpu/_extensions/test_torch_extensions.py @@ -14,7 +14,6 @@ # limitations under the License. -from collections.abc import Callable from typing import NamedTuple import pytest @@ -39,12 +38,8 @@ def test_cuda_ext_mx(): assert ext.get_cuda_ext_mx() is not None -def test_cuda_ext_iq1_s(): - assert ext.get_cuda_ext_iq1_s() is not None - - -def test_cuda_ext_iq2_xs(): - assert ext.get_cuda_ext_iq2_xs() is not None +def test_cuda_ext_ggml(): + assert ext.get_cuda_ext_ggml() is not None def _generator(): @@ -53,9 +48,10 @@ def _generator(): class _IqFormat(NamedTuple): - """One GGML IQ packing extension and the format constants its contract is defined by.""" + """One GGML IQ packer and the format constants its contract is defined by.""" - get_extension: Callable + # Name the packer is bound under on the shared GGML extension. + packer: str entries: int payload_bytes: int needs_scales: bool @@ -68,11 +64,9 @@ class _IqFormat(NamedTuple): _IQ_EXTENSIONS = ( + pytest.param(_IqFormat("iq1_s_pack", 2048, 50, False, (-1.0, 0.0, 1.0), 16.875), id="iq1_s"), pytest.param( - _IqFormat(ext.get_cuda_ext_iq1_s, 2048, 50, False, (-1.0, 0.0, 1.0), 16.875), id="iq1_s" - ), - pytest.param( - _IqFormat(ext.get_cuda_ext_iq2_xs, 512, 74, True, (8.0, 25.0, 43.0), 166.625), + _IqFormat("iq2_xs_pack", 512, 74, True, (8.0, 25.0, 43.0), 166.625), id="iq2_xs", ), ) @@ -90,16 +84,17 @@ def _grid(fmt: _IqFormat, zero: bool = False) -> torch.Tensor: def _pack(fmt: _IqFormat, extension, weight, grid, scales=None): + pack = getattr(extension, fmt.packer) if not fmt.needs_scales: - return extension.pack(weight, grid) + return pack(weight, grid) if scales is None: scales = torch.zeros(weight.numel() // 256, device=weight.device, dtype=torch.float16) - return extension.pack(weight, grid, scales) + return pack(weight, grid, scales) @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_zero_block_layout(fmt): - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) weight = torch.zeros((2, 256), device="cuda", dtype=torch.bfloat16) packed = _pack(fmt, extension, weight, _grid(fmt, zero=True)) @@ -111,7 +106,7 @@ def test_cuda_ext_iq_zero_block_layout(fmt): @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_encodes_non_zero_block(fmt): """Exercise the encode loop itself: search, reductions, and the payload writes.""" - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) weight = torch.randn((2, 256), device="cuda", dtype=torch.bfloat16, generator=_generator()) scales = (weight.float().abs().amax(dim=-1) / fmt.native_max).half() @@ -132,7 +127,7 @@ def test_cuda_ext_iq_encodes_non_zero_block(fmt): @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_rejects_unsupported_dtype(fmt): - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) weight = torch.ones((1, 256), device="cuda").to(torch.float8_e4m3fn) with pytest.raises(RuntimeError, match="supports float32, float64, float16, and bfloat16"): @@ -141,7 +136,7 @@ def test_cuda_ext_iq_rejects_unsupported_dtype(fmt): @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_rejects_row_straddling_input(fmt): - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) weight = torch.ones((512, 384), device="cuda", dtype=torch.bfloat16) with pytest.raises(RuntimeError, match="innermost dimension must be a multiple of 256"): @@ -151,24 +146,24 @@ def test_cuda_ext_iq_rejects_row_straddling_input(fmt): @pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf"), -1.0, -1e-4]) def test_cuda_ext_iq2_xs_rejects_invalid_scales(bad): """A non-finite scale decodes to garbage; a negative one inverts every decoded element.""" - extension = ext.get_cuda_ext_iq2_xs(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) fmt = _IQ_EXTENSIONS[1].values[0] weight = torch.ones((1, 256), device="cuda", dtype=torch.bfloat16) scales = torch.full((1,), bad, device="cuda", dtype=torch.float16) with pytest.raises(RuntimeError, match="scales must be finite and non-negative"): - extension.pack(weight, _grid(fmt), scales) + extension.iq2_xs_pack(weight, _grid(fmt), scales) def test_cuda_ext_iq2_xs_negative_zero_scale_packs_as_zero(): """Negative zero is a zero scale: it must take the zero-payload branch, not search.""" - extension = ext.get_cuda_ext_iq2_xs(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) fmt = _IQ_EXTENSIONS[1].values[0] weight = torch.randn((2, 256), device="cuda", dtype=torch.bfloat16, generator=_generator()) grid = _random_grid(fmt) scales = torch.tensor([-0.0, 0.0], device="cuda", dtype=torch.float16) - packed = extension.pack(weight, grid, scales) + packed = extension.iq2_xs_pack(weight, grid, scales) assert not packed.any() @@ -260,7 +255,7 @@ def test_cuda_ext_iq_encoding_is_optimal(fmt): misplaced index, local scale, delta sign, or sign bit makes the reconstruction worse than the brute-force optimum rather than merely different. """ - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) weight = torch.randn((4, 256), device="cuda", dtype=torch.float32, generator=_generator()) grid = _random_grid(fmt) scales = (weight.abs().amax(dim=-1) / fmt.native_max).half() if fmt.needs_scales else None @@ -284,7 +279,7 @@ def test_cuda_ext_iq_encoding_is_optimal(fmt): @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_input_dtype_equivalence(fmt): """Every accepted input dtype carrying identical values must pack to identical bytes.""" - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) # Multiples of 1/16 in [-4, 4) are exact in float16 and bfloat16 as well as the wider types. weight = torch.randint(-64, 64, (2, 256), device="cuda", generator=_generator()).float() / 16 grid = _random_grid(fmt) @@ -302,7 +297,7 @@ def test_cuda_ext_iq_input_dtype_equivalence(fmt): @pytest.mark.parametrize("fmt", _IQ_EXTENSIONS) def test_cuda_ext_iq_non_finite_inputs_are_zeroed(fmt): """NaN and infinity pack as zeros; finite values too large for float32 saturate instead.""" - extension = fmt.get_extension(raise_if_failed=True) + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) clean = torch.randn((2, 256), device="cuda", dtype=torch.float32, generator=_generator()) grid = _random_grid(fmt) scales = (clean.abs().amax(dim=-1) / fmt.native_max).half() if fmt.needs_scales else None