Share one CUDA encoder per IQ family

modelopt/torch/kernels/quantization/ggml/common.cuh

@@ -41,6 +41,7 @@ namespace modelopt::ggml {
 // https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kBlockSize = 256;
 constexpr int kVectorSize = 8;
+constexpr int kVectorsPerBlock = kBlockSize / kVectorSize;
 constexpr int kScaleOffset = 0;
 constexpr int kScaleBytes = 2;
 
@@ -72,18 +73,38 @@ inline void check_scalar_pack_input(const char *format, const at::Tensor &input,
               " 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.
+// Validates the codebook contract every IQ format shares, once, at each format's CUDA entry point.
+// A format that takes block scales passes them so their device is checked with the others', before
+// any shape or dtype rule.
 inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid,
-                              int64_t entries) {
+                              int64_t entries, const at::Tensor *scales = nullptr) {
+  TORCH_CHECK(input.is_cuda(), format, " packing requires a CUDA input");
+  TORCH_CHECK(grid.is_cuda(), format, " packing requires a CUDA grid");
+  TORCH_CHECK(scales == nullptr || scales->is_cuda(), format, " packing requires CUDA scales");
   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");
 }
 
+// Validates a format that takes one precomputed FP16 block scale per 256 values. The kernels copy
+// these bits straight into the block scale field: a non-finite entry produces a payload that
+// decodes to garbage, and a negative one inverts the sign of every decoded element while still
+// packing cleanly -- GGML's own encoders assert a non-negative block scale. One fused reduction, so
+// the synchronization is paid once per packed tensor, on an export path.
+inline void check_scaled_pack_inputs(const char *format, const at::Tensor &input,
+                                     const at::Tensor &grid, int64_t entries,
+                                     const at::Tensor &scales) {
+  check_pack_inputs(format, input, grid, entries, &scales);
+  TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
+                  scales.numel() == input.numel() / kBlockSize,
+              "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");
+}
+
 #ifdef __CUDACC__
 
 // Reads one input element as float32. Non-finite elements are treated as zero, and finiteness is

modelopt/torch/kernels/quantization/ggml/ggml.cpp

@@ -16,8 +16,9 @@
  */
 
 // 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.
+// so they compile into one extension and bind here. Each IQ format keeps its block layout in its
+// own translation unit, on the encoder its family shares (iq1_family.cuh or iq2_family.cuh), and
+// exposes its packer.
 
 #include "common.cuh"
 
@@ -30,80 +31,6 @@ at::Tensor q8_0_pack_cuda(at::Tensor input);
 
 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");
