[OMNIML-5899] Add Q8_0 CUDA packing kernel

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<int>::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<int>::max(), format,
-              " CUDA grid is too large");
 }
 
 #ifdef __CUDACC__

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.");
 }

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 <typename scalar_t>
+__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<uint8_t>(d_bits);
+    payload[1] = static_cast<uint8_t>(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<int>(fminf(fmaxf(rounded, -kMaxQuant), kMaxQuant));
+  payload[kScaleBytes + lane] = static_cast<uint8_t>(static_cast<int8_t>(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<scalar_t><<<static_cast<int>(num_blocks), kQ8BlockSize, 0, stream>>>(
+            input.data_ptr<scalar_t>(), num_blocks, output.data_ptr<uint8_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}

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.",

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)