diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index cd234d3423b..75cb483e4f6 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -56,25 +56,31 @@ constexpr int kIq2xxsEntries = 256; constexpr int kThreads = 256; constexpr int kWarps = kThreads / 32; -// Validates the packing contract every IQ format shares. Called from the pybind wrapper on the -// caller's tensors and again from the CUDA entry point on the materialized contiguous tensors, so -// the enforced rule and the message it reports are written once. -inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, - int64_t entries) { +// Validates the scalar-input contract shared by the vector-codebook and scalar GGML packers. +inline void check_scalar_pack_input(const char *format, const at::Tensor &input, + int64_t block_size) { const auto input_type = input.scalar_type(); TORCH_CHECK(input_type == at::kFloat || input_type == at::kDouble || input_type == at::kHalf || input_type == at::kBFloat16, format, " packing supports float32, float64, float16, and bfloat16 inputs"); TORCH_CHECK(input.numel() > 0, "input must be non-empty"); - TORCH_CHECK(input.dim() > 0 && input.size(-1) % kBlockSize == 0, - "input's innermost dimension must be a multiple of ", kBlockSize, + TORCH_CHECK(input.dim() > 0 && input.size(-1) % block_size == 0, + "input's innermost dimension must be a multiple of ", block_size, " so blocks do not straddle rows"); + TORCH_CHECK(input.numel() / block_size <= std::numeric_limits::max(), format, + " CUDA grid is too large"); +} + +// Validates the codebook contract every IQ format shares. Called from the pybind wrapper on the +// caller's tensors and again from the CUDA entry point on the materialized contiguous tensors, so +// the enforced rule and the message it reports are written once. +inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, + int64_t entries) { + check_scalar_pack_input(format, input, kBlockSize); TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == entries && grid.size(1) == kVectorSize, "grid must be float32 [", entries, ", ", kVectorSize, "]"); TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); - TORCH_CHECK(input.numel() / kBlockSize <= std::numeric_limits::max(), format, - " CUDA grid is too large"); } #ifdef __CUDACC__ diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index 5ee5b0a3fbf..08b75243ad3 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/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 q8_0_pack_cuda(at::Tensor input); namespace { @@ -71,6 +72,12 @@ 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 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); + return q8_0_pack_cuda(input.contiguous()); +} + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { @@ -95,4 +102,9 @@ 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("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. " + "Non-finite input elements are treated as zero during packing, and finite elements " + "outside the float32 range saturate."); } diff --git a/modelopt/torch/kernels/quantization/ggml/q8_0.cu b/modelopt/torch/kernels/quantization/ggml/q8_0.cu new file mode 100644 index 00000000000..5f9707bfcc7 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/q8_0.cu @@ -0,0 +1,78 @@ +/* + * 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; + +// This CUDA encoder was independently implemented for ModelOpt. It follows the Q8_0 packed +// payload and scalar quantization formula defined by GGML; no llama.cpp implementation code is +// incorporated here. Format reference: +// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +// Scalar-formula reference: +// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-quants.c +constexpr int kQ8BlockSize = 32; +constexpr int kPayloadBytes = kScaleBytes + kQ8BlockSize; +constexpr float kMaxQuant = 127.0f; + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, uint8_t *output) { + const int lane = threadIdx.x; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + const float value = load_float(input + block * kQ8BlockSize + lane); + float amax = fabsf(value); +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + amax = fmaxf(amax, __shfl_down_sync(0xffffffff, amax, delta)); + + float d = 0.0f; + uint8_t *payload = output + block * kPayloadBytes; + if (lane == 0) { + d = fminf(amax / kMaxQuant, 65504.0f); + const uint16_t d_bits = __half_as_ushort(__float2half_rn(d)); + payload[0] = static_cast(d_bits); + payload[1] = static_cast(d_bits >> 8); + } + d = __shfl_sync(0xffffffff, d, 0); + const float rounded = d > 0.0f ? roundf(value / d) : 0.0f; + const int quant = static_cast(fminf(fmaxf(rounded, -kMaxQuant), kMaxQuant)); + payload[kScaleBytes + lane] = static_cast(static_cast(quant)); +} + +} // namespace + +at::Tensor q8_0_pack_cuda(at::Tensor input) { + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + modelopt::ggml::check_scalar_pack_input("Q8_0", input, kQ8BlockSize); + const int64_t num_blocks = input.numel() / kQ8BlockSize; + 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(), "q8_0_pack", [&] { + encode<<(num_blocks), kQ8BlockSize, 0, stream>>>( + input.data_ptr(), num_blocks, output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 0b2238e452b..d20c80a5eee 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -80,11 +80,12 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): def get_cuda_ext_ggml(raise_if_failed: bool = False): - """Return the GGML-compatible IQ packing extension, exposing one packer per IQ format. + """Return the GGML-compatible packing extension, exposing one packer per format. 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)``. + 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)``. """ if not hasattr(get_cuda_ext_ggml, "extension") or ( raise_if_failed and get_cuda_ext_ggml.extension is None @@ -96,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 / "q8_0.cu", ], cuda_version_specifiers=">=11.8", fail_msg="GGML IQ CUDA packing extension is unavailable.", diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py index 68ff3f2e7e5..8c10a81ad4f 100644 --- a/tests/gpu/_extensions/test_torch_extensions.py +++ b/tests/gpu/_extensions/test_torch_extensions.py @@ -319,3 +319,92 @@ def test_cuda_ext_iq_non_finite_inputs_are_zeroed(fmt): _pack(fmt, extension, huge, grid, scales=scales), _pack(fmt, extension, zeroed.double(), grid, scales=scales), ) + + +def test_cuda_ext_q8_0_zero_and_roundf_layout(): + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + weight = torch.zeros((2, 32), device="cuda", dtype=torch.float32) + weight[1, :5] = torch.tensor([127.0, 0.5, -0.5, 1.5, -1.5], device="cuda") + + packed = extension.q8_0_pack(weight).cpu() + + assert packed.shape == (2, 34) + assert not packed[0].any() + assert packed[1, :2].contiguous().view(torch.float16).item() == 1.0 + assert packed[1, 2:7].contiguous().view(torch.int8).tolist() == [127, 1, -1, 2, -2] + + +def test_cuda_ext_q8_0_dequantizes_with_small_error(): + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + weight = torch.randn((4, 32), device="cuda", dtype=torch.bfloat16, generator=_generator()) + + packed = extension.q8_0_pack(weight).cpu() + d = packed[:, :2].contiguous().view(torch.float16).float().flatten() + quants = packed[:, 2:].contiguous().view(torch.int8).float() + decoded = d.unsqueeze(1) * quants + cpu_weight = weight.cpu().float() + normalized_mse = (decoded - cpu_weight).square().mean() / cpu_weight.square().mean() + + assert normalized_mse < 1e-4 + + +def test_cuda_ext_q8_0_input_dtype_equivalence(): + """Exactly representable values must produce identical bytes for every accepted dtype.""" + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + weight = torch.randint(-64, 64, (2, 32), device="cuda", generator=_generator()).float() / 16 + + payloads = [ + extension.q8_0_pack(weight.to(dtype)) + for dtype in (torch.float32, torch.float64, torch.float16, torch.bfloat16) + ] + + for dtype, payload in zip((torch.float64, torch.float16, torch.bfloat16), payloads[1:]): + assert torch.equal(payloads[0], payload), f"{dtype} disagrees with float32" + + +def test_cuda_ext_q8_0_nonfinite_and_float64_overflow_policy(): + """Non-finite inputs become zero; finite float64 overflow saturates to float32.""" + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + clean = torch.randn((2, 32), device="cuda", dtype=torch.float32, generator=_generator()) + zeroed = clean.clone() + zeroed[0, 5], zeroed[0, 20], zeroed[1, 17] = 0.0, 0.0, 0.0 + spoiled = clean.clone() + spoiled[0, 5], spoiled[0, 20], spoiled[1, 17] = float("nan"), float("inf"), float("-inf") + + assert torch.equal(extension.q8_0_pack(spoiled), extension.q8_0_pack(zeroed)) + + huge = zeroed.double() + huge[1, 7] = 1e100 + saturated = zeroed.double() + saturated[1, 7] = torch.finfo(torch.float32).max + huge_payload = extension.q8_0_pack(huge) + assert torch.equal(huge_payload, extension.q8_0_pack(saturated)) + assert not torch.equal(huge_payload, extension.q8_0_pack(zeroed.double())) + + +def test_cuda_ext_q8_0_fp16_scale_saturation_and_underflow(): + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + weight = torch.zeros((2, 32), device="cuda", dtype=torch.float32) + weight[0, 0] = 2 * 127 * torch.finfo(torch.float16).max + weight[1, 0] = 127 * 2.0**-26 + + packed = extension.q8_0_pack(weight).cpu() + scales = packed[:, :2].contiguous().view(torch.float16).float().flatten() + quants = packed[:, 2:].contiguous().view(torch.int8) + decoded = scales.unsqueeze(1) * quants.float() + + assert scales.tolist() == [torch.finfo(torch.float16).max, 0.0] + assert quants[:, 0].tolist() == [127, 127] + assert decoded[0, 0].item() == torch.finfo(torch.float16).max * 127 + assert not decoded[1].any() + + +def test_cuda_ext_q8_0_rejects_invalid_input(): + extension = ext.get_cuda_ext_ggml(raise_if_failed=True) + unsupported = torch.ones((1, 32), device="cuda").to(torch.float8_e4m3fn) + row_straddling = torch.ones((2, 48), device="cuda", dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="supports float32, float64, float16, and bfloat16"): + extension.q8_0_pack(unsupported) + with pytest.raises(RuntimeError, match="innermost dimension must be a multiple of 32"): + extension.q8_0_pack(row_straddling)