-  TORCH_CHECK(scales.is_cuda(), "IQ2_XS packing requires CUDA scales");
-  modelopt::ggml::check_pack_inputs("IQ2_XS", input, grid, modelopt::ggml::kIq2xsEntries);
-  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]");
-  // The kernel copies these bits straight into the GGML block scale field. A non-finite entry
-  // would produce a payload that decodes to garbage, and a negative one inverts the sign of every
-  // decoded element while still packing cleanly -- GGML's own encoders assert a non-negative block
-  // scale. One fused reduction, so the synchronization is paid once per packed tensor, on an
-  // export path.
-  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_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
-}
-
-at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
-  TORCH_CHECK(input.is_cuda(), "IQ2_XXS packing requires a CUDA input");
-  TORCH_CHECK(grid.is_cuda(), "IQ2_XXS packing requires a CUDA grid");
-  TORCH_CHECK(scales.is_cuda(), "IQ2_XXS packing requires CUDA scales");
-  modelopt::ggml::check_pack_inputs("IQ2_XXS", input, grid, modelopt::ggml::kIq2xxsEntries);
-  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]");
-  // Same rule as IQ2_XS: these bits become the block scale verbatim, and a negative or non-finite
-  // one packs cleanly while decoding to garbage.
-  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_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 iq1_m_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
-  TORCH_CHECK(input.is_cuda(), "IQ1_M packing requires a CUDA input");
-  TORCH_CHECK(grid.is_cuda(), "IQ1_M packing requires a CUDA grid");
-  TORCH_CHECK(scales.is_cuda(), "IQ1_M packing requires CUDA scales");
-  modelopt::ggml::check_pack_inputs("IQ1_M", input, grid, modelopt::ggml::kIq1sEntries);
-  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 iq1_m_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);
@@ -113,36 +40,36 @@ at::Tensor q8_0_pack(at::Tensor input) {
 } // namespace
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
-  module.def("iq1_s_pack", &iq1_s_pack,
+  module.def("iq1_s_pack", &iq1_s_pack_cuda,
              "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,
+  module.def("iq2_xs_pack", &iq2_xs_pack_cuda,
              "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 "
              "[numel / 256]. "
              "Returns uint8 [numel / 256, 74] 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_xxs_pack", &iq2_xxs_pack,
+  module.def("iq2_xxs_pack", &iq2_xxs_pack_cuda,
              "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
              "dimension is a multiple of 256. The grid must be float32 [256, 8] holding "
              "non-negative codebook magnitudes, and scales must be finite non-negative float16 "
              "[numel / 256]. "
              "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,
+  module.def("iq2_s_pack", &iq2_s_pack_cuda,
              "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("iq1_m_pack", &iq1_m_pack,
+  module.def("iq1_m_pack", &iq1_m_pack_cuda,
              "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], and scales must "
              "be finite non-negative float16 [numel / 256]; IQ1_M carries that scale in the four "

modelopt/torch/kernels/quantization/ggml/iq1_family.cuh

@@ -0,0 +1,191 @@
+/*
+ * 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.
+ */
+
+// The IQ1 formats share one encoder. Every 8-value vector is approximated by scale * (q + delta),
+// with q from the 2048-entry ternary grid and delta = +/- 1/8; the grid is 64 KiB, past the 48 KiB
+// static shared-memory limit, so it is read from global memory and left to the cache. Every group
+// of vectors picks one of Format::kChoices options, with local scale choice % 8. A format
+// describes the rest with a Format type:
+//
+//   kGroups, kVectorsPerGroup, kChoices, kPayloadBytes  block layout
+//   kSharedShift  the choice also fixes the delta for the whole group, as choice / 8 (IQ1_S);
+//                 otherwise each vector picks its own (IQ1_M)
+//   begin(payload, d_bits)  false when the block is all zero and needs nothing more written
+//   store(payload, picks, choices, d_bits)  writes every vector's pick -- its grid entry, with
+//                 its own shift in bit kIq1EntryBits when the shift is per vector -- and every
+//                 group's choice; called by every thread once the search is done
+
+#pragma once
+
+#include "common.cuh"
+
+namespace modelopt::ggml {
+
+#ifdef __CUDACC__
+
+constexpr float kIq1Delta = 0.125f;
+constexpr int kIq1EntryBits = 11; // 2048 entries; a per-vector shift sits just above them
+
+__device__ __forceinline__ float iq1_delta(int shift) { return shift ? -kIq1Delta : kIq1Delta; }
+
+template <typename Format, typename scalar_t>
+__global__ void iq1_encode(const scalar_t *input, int64_t num_blocks, const float *grid,
+                           const __half *scales, uint8_t *output) {
+  constexpr int kGroups = Format::kGroups;
+  constexpr int kVectorsPerGroup = Format::kVectorsPerGroup;
+  constexpr int kChoices = Format::kChoices;
+  constexpr bool kSharedShift = Format::kSharedShift;
+  static_assert(kIq1sEntries % kThreads == 0, "every thread must visit the same number of entries");
+  static_assert(kIq1sEntries == 1 << kIq1EntryBits, "the shift bit sits just above the entry");
+  static_assert(kGroups * kVectorsPerGroup == kVectorsPerBlock, "groups must cover the block");
+
+  __shared__ float warp_best[kWarps * kChoices];
+  __shared__ float group_error[kChoices];
+  __shared__ unsigned long long warp_keys[kWarps];
+  __shared__ int selected_choice;
+  __shared__ uint8_t choices[kGroups];
+  __shared__ uint16_t picks[kVectorsPerBlock];
+
+  const int tid = threadIdx.x;
+  const int64_t block = blockIdx.x;
+  if (block >= num_blocks)
+    return;
+
+  const scalar_t *source = input + block * kBlockSize;
+  uint8_t *payload = output + block * Format::kPayloadBytes;
+  const __half d_half = scales[block];
+  const uint16_t d_bits = __half_as_ushort(d_half);
+  const float d = __half2float(d_half);
+  if (!Format::begin(payload, d_bits))
+    return;
+
+#pragma unroll 1
+  for (int group = 0; group < kGroups; ++group) {
+    if (tid < kChoices)
+      group_error[tid] = 0.0f;
+    __syncthreads();
+
+    // Score every choice: each vector's best error under it, summed over the group. With a
+    // per-vector shift, a vector takes the better of the two before the group chooses.
+#pragma unroll
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
+      float x[kVectorSize];
+      float xnorm, xsum;
+      load_vector(source + (group * kVectorsPerGroup + vector) * kVectorSize, x, xnorm, xsum);
+      float local_best[kChoices];
+#pragma unroll
+      for (int choice = 0; choice < kChoices; ++choice)
+        local_best[choice] = FLT_MAX;
+      for (int entry = tid; entry < kIq1sEntries; entry += blockDim.x) {
+        float dot, qnorm, qsum;
+        grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum);
+#pragma unroll
+        for (int choice = 0; choice < kChoices; ++choice) {
+          const float scale = d * (2 * (choice & 7) + 1);
+          if constexpr (kSharedShift) {
+            local_best[choice] =
+                fminf(local_best[choice],
+                      shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, iq1_delta(choice >> 3)));
+          } else {
+#pragma unroll
+            for (int shift = 0; shift < 2; ++shift)
+              local_best[choice] =
+                  fminf(local_best[choice],
+                        shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, iq1_delta(shift)));
+          }
+        }
+      }
+      block_min_accumulate<kChoices>(local_best, warp_best, group_error);
+    }
+
+    if (tid == 0) {
+      selected_choice = 0;
+      float best = group_error[0];
+#pragma unroll
+      for (int choice = 1; choice < kChoices; ++choice) {
+        if (group_error[choice] < best) {
+          best = group_error[choice];
+          selected_choice = choice;
+        }
+      }
+      choices[group] = static_cast<uint8_t>(selected_choice);
+    }
+    __syncthreads();
+    const float selected_scale = d * (2 * (selected_choice & 7) + 1);
+
+    // Under the chosen option, each vector takes its best entry. A per-vector shift sits above
+    // the entry index in the key, so a tie prefers the lower shift and then the lower entry.
+#pragma unroll
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
+      const int slot = group * kVectorsPerGroup + vector;
+      float x[kVectorSize];
+      float xnorm, xsum;
+      load_vector(source + slot * kVectorSize, x, xnorm, xsum);
+      unsigned long long key = ~0ULL;
+      for (int entry = tid; entry < kIq1sEntries; entry += blockDim.x) {
+        float dot, qnorm, qsum;
+        grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum);
+        if constexpr (kSharedShift) {
+          const float error = shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale,
+                                            iq1_delta(selected_choice >> 3));
+          const unsigned long long candidate = error_key(error, entry);
+          key = candidate < key ? candidate : key;
+        } else {
+#pragma unroll
+          for (int shift = 0; shift < 2; ++shift) {
+            const float error =
+                shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, iq1_delta(shift));
+            const unsigned long long candidate = error_key(error, (shift << kIq1EntryBits) | entry);
+            key = candidate < key ? candidate : key;
+          }
+        }
+      }
+      key = block_min_key(key, warp_keys);
+      if (tid == 0)
+        picks[slot] = static_cast<uint16_t>(key & ((1u << (kIq1EntryBits + 1)) - 1));
+    }
+  }
+  __syncthreads();
+  Format::store(payload, picks, choices, d_bits);
+}
+
+// Packs one IQ1 format from per-block FP16 scales, once its inputs have been validated.
+template <typename Format>
+at::Tensor iq1_encode_blocks(const at::Tensor &input, const at::Tensor &grid,
+                             const at::Tensor &scales) {
+  const auto values = input.contiguous();
+  const auto table = grid.contiguous();
+  const auto block_scales = scales.contiguous();
+  c10::cuda::CUDAGuard guard(values.device());
+  const int64_t num_blocks = values.numel() / kBlockSize;
+  at::Tensor output =
+      at::empty({num_blocks, Format::kPayloadBytes}, values.options().dtype(at::kByte));
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, values.scalar_type(), "iq1_pack", [&] {
+        iq1_encode<Format, scalar_t><<<static_cast<int>(num_blocks), kThreads, 0, stream>>>(
+            values.data_ptr<scalar_t>(), num_blocks, table.data_ptr<float>(),
+            reinterpret_cast<const __half *>(block_scales.data_ptr<at::Half>()),
+            output.data_ptr<uint8_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}
+
+#endif // __CUDACC__
+
+} // namespace modelopt::ggml